SkyCoord for flarePosition

Fits de/serialize hook for converting into fits serializeable icrs frame

Author: nicHoch

stixcore/io/product_processors/fits/processors.py

@@ -1141,7 +1141,8 @@ def write_fits(self, prod, *, version=0):
         elif fitspath_complete.exists():
             logger.warning("Complete Fits file %s exists will be overridden", fitspath.name)
 
-        data = prod.data
+        data = prod.data.copy()
+        prod.on_serialize(data)
 
         primary_header, header_override = self.generate_primary_header(filename, prod, version=version)
 

stixcore/products/level3/flarelist.py

@@ -97,24 +97,12 @@ def add_flare_position(
         # SkyCoord(HeliographicStonyhurst(0 * u.deg, 0 * u.deg))
         # SkyCoord(0 * u.deg, 0 * u.deg, frame=helio_frame)
 
-        n = len(data)
-
-        data["flareposition_obs_hgs_x"] = Column(
-            np.zeros(n, dtype=float) * u.km, description="HeliographicStonyhurst X of observer"
-        )
-        data["flareposition_obs_hgs_y"] = Column(
-            np.zeros(n, dtype=float) * u.km, description="HeliographicStonyhurst Y of observer"
-        )
-        data["flareposition_obs_hgs_z"] = Column(
-            np.zeros(n, dtype=float) * u.km, description="HeliographicStonyhurst Z of observer"
-        )
-        data["flareposition_hp_tx"] = Column(np.zeros(n, dtype=float) * u.arcsec, description="Helioprojective Tx")
-        data["flareposition_hp_ty"] = Column(np.zeros(n, dtype=float) * u.arcsec, description="Helioprojective Ty")
-
         data["anc_ephemeris_path"] = Column(" " * 500, dtype=str, description="TDB")
         data["cpd_path"] = Column(" " * 500, dtype=str, description="TDB")
         data["_position_status"] = Column(False, dtype=bool, description="TDB")
         data["_position_message"] = Column(" " * 500, dtype=str, description="TDB")
+        tx_list, ty_list = [], []
+        solo_x_list, solo_y_list, solo_z_list, peak_time_list = [], [], [], []
 
         to_remove = []
         pass_filter = 0
@@ -127,12 +115,11 @@ def add_flare_position(
         day_asp_ephemeris_cache = dict()
 
         for i, row in enumerate(data):
-            if filter_function(row):  # and i < 200:
+            peak_time = row[peak_time_colname]
+            start_time = row[start_time_colname]
+            end_time = row[end_time_colname]
+            if filter_function(row):  # and i < 60:
                 pass_filter += 1
-                peak_time = row[peak_time_colname]
-                start_time = row[start_time_colname]
-                end_time = row[end_time_colname]
-
                 day = peak_time.to_datetime().date()
 
                 if day in day_asp_ephemeris_cache:
@@ -249,21 +236,54 @@ def add_flare_position(
                     solo = HeliographicStonyhurst(*solo_xyz, obstime=peak_time, representation_type="cartesian")
                     with SphericalScreen(solo, only_off_disk=True):
                         center_hpc = coord.transform_to(Helioprojective(observer=solo))
-
-                    data[i]["flareposition_obs_hgs_x"] = solo_xyz[0].to(u.km)
-                    data[i]["flareposition_obs_hgs_y"] = solo_xyz[1].to(u.km)
-                    data[i]["flareposition_obs_hgs_z"] = solo_xyz[2].to(u.km)
-                    data[i]["flareposition_hp_tx"] = center_hpc.Tx.to(u.arcsec)
-                    data[i]["flareposition_hp_ty"] = center_hpc.Ty.to(u.arcsec)
+                        tx_list.append(center_hpc.Tx)
+                        ty_list.append(center_hpc.Ty)
+                    solo_x_list.append(solo.cartesian.x)
+                    solo_y_list.append(solo.cartesian.y)
+                    solo_z_list.append(solo.cartesian.z)
+                    peak_time_list.append(peak_time)
 
                     data[i]["_position_status"] = True
                     data[i]["_position_message"] = "OK"
                 except Exception as e:
                     data[i]["_position_status"] = False
                     data[i]["_position_message"] = f"Error: {type(e)}"
-                    logger.warn(f"Error calculating flare position for flare at time {start_time} : {end_time}: {e}")
+                    logger.warning(f"Error calculating flare position for flare at time {start_time} : {end_time}: {e}")
+                    tx_list.append(np.nan * u.arcsec)
+                    ty_list.append(np.nan * u.arcsec)
+                    solo_x_list.append(np.nan * u.km)
+                    solo_y_list.append(np.nan * u.km)
+                    solo_z_list.append(np.nan * u.km)
+                    peak_time_list.append(peak_time)
             else:
                 to_remove.append(i)
+                tx_list.append(np.nan * u.arcsec)
+                ty_list.append(np.nan * u.arcsec)
+                solo_x_list.append(np.nan * u.km)
+                solo_y_list.append(np.nan * u.km)
+                solo_z_list.append(np.nan * u.km)
+                peak_time_list.append(peak_time)
+                data[i]["_position_status"] = False
+                data[i]["_position_message"] = "flare did not pass the filter function"
+
+        solo_times = Time(peak_time_list)
+        hgs_coords = SkyCoord(
+            u.Quantity(solo_x_list),
+            u.Quantity(solo_y_list),
+            u.Quantity(solo_z_list),
+            frame=HeliographicStonyhurst(obstime=solo_times),
+            representation_type="cartesian",
+        )
+
+        hp_coords = SkyCoord(
+            u.Quantity(tx_list), u.Quantity(ty_list), frame=Helioprojective(obstime=solo_times, observer=hgs_coords)
+        )
+
+        data["location_hgs"] = hgs_coords
+        # description="Flare location in Heliographic Stonyhurst coordinates"
+
+        data["location_hp"] = hp_coords
+        # description="Flare location in Helioprojective coordinates"
 
         if not keep_all_flares:
             data.remove_rows(to_remove)
@@ -276,6 +296,34 @@ def add_flare_position(
             f"finally {len(data)} flares remaining"
         )
 
+    def on_serialize(self, data):
+        for col_name in ("location_hgs", "location_hp"):
+            if col_name in data.colnames:
+                icrs = data[col_name].icrs
+                icrs_coord = SkyCoord(icrs.ra, icrs.dec, icrs.distance, frame="icrs")
+                col_idx = data.colnames.index(col_name)
+                data.remove_column(col_name)
+                data.add_column(icrs_coord, name=col_name, index=col_idx)
+        s = super()
+        if hasattr(s, "on_serialize"):
+            s.on_serialize(data)
+
+    def on_deserialize(self, data, *, peak_time_colname=None):
+        peak_col = peak_time_colname or self.peak_time_colname
+        if peak_col not in data.colnames:
+            logger.warning(f"on_deserialize: column '{peak_col}' not found, skipping location transform")
+        else:
+            obstime = Time(data[peak_col])
+            if "location_hgs" in data.colnames:
+                data["location_hgs"] = data["location_hgs"].transform_to(HeliographicStonyhurst(obstime=obstime))
+            if "location_hp" in data.colnames:
+                data["location_hp"] = data["location_hp"].transform_to(
+                    Helioprojective(obstime=obstime, observer=data["location_hgs"])
+                )
+        s = super()
+        if hasattr(s, "on_deserialize"):
+            s.on_deserialize(data)
+
 
 class FlareSOOPMixin:
     """_summary_"""
@@ -611,6 +659,7 @@ def __init__(self, *, service_type=0, service_subtype=0, ssid=3, data, month, **
 
         self.name = FlarelistSDCLoc.NAME
         self.ssid = 3
+        self.peak_time_colname = "peak_UTC"
 
     def enhance_from_product(self, in_prod: GenericProduct):
         pass
@@ -628,7 +677,7 @@ def add_flare_position(cls, data, fido_client: STIXClient, *, month=None):
             peak_time_colname="peak_UTC",
             start_time_colname="start_UTC",
             end_time_colname="end_UTC",
-            keep_all_flares=False,
+            keep_all_flares=True,
             month=month,
         )
 
@@ -745,6 +794,7 @@ def __init__(self, *, service_type=0, service_subtype=0, ssid=7, data, month, **
 
         self.name = FlarelistSCLoc.NAME
         self.ssid = 7
+        self.peak_time_colname = "peak_UTC"
 
     def enhance_from_product(self, in_prod: GenericProduct):
         pass

stixcore/products/product.py

@@ -315,6 +315,9 @@ def __call__(self, *args, **kwargs):
                     month=month,
                 )
 
+                if hasattr(p, "on_deserialize") and callable(getattr(p, "on_deserialize")):
+                    p.on_deserialize(p.data)
+
                 # store the old fits header for later reuse
                 if isinstance(p, (L1Mixin, L2Mixin)):
                     p.fits_header = pri_header
@@ -562,6 +565,18 @@ def max_exposure(self):
         # default for FITS HEADER
         return 0.0
 
+    def on_serialize(self, data):
+        """Hook called before writing data to FITS. Mixins override and chain via super()."""
+        s = super()
+        if hasattr(s, "on_serialize"):
+            s.on_serialize(data)
+
+    def on_deserialize(self, data):
+        """Hook called after reading data from FITS. Mixins override and chain via super()."""
+        s = super()
+        if hasattr(s, "on_deserialize"):
+            s.on_deserialize(data)
+
     def find_parent_products(self, root):
         """
         Convenient way to get access to the parent products.

stixcore/products/tests/test_flarelist.py

@@ -0,0 +1,86 @@
+from datetime import date
+
+import numpy as np
+import pytest
+from sunpy.coordinates import HeliographicStonyhurst, Helioprojective
+
+import astropy.units as u
+from astropy.coordinates import SkyCoord
+from astropy.io import fits
+from astropy.table import QTable
+from astropy.time import Time
+
+from stixcore.io.product_processors.fits.processors import FitsL3Processor
+from stixcore.products.level3.flarelist import FlarelistSDCLoc
+from stixcore.products.product import Product
+
+N = 10
+
+
+@pytest.fixture
+def flare_data():
+    peak_times = Time("2022-01-01T12:00:00") + np.arange(N) * 600 * u.s
+
+    hgs_coords = SkyCoord(
+        lon=np.linspace(0, 30, N) * u.deg,
+        lat=np.linspace(-5, 5, N) * u.deg,
+        radius=np.ones(N) * 1.0 * u.AU,
+        frame=HeliographicStonyhurst(obstime=peak_times),
+    )
+
+    hp_coords = SkyCoord(
+        Tx=np.linspace(-300, 300, N) * u.arcsec,
+        Ty=np.linspace(-200, 200, N) * u.arcsec,
+        frame=Helioprojective(obstime=peak_times, observer=hgs_coords),
+    )
+
+    data = QTable()
+    data["peak_UTC"] = peak_times
+    data["start_UTC"] = peak_times - 60 * u.s
+    data["end_UTC"] = peak_times + 60 * u.s
+    data["duration"] = np.ones(N) * 120 * u.s
+    data["lc_peak"] = np.ones((N, 5)) * u.ct / u.s
+    data["location_hgs"] = hgs_coords
+    data["location_hp"] = hp_coords
+
+    return data
+
+
+def test_flarelist_sdcloc_location_roundtrip(flare_data, tmp_path):
+    prod = FlarelistSDCLoc(
+        data=flare_data,
+        month=date(2022, 1, 1),
+        control=QTable(),
+    )
+
+    # minimal header bypasses the Spice-dependent header generation chain
+    header = fits.Header()
+    header["LEVEL"] = "L3"
+    header["STYPE"] = 0
+    header["SSTYPE"] = 0
+    header["SSID"] = 3
+    header["DATE-BEG"] = "2022-01-01T00:00:00"
+    prod.fits_header = header
+
+    # energy/additional_header_keywords are not set for freshly created products
+    prod.energy = None
+    prod._additional_header_keywords = []
+
+    orig_hgs_lon = prod.data["location_hgs"].lon.copy()
+    orig_hgs_lat = prod.data["location_hgs"].lat.copy()
+    orig_hp_tx = prod.data["location_hp"].Tx.copy()
+    orig_hp_ty = prod.data["location_hp"].Ty.copy()
+
+    # write via FitsL3Processor — calls on_serialize internally, prod.data unchanged
+    writer = FitsL3Processor(tmp_path)
+    written_file_name = writer.write_fits(prod)
+    assert len(written_file_name) == 1
+
+    # read back via Product factory — calls on_deserialize internally
+    recovered = Product(written_file_name[0])
+
+    assert isinstance(recovered, FlarelistSDCLoc)
+    assert u.allclose(recovered.data["location_hgs"].lon, orig_hgs_lon, atol=1e-6 * u.deg)
+    assert u.allclose(recovered.data["location_hgs"].lat, orig_hgs_lat, atol=1e-6 * u.deg)
+    assert u.allclose(recovered.data["location_hp"].Tx, orig_hp_tx, atol=1e-3 * u.arcsec)
+    assert u.allclose(recovered.data["location_hp"].Ty, orig_hp_ty, atol=1e-3 * u.arcsec)