ENH(user): Add FITS catalogue writer tool #156

.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

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

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

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",