add tests

punch-mission · Nov 2, 2024 · 388d947 · 388d947
1 parent c22a3c5
commit 388d947
Show file tree

Hide file tree

Showing 6 changed files with 73 additions and 44 deletions.
diff --git a/docs/source/example.ipynb b/docs/source/example.ipynb
diff --git a/regularizepsf/psf.py b/regularizepsf/psf.py
@@ -281,8 +281,8 @@ def save(self, path: pathlib.Path) -> None:
             fits.HDUList([fits.PrimaryHDU(),
                           fits.CompImageHDU(np.array(self.coordinates), name="coordinates"),
                           fits.CompImageHDU(self.values, name="values"),
-                          fits.CompImageHDU(self.fft_evaluations.real, name="fft_real"),
-                          fits.CompImageHDU(self.fft_evaluations.imag, name="fft_imag")
+                          fits.CompImageHDU(self.fft_evaluations.real, name="fft_real", quantize_level=32),
+                          fits.CompImageHDU(self.fft_evaluations.imag, name="fft_imag", quantize_level=32),
                           ]).writeto(path)
         else:
             raise NotImplementedError(f"Unsupported file type {path.suffix}. Change to .h5 or .fits.")
@@ -310,7 +310,7 @@ def load(cls, path: pathlib.Path) -> ArrayPSF:
             values_cube = IndexedCube(coordinates, values)
             fft_cube = IndexedCube(coordinates, fft_evaluations)
         elif path.suffix == ".fits":
-            with fits.open(path, "r") as hdul:
+            with fits.open(path) as hdul:
                 coordinates_index = hdul.index_of("coordinates")
                 coordinates = [tuple(c) for c in hdul[coordinates_index].data]
 

diff --git a/regularizepsf/transform.py b/regularizepsf/transform.py
@@ -167,8 +167,10 @@ def save(self, path: pathlib.Path) -> None:
         elif path.suffix == ".fits":
             fits.HDUList([fits.PrimaryHDU(),
                           fits.CompImageHDU(np.array(self.coordinates), name="coordinates"),
-                          fits.CompImageHDU(self._transfer_kernel.values.real, name="transfer_real"),
-                          fits.CompImageHDU(self._transfer_kernel.values.imag, name="transfer_imag")]).writeto(path)
+                          fits.CompImageHDU(self._transfer_kernel.values.real,
+                                            name="transfer_real", quantize_level=32),
+                          fits.CompImageHDU(self._transfer_kernel.values.imag,
+                                            name="transfer_imag", quantize_level=32)]).writeto(path)
         else:
             raise NotImplementedError(f"Unsupported file type {path.suffix}. Change to .h5 or .fits.")
     @classmethod

diff --git a/regularizepsf/util.py b/regularizepsf/util.py
@@ -168,5 +168,5 @@ def __eq__(self, other: IndexedCube) -> bool:
         return (
             self.coordinates == other.coordinates
             and self.sample_shape == other.sample_shape
-            and np.allclose(self.values, other.values)
+            and np.allclose(self.values, other.values, rtol=1e-04, atol=1e-06)
         )
diff --git a/tests/test_psf.py b/tests/test_psf.py
@@ -14,15 +14,16 @@
 from tests.helper import make_gaussian
 
 
-def test_arraypsf_saves_and_loads(tmp_path):
+@pytest.mark.parametrize("extension", ["fits", "h5"])
+def test_arraypsf_saves_and_loads(tmp_path, extension):
     """Can save and reload an ArrayPSF"""
     coordinates = [(0, 0), (1, 1), (2, 2)]
     gauss = make_gaussian(128, fwhm=3)
     values = np.stack([gauss for _ in coordinates])
 
     source = ArrayPSF(IndexedCube(coordinates, values))
 
-    path = tmp_path / "psf.h5"
+    path = tmp_path / f"psf.{extension}"
 
     source.save(path)
     reloaded = ArrayPSF.load(path)

diff --git a/tests/test_transform.py b/tests/test_transform.py
@@ -8,6 +8,24 @@
 from tests.helper import make_gaussian
 
 
+@pytest.mark.parametrize("extension", ["fits", "h5"])
+def test_transform_saves_and_loads(tmp_path, extension):
+    """Can save and reload an ArrayPSF"""
+    coordinates = [(0, 0), (1, 1), (2, 2)]
+    gauss = make_gaussian(128, fwhm=3)
+    values = np.stack([gauss for _ in coordinates])
+
+    source = ArrayPSF(IndexedCube(coordinates, values))
+    target = ArrayPSF(IndexedCube(coordinates, values))
+    transform = ArrayPSFTransform.construct(source, target, 1.0, 0.1)
+
+    path = tmp_path / f"transform.{extension}"
+
+    transform.save(path)
+    reloaded = ArrayPSFTransform.load(path)
+
+    assert transform == reloaded
+
 def test_transform_apply():
     """Test that applying an identity transform does not change the values."""
     size = 256