Merge pull request #62 from punch-mission/improve-wcs-metadata-handling

Improve wcs metadata handling

punchbowl/data.py

@@ -11,16 +11,21 @@
 import matplotlib as mpl
 import numpy as np
 import pandas as pd
+import sunpy.map
 import yaml
+from astropy.coordinates import GCRS, EarthLocation, SkyCoord
 from astropy.io import fits
 from astropy.io.fits import Header
 from astropy.nddata import StdDevUncertainty
+from astropy.time import Time
 from astropy.wcs import WCS
 from dateutil.parser import parse as parse_datetime
 from ndcube import NDCube
+from sunpy.coordinates import frames, sun
+from sunpy.coordinates.sun import _sun_north_angle_to_z
+from sunpy.map import solar_angular_radius
 
 from punchbowl.errors import MissingMetadataError
-from punchbowl.exceptions import InvalidDataError
 
 _ROOT = os.path.abspath(os.path.dirname(__file__))
 
@@ -62,6 +67,56 @@ def load_spacecraft_def(path: t.Optional[str] = None) -> dict[str, t.Any]:
         return yaml.safe_load(f)
 
 
+def extract_crota_from_wcs(wcs, is_3d=False):
+    index = 1 if is_3d else 0
+    return np.arctan2(wcs.wcs.pc[index+1, index], wcs.wcs.pc[index, index]) * u.rad
+
+
+def calculate_helio_wcs_from_celestial(wcs_celestial, date_obs, data_shape):
+    is_3d = len(data_shape) == 3
+    index = 1 if is_3d else 0
+
+    # we're at the center of the Earth
+    test_loc = EarthLocation.from_geocentric(0, 0, 0, unit=u.m)
+    test_gcrs = SkyCoord(test_loc.get_gcrs(date_obs))
+
+    # follow the SunPy tutorial from here
+    # https://docs.sunpy.org/en/stable/generated/gallery/units_and_coordinates/radec_to_hpc_map.html#sphx-glr-generated-gallery-units-and-coordinates-radec-to-hpc-map-py
+    reference_coord = SkyCoord(wcs_celestial.wcs.crval[index] * u.Unit(wcs_celestial.wcs.cunit[index]),
+                               wcs_celestial.wcs.crval[index+1] * u.Unit(wcs_celestial.wcs.cunit[index+1]),
+                               frame='gcrs',
+                               obstime=date_obs,
+                               obsgeoloc=test_gcrs.cartesian,
+                               obsgeovel=test_gcrs.velocity.to_cartesian(),
+                               distance=test_gcrs.hcrs.distance
+                               )
+
+    reference_coord_arcsec = reference_coord.transform_to(frames.Helioprojective(observer=test_gcrs))
+
+    cdelt1 = (np.abs(wcs_celestial.wcs.cdelt[index]) * u.deg).to(u.arcsec)
+    cdelt2 = (np.abs(wcs_celestial.wcs.cdelt[index+1]) * u.deg).to(u.arcsec)
+
+    geocentric = GCRS(obstime=date_obs)
+    p_angle = _sun_north_angle_to_z(geocentric)
+
+    crota = extract_crota_from_wcs(wcs_celestial, is_3d=is_3d)
+
+    new_header = sunpy.map.make_fitswcs_header(data_shape[index:], reference_coord_arcsec,
+                                               reference_pixel=u.Quantity([wcs_celestial.wcs.crpix[index] - 1,
+                                                                           wcs_celestial.wcs.crpix[index+1] - 1] * u.pixel),
+                                               scale=u.Quantity([cdelt1, cdelt2] * u.arcsec / u.pix),
+                                               rotation_angle=-p_angle-crota,
+                                               observatory='PUNCH',
+                                               projection_code='ARC')
+
+    wcs_helio = WCS(new_header)
+
+    if is_3d:
+        wcs_helio = astropy.wcs.utils.add_stokes_axis_to_wcs(wcs_helio, 0)
+
+    return wcs_helio, p_angle
+
+
 HistoryEntry = namedtuple("HistoryEntry", "datetime, source, comment")
 
 
@@ -611,9 +666,18 @@ def load_template(cls,  # noqa: C901, not too complex
 
     @property
     def sections(self) -> t.List[str]:
-        """returns header keys"""
+        """returns header sections"""
         return list(self._contents.keys())
 
+    @property
+    def fits_keys(self) -> t.List[str]:
+        """returns fits keys in header template"""
+
+        def flatten(xss):
+            return [x for xs in xss for x in xs]
+
+        return flatten([list(self._contents[section_name].keys()) for section_name in self._contents.keys()])
+
     @property
     def history(self) -> History:
         """returns header history"""
@@ -641,7 +705,7 @@ def _validate_key_is_str(key: str) -> None:
         if not isinstance(key, str):
             raise TypeError(f"Keys for NormalizedMetadata must be strings. You provided {type(key)}.")
         if len(key) > 8:
-            raise ValueError("Keys must be <= 8 characters long")
+            raise ValueError(f"Keys must be <= 8 characters long, received {key}")
 
     def __setitem__(self, key: str, value: t.Any) -> None:
         """
@@ -878,7 +942,7 @@ def filename_base(self) -> str:
             "PUNCH_L" + file_level + "_" + type_code + obscode + "_" + date_string
         )
 
-    def write(self, filename: str, overwrite: bool = True) -> None:
+    def write(self, filename: str, overwrite: bool = True, skip_wcs_conversion: bool=False) -> None:
         """Write PUNCHData elements to file
 
         Parameters
@@ -900,7 +964,7 @@ def write(self, filename: str, overwrite: bool = True) -> None:
         """
 
         if filename.endswith(".fits"):
-            self._write_fits(filename, overwrite=overwrite)
+            self._write_fits(filename, overwrite=overwrite, skip_wcs_conversion=skip_wcs_conversion)
         elif filename.endswith((".png", ".jpg", ".jpeg")):
             self._write_ql(filename, overwrite=overwrite)
         else:
@@ -909,7 +973,55 @@ def write(self, filename: str, overwrite: bool = True) -> None:
                 f"Found: {os.path.splitext(filename)[1]}"
             )
 
-    def _write_fits(self, filename: str, overwrite: bool=True) -> None:
+    def construct_wcs_header_fields(self) -> Header:
+        """Computes primary and secondary WCS header cards to add to a data object
+
+        Returns
+        -------
+        Header
+
+        """
+        date_obs = Time(self.meta.datetime)
+
+        celestial_wcs_header = self.wcs.to_header()
+        output_header = astropy.io.fits.Header()
+
+        unused_keys = ['DATE-OBS', 'DATE-BEG', 'DATE-AVG', 'DATE-END', 'DATE',
+                       'MJD-OBS', 'TELAPSE', 'RSUN_REF', 'TIMESYS']
+
+        helio_wcs, p_angle = calculate_helio_wcs_from_celestial(wcs_celestial=self.wcs,
+                                                       date_obs=date_obs,
+                                                       data_shape=self.data.shape)
+
+        helio_wcs_header = helio_wcs.to_header()
+
+        for key in unused_keys:
+            if key in celestial_wcs_header:
+                del celestial_wcs_header[key]
+            if key in helio_wcs_header:
+                del helio_wcs_header[key]
+
+        if self.meta['CTYPE1'] is not None:
+            for key, value in helio_wcs.to_header().items():
+                output_header[key] = value
+        if self.meta['CTYPE1A'] is not None:
+            for key, value in celestial_wcs_header.items():
+                output_header[key + 'A'] = value
+
+        index = 1 if len(self.data.shape) == 3 else 0
+        center_helio_coord = SkyCoord(helio_wcs.wcs.crval[index]*u.deg,
+                                      helio_wcs.wcs.crval[index+1]*u.deg,
+                                      frame=frames.Helioprojective,
+                                      obstime=date_obs,
+                                      observer='earth')
+
+        output_header['RSUN_ARC'] = solar_angular_radius(center_helio_coord).value
+        output_header['SOLAR_EP'] = p_angle.value
+        output_header['CAR_ROT'] = float(sun.carrington_rotation_number(t=date_obs))
+
+        return output_header
+
+    def _write_fits(self, filename: str, overwrite: bool=True, skip_wcs_conversion: bool=False) -> None:
         """Write PUNCHData elements to FITS files
 
         Parameters
@@ -929,26 +1041,28 @@ def _write_fits(self, filename: str, overwrite: bool=True) -> None:
         header = self.meta.to_fits_header()
 
         # update the header with the WCS
-        wcs_header = self.wcs.to_header()
-        for k, v in wcs_header.items():
-            if k in header:
-                header[k] = v
+        if skip_wcs_conversion:
+            wcs_header = self.wcs.to_header()
+        else:
+            wcs_header = self.construct_wcs_header_fields()
+        for key, value in wcs_header.items():
+            if key in header:
+                header[key] = (self.meta[key]._datatype)(value)
+                self.meta[key] = (self.meta[key]._datatype)(value)
 
         hdul_list = []
-
         hdu_dummy = fits.PrimaryHDU()
         hdul_list.append(hdu_dummy)
 
-        hdu_data = fits.CompImageHDU(data=self.data, header=header)
+        hdu_data = fits.CompImageHDU(data=self.data, header=header, name='Primary data array')
         hdul_list.append(hdu_data)
 
         if self.uncertainty is not None:
-            hdu_uncertainty = fits.CompImageHDU(data=self.uncertainty.array)
-            hdu_uncertainty.header["BITPIX"] = 8
+            hdu_uncertainty = fits.CompImageHDU(data=self.uncertainty.array, name='Uncertainty array')
+            hdu_uncertainty.header['BITPIX'] = 8
             # write WCS to uncertainty header
-            for k, v in wcs_header.items():
-                if k in hdu_uncertainty.header:
-                    hdu_uncertainty.header[k] = v
+            for key, value in wcs_header.items():
+                hdu_uncertainty.header[key] = value
             hdul_list.append(hdu_uncertainty)
 
         hdul = fits.HDUList(hdul_list)
@@ -993,27 +1107,34 @@ def _update_statistics(self) -> None:
 
         # TODO - Determine DSATVAL omniheader value in calibrated units for L1+
 
-        if not np.any(self.data) or np.all(np.isnan(self.data)) or np.all(np.isinf(self.data)):
-            raise InvalidDataError("Input data array expected to contain real, non-zero data.")
+        # if not np.any(self.data) or np.all(np.isnan(self.data)) or np.all(np.isinf(self.data)):
+        #     raise InvalidDataError("Input data array expected to contain real, non-zero data.")
 
         self.meta["DATAZER"] = len(np.where(self.data == 0)[0])
 
         self.meta["DATASAT"] = len(np.where(self.data >= self.meta["DSATVAL"].value)[0])
 
         nonzero_data = self.data[np.where(self.data != 0)].flatten()
 
-        self.meta["DATAAVG"] = np.mean(nonzero_data).item()
-        self.meta["DATAMDN"] = np.median(nonzero_data).item()
-        self.meta["DATASIG"] = np.std(nonzero_data).item()
+        if len(nonzero_data) > 0:
+            self.meta["DATAAVG"] = np.nanmean(nonzero_data).item()
+            self.meta["DATAMDN"] = np.nanmedian(nonzero_data).item()
+            self.meta["DATASIG"] = np.nanstd(nonzero_data).item()
+        else:
+            self.meta['DATAAVG'] = -999.0
+            self.meta['DATAMDN'] = -999.0
+            self.meta['DATASIG'] = -999.0
 
         percentile_percentages = [1, 10, 25, 50, 75, 90, 95, 98, 99]
-        percentile_values = np.percentile(nonzero_data, percentile_percentages)
+        percentile_values = np.nanpercentile(nonzero_data, percentile_percentages)
+        if np.any(np.isnan(percentile_values)):  # report nan if any of the values are nan
+            percentile_values = [-999.0 for _ in percentile_percentages]
 
         for percent, value in zip(percentile_percentages, percentile_values):
             self.meta[f"DATAP{percent:02d}"] = value
 
-        self.meta["DATAMIN"] = float(self.data.min().item())
-        self.meta["DATAMAX"] = float(self.data.max().item())
+        self.meta["DATAMIN"] = float(np.nanmin(self.data))
+        self.meta["DATAMAX"] = float(np.nanmax(self.data))
 
     def duplicate_with_updates(self,
                                data: t.Optional[np.ndarray] = None,

punchbowl/data/Level0.yaml

@@ -13,6 +13,15 @@ Level:
     Instrument and Spacecraft State:
         omits: [OBSLAYR2, OBSLAYR3]
     World Coordinate System:
+        overrides:
+            WCSAXES: 2
+            CTYPE1: HPLN-ARC
+            CTYPE2: HPLT-ARC
+            WCSAXESA: 2
+            CTYPE1A: HPLN-ARC
+            CTYPE2A: HPLT-ARC
+        omits: [CRPIX3, PC1_3, PC2_3, PC3_1, PC3_2, PC3_3, CDELT3, CUNIT3, CTYPE3, CRVAL3, CNAME1,
+                CRPIX3A, PC1_3A, PC2_3A, PC3_1A, PC3_2A, PC3_3A, CDELT3A, CUNIT3A, CTYPE3A, CRVAL3A, CNAME1A]
     Camera and Readout State:
     Onboard Image Processing:
     Calibration Data:

punchbowl/data/Level1.yaml

@@ -13,6 +13,15 @@ Level:
     Instrument and Spacecraft State:
         omits: [OBSLAYR2, OBSLAYR3]
     World Coordinate System:
+        overrides:
+            WCSAXES: 2
+            CTYPE1: HPLN-ARC
+            CTYPE2: HPLT-ARC
+            WCSAXESA: 2
+            CTYPE1A: HPLN-ARC
+            CTYPE2A: HPLT-ARC
+        omits: [CRPIX3, PC1_3, PC2_3, PC3_1, PC3_2, PC3_3, CDELT3, CUNIT3, CTYPE3, CRVAL3, CNAME1,
+                CRPIX3A, PC1_3A, PC2_3A, PC3_1A, PC3_2A, PC3_3A, CDELT3A, CUNIT3A, CTYPE3A, CRVAL3A, CNAME1A]
     Camera and Readout State:
     Onboard Image Processing:
     Calibration Data:
@@ -139,9 +148,9 @@ Products:
         kinds: [Calibration]
         overrides:
             NAXIS: 3
-            NAXIS1: 5
+            NAXIS1: 2048
             NAXIS2: 2048
-            NAXIS3: 2048
+            NAXIS3: 5
             TITLE: PUNCH Level-1 {craftname} Flat-field Parameter map (quartic polynomial coefficients)
             OBSTYPE: '{crafttype} flat-field parameter map (quartic polynomial coefficients)'
             TYPECODE: FQ

punchbowl/data/Level2.yaml

@@ -13,6 +13,11 @@ Level:
     Temporal Information:
     Instrument and Spacecraft State:
     World Coordinate System:
+        overrides:
+            CRPIX1: 0.0
+            CDELT1: 1.0
+            CRPIX1A: 0.0
+            CDELT1A: 1.0
     Camera and Readout State:
     Onboard Image Processing:
     Calibration Data:
@@ -29,9 +34,9 @@ Kinds:
     Polarized:
         overrides:
             NAXIS: 3
-            NAXIS1: 3
+            NAXIS1: 4096
             NAXIS2: 4096
-            NAXIS3: 4096
+            NAXIS3: 3
             OBS-MODE: Polar_MZP
             OBSLAYR1: Polar_M
             OBSLAYR2: Polar_Z
@@ -45,7 +50,23 @@ Kinds:
             NAXIS2: 4096
             OBS-MODE: Unpolarized
             POLAR: -999
-        omits: [NAXIS3, FILTER, OBSLAYR1, OBSLAYR2, OBSLAYR3]
+            CTYPE1: 'HPLN-ARC'
+            CTYPE2: 'HPLT-ARC'
+            WCSAXES: 2
+            WCSAXESA: 2
+            CDELT1: 0.0225
+            CDELT2: 0.0225
+            CRPIX1: 2047.5
+            CRPIX2: 2047.5
+            CDELT1A: -0.025
+            CDELT2A: 0.025
+            CRPIX1A: 2047.5
+            CRPIX2A: 2047.5
+            CTYPE1A: "RA---ARC"
+            CTYPE2A: "DEC--ARC"
+        omits: [NAXIS3, FILTER, OBSLAYR1, OBSLAYR2, OBSLAYR3 ,CRPIX3,
+                PC1_3, PC2_3, PC3_1, PC3_2, PC3_3, CDELT3, CUNIT3, CTYPE3, CRVAL3, CNAME1,
+                CRPIX3A, PC1_3A, PC2_3A, PC3_1A, PC3_2A, PC3_3A, CDELT3A, CUNIT3A, CTYPE3A, CRVAL3A, CNAME1A ]
 
     Calibration:
         overrides:
@@ -70,8 +91,8 @@ Products:
     PNN:
         kinds: [Polarized]
         overrides:
+            NAXIS1: 2048
             NAXIS2: 2048
-            NAXIS3: 2048
             TITLE: PUNCH Level-2 Polarized NFI Image
             OBSTYPE: Polarized NFI Image
             TYPECODE: PN