updated test_starfield.py

punch-mission · Nov 12, 2024 · 6d18d98 · 6d18d98
1 parent d94047d
commit 6d18d98
Showing 1 changed file with 14 additions and 11 deletions.
diff --git a/simpunch/tests/test_starfield.py b/simpunch/tests/test_starfield.py
@@ -12,11 +12,11 @@
 
 @pytest.fixture
 def sample_ndcube() -> NDCube:
-    def _sample_ndcube(shape: tuple, code:str = "PM1", level:str = "0") -> NDCube:
+    def _sample_ndcube(shape: tuple, code: str = "PM1", level: str = "0") -> NDCube:
         data = np.random.random(shape).astype(np.float32)
         sqrt_abs_data = np.sqrt(np.abs(data))
         uncertainty = StdDevUncertainty(np.interp(sqrt_abs_data, (sqrt_abs_data.min(), sqrt_abs_data.max()),
-                                                  (0,1)).astype(np.float32))
+                                                  (0, 1)).astype(np.float32))
         wcs = WCS(naxis=2)
         wcs.wcs.ctype = "HPLN-ARC", "HPLT-ARC"
         wcs.wcs.cunit = "deg", "deg"
@@ -29,11 +29,12 @@ def _sample_ndcube(shape: tuple, code:str = "PM1", level:str = "0") -> NDCube:
         meta["DATE-OBS"] = str(datetime(2024, 2, 22, 16, 0, 1))
         meta["FILEVRSN"] = "1"
         return NDCube(data=data, uncertainty=uncertainty, wcs=wcs, meta=meta)
+
     return _sample_ndcube
 
+
 @pytest.fixture
-def sample_ndcollection() -> NDCollection:
-    shape = (2048, 2048)
+def sample_ndcollection(shape: tuple) -> NDCollection:
     wcs = WCS(naxis=2)
     wcs.wcs.ctype = "HPLN-ARC", "HPLT-ARC"
     wcs.wcs.cunit = "deg", "deg"
@@ -42,24 +43,26 @@ def sample_ndcollection() -> NDCollection:
     wcs.wcs.crval = 1, 1
     wcs.wcs.cname = "HPC lon", "HPC lat"
 
-    input_data = NDCube(np.random.random(shape).astype(np.float32),  wcs=wcs)
+    input_data = NDCube(np.random.random(shape).astype(np.float32), wcs=wcs)
     return NDCollection(
-            [("-60.0 deg", input_data),
-             ("0.0 deg", input_data),
-             ("60.0 deg", input_data)],
-            aligned_axes="all")
+        [("-60.0 deg", input_data),
+         ("0.0 deg", input_data),
+         ("60.0 deg", input_data)],
+        aligned_axes="all")
+
 
 def test_starfield(sample_ndcube: NDCube) -> None:
     """Test starfield generation."""
-    input_data = sample_ndcube((2048,2048))
+    input_data = sample_ndcube((2048, 2048))
 
     output_data = add_starfield_clear(input_data)
 
     assert isinstance(output_data, NDCube)
 
+
 def test_polarized_starfield(sample_ndcollection: NDCollection) -> None:
     """Test polarized starfield generation."""
-    input_data = sample_ndcollection()
+    input_data = sample_ndcollection((2048, 2048))
 
     output_data = add_starfield_polarized(input_data)