Merge pull request #337 from punch-mission/wcs-array-shape

WCS array shape

punchbowl/data/tests/test_history.py

@@ -1,7 +1,6 @@
 from datetime import datetime
 
 import pytest
-from astropy.wcs import WCS
 from pytest import fixture
 
 from punchbowl.data.history import History, HistoryEntry

punchbowl/data/tests/test_io.py

@@ -204,7 +204,7 @@ def test_load_punchdata_with_history(tmpdir):
     os.remove(file_path)
 
 
-def make_empty_distortion_model(num_bins: int, image: np.ndarray) -> (DistortionLookupTable, DistortionLookupTable):
+def make_empty_distortion_model(num_bins: int, image: np.ndarray) -> tuple:
     """ Create an empty distortion table
 
     Parameters

punchbowl/data/tests/test_meta.py

@@ -1,20 +1,10 @@
 import os
 from datetime import datetime
 
-import astropy
-import astropy.units as u
-import numpy as np
 import pytest
-from astropy.coordinates import GCRS, ICRS, EarthLocation, SkyCoord, get_sun
 from astropy.io import fits
-from astropy.nddata import StdDevUncertainty
-from astropy.time import Time
-from astropy.wcs import WCS
-from ndcube import NDCube
-from pytest import fixture
-from sunpy.coordinates import frames
-
-from punchbowl.data.history import History, HistoryEntry
+
+from punchbowl.data.history import HistoryEntry
 from punchbowl.data.meta import MetaField, NormalizedMetadata, load_spacecraft_def
 
 TESTDATA_DIR = os.path.dirname(__file__)
@@ -29,19 +19,19 @@
 def test_metafield_creation_keyword_too_long():
     """ cannot create an invalid metafield"""
     with pytest.raises(ValueError):
-        mf = MetaField("TOO LONG KEYWORD",
-                       "What's up?", 3, int, False, True, -99)
+        MetaField("TOO LONG KEYWORD",
+                  "What's up?", 3, int, False, True, -99)
 
 
 def test_metafield_creation_kinds_do_not_match():
     """the value, default, and kind must all agree"""
     with pytest.raises(TypeError):
-        mf = MetaField("HI",
-                       "What's up?", 3, str, False, True, "hi there")
+        MetaField("HI",
+                  "What's up?", 3, str, False, True, "hi there")
 
     with pytest.raises(TypeError):
-        mf = MetaField("TOO LONG KEYWORD",
-                       "What's up?", "hi there", str, False, True, -99)
+        MetaField("TOO LONG KEYWORD",
+                  "What's up?", "hi there", str, False, True, -99)
 
 
 def test_metafield_update():
@@ -50,7 +40,7 @@ def test_metafield_update():
     assert mf.keyword == "HI"
     assert mf.comment == "What's up?"
     assert mf.value == 3
-    assert mf._datatype == int
+    assert mf._datatype is int
     assert not mf.nullable
     assert mf._mutable
     assert mf.default == -99
@@ -68,7 +58,7 @@ def test_metafield_not_mutable():
     assert mf.keyword == "HI"
     assert mf.comment == "What's up?"
     assert mf.value == 3
-    assert mf._datatype == int
+    assert mf._datatype is int
     assert not mf.nullable
     assert not mf._mutable
     assert mf.default == -99
@@ -83,7 +73,7 @@ def test_metafield_wrong_kind_for_update():
     assert mf.keyword == "HI"
     assert mf.comment == "What's up?"
     assert mf.value == 3
-    assert mf._datatype == int
+    assert mf._datatype is int
     assert not mf.nullable
     assert mf._mutable
     assert mf.default == -99

punchbowl/data/tests/test_units.py

@@ -17,6 +17,7 @@ def test_calculate_image_pixel_area_rotated_input():
         wcs.wcs.crval = 0, 0
         wcs.wcs.crpix = 256, 256
         wcs.wcs.cdelt = 0.05, 0.05
+        wcs.array_shape = shape
 
         area_maps.append(units.calculate_image_pixel_area(wcs, shape))
 
@@ -33,6 +34,7 @@ def test_calculate_image_pixel_area_output_shape():
     wcs.wcs.crval = 0, 0
     wcs.wcs.crpix = 512, 480
     wcs.wcs.cdelt = 0.05, 0.05
+    wcs.array_shape = shape
 
     area_map = units.calculate_image_pixel_area(wcs, shape)
     assert area_map.shape == shape

punchbowl/data/wcs.py

@@ -28,7 +28,7 @@ def extract_crota_from_wcs(wcs: WCS) -> u.deg:
 
 def calculate_helio_wcs_from_celestial(wcs_celestial: WCS,
                                        date_obs: datetime,
-                                       data_shape: tuple[int, int]) -> (WCS, float):
+                                       data_shape: tuple[int, int]) -> tuple[WCS, float]:
     """Calculate the helio WCS from a celestial WCS."""
     is_3d = len(data_shape) == 3
 
@@ -80,13 +80,13 @@ def calculate_helio_wcs_from_celestial(wcs_celestial: WCS,
     return wcs_helio, p_angle
 
 
-def get_sun_ra_dec(dt: datetime) -> (float, float):
+def get_sun_ra_dec(dt: datetime) -> tuple[float, float]:
     """Get the position of the Sun in right ascension and declination."""
     position = get_sun(Time(str(dt), scale="utc"))
     return position.ra.value, position.dec.value
 
 
-def calculate_pc_matrix(crota: float, cdelt: (float, float)) -> np.ndarray:
+def calculate_pc_matrix(crota: float, cdelt: tuple[float, float]) -> np.ndarray:
     """
     Calculate a PC matrix given CROTA and CDELT.
 
@@ -311,15 +311,18 @@ def calculate_celestial_wcs_from_helio(wcs_helio: WCS, date_obs: datetime, data_
     return wcs_celestial
 
 
-def load_trefoil_wcs() -> (astropy.wcs.WCS, (int, int)):
+def load_trefoil_wcs() -> tuple[astropy.wcs.WCS, tuple[int, int]]:
     """Load Level 2 trefoil world coordinate system and shape."""
     trefoil_wcs = WCS(os.path.join(_ROOT, "data", "trefoil_wcs.fits"))
     trefoil_wcs.wcs.ctype = "HPLN-ARC", "HPLT-ARC"  # TODO: figure out why this is necessary, seems like a bug
     trefoil_shape = (4096, 4096)
+
+    trefoil_wcs.array_shape = trefoil_shape
+
     return trefoil_wcs, trefoil_shape
 
 
-def load_quickpunch_mosaic_wcs() -> (astropy.wcs.WCS, (int, int)):
+def load_quickpunch_mosaic_wcs() -> tuple[astropy.wcs.WCS, tuple[int, int]]:
     """Load Level quickPUNCH mosaic world coordinate system and shape."""
     quickpunch_mosaic_shape = (1024, 1024)
     quickpunch_mosaic_wcs = WCS(naxis=2)
@@ -329,10 +332,12 @@ def load_quickpunch_mosaic_wcs() -> (astropy.wcs.WCS, (int, int)):
     quickpunch_mosaic_wcs.wcs.cdelt = 0.045, 0.045
     quickpunch_mosaic_wcs.wcs.ctype = "HPLN-ARC", "HPLT-ARC"
 
+    quickpunch_mosaic_wcs.array_shape = quickpunch_mosaic_shape
+
     return quickpunch_mosaic_wcs, quickpunch_mosaic_shape
 
 
-def load_quickpunch_nfi_wcs() -> (astropy.wcs.WCS, (int, int)):
+def load_quickpunch_nfi_wcs() -> tuple[astropy.wcs.WCS, tuple[int, int]]:
     """Load Level quickPUNCH NFI world coordinate system and shape."""
     quickpunch_nfi_shape = (1024, 1024)
     quickpunch_nfi_wcs = WCS(naxis=2)
@@ -342,4 +347,6 @@ def load_quickpunch_nfi_wcs() -> (astropy.wcs.WCS, (int, int)):
     quickpunch_nfi_wcs.wcs.cdelt = 30 / 3600 * 2, 30 / 3600 * 2
     quickpunch_nfi_wcs.wcs.ctype = "HPLN-TAN", "HPLT-TAN"
 
+    quickpunch_nfi_wcs.array_shape = quickpunch_nfi_shape
+
     return quickpunch_nfi_wcs, quickpunch_nfi_shape

punchbowl/level1/tests/test_deficient_pixel.py

@@ -5,10 +5,9 @@
 from astropy.nddata import StdDevUncertainty
 from astropy.wcs import WCS
 from ndcube import NDCube
-from prefect.logging import disable_run_logger
 
 from punchbowl.data import NormalizedMetadata
-from punchbowl.level1.deficient_pixel import remove_deficient_pixels, remove_deficient_pixels_task
+from punchbowl.level1.deficient_pixel import remove_deficient_pixels
 
 THIS_DIRECTORY = pathlib.Path(__file__).parent.resolve()
 
@@ -34,6 +33,7 @@ def sample_bad_pixel_map(shape: tuple = (2048, 2048), n_bad_pixels: int = 20) ->
     wcs.wcs.cdelt = 0.02, 0.02
     wcs.wcs.crpix = 1024, 1024
     wcs.wcs.crval = 0, 24.75
+    wcs.array_shape = shape
 
     meta = NormalizedMetadata({"TYPECODE": "CL", "LEVEL": "1", "OBSRVTRY": "0", "DATE-OBS": "2008-01-03 08:57:00"})
     return NDCube(data=bad_pixel_map, uncertainty=uncertainty, wcs=wcs, meta=meta)
@@ -56,6 +56,7 @@ def perfect_pixel_map(shape: tuple = (2048, 2048)) -> NDCube:
     wcs.wcs.cdelt = 0.02, 0.02
     wcs.wcs.crpix = 1024, 1024
     wcs.wcs.crval = 0, 24.75
+    wcs.array_shape = shape
 
     meta = NormalizedMetadata({"TYPECODE": "CL", "LEVEL": "1", "OBSRVTRY": "0", "DATE-OBS": "2008-01-03 08:57:00"})
     return NDCube(data=bad_pixel_map, uncertainty=uncertainty, wcs=wcs, meta=meta)
@@ -80,6 +81,7 @@ def one_bad_pixel_map(shape: tuple = (2048, 2048)) -> NDCube:
     wcs.wcs.cdelt = 0.02, 0.02
     wcs.wcs.crpix = 1024, 1024
     wcs.wcs.crval = 0, 24.75
+    wcs.array_shape = shape
 
     meta = NormalizedMetadata({"TYPECODE": "CL", "LEVEL": "1", "OBSRVTRY": "0", "DATE-OBS": "2008-01-03 08:57:00"})
     return NDCube(data=bad_pixel_map, uncertainty=uncertainty, wcs=wcs, meta=meta)
@@ -112,6 +114,7 @@ def nine_bad_pixel_map(shape: tuple = (2048, 2048)) -> NDCube:
     wcs.wcs.cdelt = 0.02, 0.02
     wcs.wcs.crpix = 1024, 1024
     wcs.wcs.crval = 0, 24.75
+    wcs.array_shape = shape
 
     meta = NormalizedMetadata({"TYPECODE": "CL", "LEVEL": "1", "OBSRVTRY": "0", "DATE-OBS": "2008-01-03 08:57:00"})
     return NDCube(data=bad_pixel_map, uncertainty=uncertainty, wcs=wcs, meta=meta)
@@ -134,6 +137,7 @@ def increasing_pixel_data(shape: tuple = (2048, 2048)) -> NDCube:
     wcs.wcs.cdelt = 0.02, 0.02
     wcs.wcs.crpix = 1024, 1024
     wcs.wcs.crval = 0, 24.75
+    wcs.array_shape = shape
 
     meta = NormalizedMetadata({"TYPECODE": "CL", "LEVEL": "1", "OBSRVTRY": "0", "DATE-OBS": "2008-01-03 08:57:00"})
     return NDCube(data=data, uncertainty=uncertainty, wcs=wcs, meta=meta)
@@ -154,6 +158,7 @@ def sample_punchdata(shape: tuple = (2048, 2048)) -> NDCube:
     wcs.wcs.cdelt = 0.02, 0.02
     wcs.wcs.crpix = 1024, 1024
     wcs.wcs.crval = 0, 24.75
+    wcs.array_shape = shape
 
     meta = NormalizedMetadata({"TYPECODE": "CL", "LEVEL": "1", "OBSRVTRY": "0", "DATE-OBS": "2008-01-03 08:57:00"})
     return NDCube(data=data, uncertainty=uncertainty, wcs=wcs, meta=meta)

punchbowl/level1/tests/test_sqrt.py

@@ -28,6 +28,7 @@ def sample_punchdata():
     wcs.wcs.crpix = 1024, 1024
     wcs.wcs.crval = 0, 0
     wcs.wcs.cname = "HPC lon", "HPC lat"
+    wcs.array_shape = data.shape
     meta = NormalizedMetadata.load_template("PM1", "0")
     meta['DATE-OBS'] = str(datetime(2023, 1, 1, 0, 0, 1))