diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index a8b5636e..1324c61e 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -20,7 +20,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: ['3.7', '3.8', '3.9', '3.10', '3.11', '3.12'] + python-version: ['3.8', '3.9', '3.10', '3.11', '3.12'] steps: - uses: actions/checkout@v4 - uses: actions/setup-python@v5 diff --git a/glass/test/test_fits.py b/glass/test/test_fits.py new file mode 100644 index 00000000..25549388 --- /dev/null +++ b/glass/test/test_fits.py @@ -0,0 +1,89 @@ +import pytest + +# check if fitsio is available for testing +import importlib.util +if importlib.util.find_spec("fitsio") is not None: + HAVE_FITSIO = True +else: + HAVE_FITSIO = False + +import glass.user as user +import numpy as np + + +def _test_append(fits, data, names): + '''Write routine for FITS test cases''' + cat_name = 'CATALOG' + if cat_name not in fits: + fits.write_table(data, names=names, extname=cat_name) + else: + hdu = fits[cat_name] + hdu.write(data, names=names, firstrow=hdu.get_nrows()) + + +delta = 0.001 # Number of points in arrays +my_max = 1000 # Typically number of galaxies in loop +except_int = 750 # Where test exception occurs in loop +filename = "MyFile.Fits" + + +@pytest.mark.skipif(not HAVE_FITSIO, reason="test requires fitsio") +def test_basic_write(tmp_path): + import fitsio + d = tmp_path / "sub" + d.mkdir() + filename_gfits = "gfits.fits" # what GLASS creates + filename_tfits = "tfits.fits" # file created on the fly to test against + + with user.write_catalog(d / filename_gfits, ext="CATALOG") as out, fitsio.FITS(d / filename_tfits, "rw", clobber=True) as myFits: + for i in range(0, my_max): + array = np.arange(i, i+1, delta) # array of size 1/delta + array2 = np.arange(i+1, i+2, delta) # array of size 1/delta + out.write(RA=array, RB=array2) + arrays = [array, array2] + names = ['RA', 'RB'] + _test_append(myFits, arrays, names) + + with fitsio.FITS(d / filename_gfits) as g_fits, fitsio.FITS(d / filename_tfits) as t_fits: + glass_data = g_fits[1].read() + test_data = t_fits[1].read() + assert glass_data['RA'].size == test_data['RA'].size + assert glass_data['RB'].size == test_data['RA'].size + + +@pytest.mark.skipif(not HAVE_FITSIO, reason="test requires fitsio") +def test_write_exception(tmp_path): + d = tmp_path / "sub" + d.mkdir() + + try: + with user.write_catalog(d / filename, ext="CATALOG") as out: + for i in range(0, my_max): + if i == except_int: + raise Exception("Unhandled exception") + array = np.arange(i, i+1, delta) # array of size 1/delta + array2 = np.arange(i+1, i+2, delta) # array of size 1/delta + out.write(RA=array, RB=array2) + + except Exception: + import fitsio + with fitsio.FITS(d / filename) as hdul: + data = hdul[1].read() + assert data['RA'].size == except_int/delta + assert data['RB'].size == except_int/delta + + fitsMat = data['RA'].reshape(except_int, int(1/delta)) + fitsMat2 = data['RB'].reshape(except_int, int(1/delta)) + for i in range(0, except_int): + array = np.arange(i, i+1, delta) # re-create array to compare to read data + array2 = np.arange(i+1, i+2, delta) + assert array.tolist() == fitsMat[i].tolist() + assert array2.tolist() == fitsMat2[i].tolist() + + +@pytest.mark.skipif(not HAVE_FITSIO, reason="test requires fitsio") +def test_out_filename(tmp_path): + import fitsio + fits = fitsio.FITS(filename, "rw", clobber=True) + writer = user._FitsWriter(fits) + assert writer.fits._filename == filename diff --git a/glass/user.py b/glass/user.py index af2f13bd..7ff63b74 100644 --- a/glass/user.py +++ b/glass/user.py @@ -10,15 +10,17 @@ library. -Input and output +Input and Output ---------------- .. autofunction:: save_cls .. autofunction:: load_cls +.. autofunction:: write_catalog ''' import numpy as np +from contextlib import contextmanager def save_cls(filename, cls): @@ -45,3 +47,61 @@ def load_cls(filename): values = npz['values'] split = npz['split'] return np.split(values, split) + + +class _FitsWriter: + '''Writer that creates a FITS file. Initialised with the fits object and extention name.''' + + def __init__(self, fits, ext=None): + '''Create a new, uninitialised writer.''' + self.fits = fits + self.ext = ext + + def _append(self, data, names=None): + '''Internal method where the FITS writing is done''' + + if self.ext is None or self.ext not in self.fits: + self.fits.write_table(data, names=names, extname=self.ext) + if self.ext is None: + self.ext = self.fits[-1].get_extnum() + else: + hdu = self.fits[self.ext] + # not using hdu.append here because of incompatibilities + hdu.write(data, names=names, firstrow=hdu.get_nrows()) + + def write(self, data=None, /, **columns): + '''Writes to FITS by calling the internal _append method. + Pass either a positional variable (data) + or multiple named arguments (**columns)''' + + # if data is given, write it as it is + if data is not None: + self._append(data) + + # if keyword arguments are given, treat them as names and columns + if columns: + names, values = list(columns.keys()), list(columns.values()) + self._append(values, names) + + +@contextmanager +def write_catalog(filename, *, ext=None): + """ + Write a catalogue into a FITS file, where *ext* is the optional + name of the extension. To be used as a context manager:: + + # create the catalogue writer + with write_catalog("catalog.fits") as out: + ... + # write catalogue columns RA, DEC, E1, E2, WHT with given arrays + out.write(RA=lon, DEC=lat, E1=eps1, E2=e2, WHT=w) + + .. note:: + Requires the ``fitsio`` package. + + """ + import fitsio + with fitsio.FITS(filename, "rw", clobber=True) as fits: + fits.write(None) + writer = _FitsWriter(fits, ext) + yield writer diff --git a/pyproject.toml b/pyproject.toml index af1e1db7..4cdc575b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,7 +6,7 @@ build-backend = "hatchling.build" name = "glass" description = "Generator for Large Scale Structure" readme = "README.md" -requires-python = ">=3.6" +requires-python = ">=3.8" license = {file = "LICENSE"} maintainers = [ {name = "Nicolas Tessore", email = "n.tessore@ucl.ac.uk"}, @@ -29,6 +29,7 @@ dynamic = ["version"] test = [ "pytest", "scipy", + "fitsio", ] docs = [ "sphinx",