Merge pull request #932 from punch-mission/starfield

Properly interpolate starfield models and use correct time windows

Author: jmbhughes

changelog/932.feature.rst

@@ -0,0 +1 @@
+Allows interpolation between starfield models.

punchbowl/auto/flows/level3.py

@@ -91,12 +91,13 @@ def level3_PTM_query_ready_files(session, pipeline_config: dict, reference_time=
                                                      File.file_type == "PI")).order_by(File.date_obs.asc()).all()
     logger.info(f"{len(all_ready_files)} Level 3 PIM files need to be processed.")
 
+    starfield_window = pipeline_config["flows"]["level3_PTM"]["starfield_window"]
+
     actually_ready_files = []
     for f in all_ready_files:
-        # TODO put magic numbers in config
-        valid_starfields = get_valid_starfields(session, f, timedelta_window=timedelta(days=14), file_type="PS")
+        valid_starfields = get_valid_starfields(session, f, timedelta_window=timedelta(days=starfield_window), file_type="PS")
 
-        if len(valid_starfields) >= 1:
+        if len(valid_starfields) >= 2:
             actually_ready_files.append(f)
             if len(actually_ready_files) >= max_n:
                 break
@@ -113,15 +114,21 @@ def level3_PTM_construct_flow_info(level2_files: list[File], level3_file: File,
     state = "planned"
     creation_time = datetime.now()
     priority = pipeline_config["flows"][flow_type]["priority"]["initial"]
+    starfield_window = pipeline_config["flows"][flow_type]["starfield_window"]
+
+    starfield_files = get_valid_starfields(session,
+                                            level2_files[0],
+                                            timedelta_window=timedelta(days=starfield_window),
+                                            file_type="PS")
+
+    before_starfield_path = get_closest_before_file(level2_files[0], starfield_files).filename()
+    after_starfield_path = get_closest_after_file(level2_files[0], starfield_files).filename()
 
-    starfield = get_closest_file(level2_files[0],
-                                 get_valid_starfields(session,
-                                                      level2_files[0],
-                                                      timedelta_window=timedelta(days=14)))
     call_data = json.dumps(
         {
             "data_list": [level2_file.filename() for level2_file in level2_files],
-            "starfield_background_path": starfield.filename(),
+            "before_starfield_path": before_starfield_path,
+            "after_starfield_path": after_starfield_path,
         },
     )
     return Flow(
@@ -166,7 +173,7 @@ def level3_PTM_scheduler_flow(pipeline_config_path=None, session=None, reference
 
 
 def level3_PTM_call_data_processor(call_data: dict, pipeline_config, session=None) -> dict:
-    for key in ["data_list", "starfield_background_path"]:
+    for key in ["data_list", "before_starfield_path", "after_starfield_path"]:
         call_data[key] = file_name_to_full_path(call_data[key], pipeline_config["root"])
     return call_data
 
@@ -408,12 +415,13 @@ def level3_CTM_query_ready_files(session, pipeline_config: dict, reference_time=
                                                      File.file_type == "CI")).order_by(File.date_obs.asc()).all()
     logger.info(f"{len(all_ready_files)} Level 3 CIM files need to be processed.")
 
+    starfield_window = pipeline_config["flows"]["level3_CTM"]["starfield_window"]
+
     actually_ready_files = []
     for f in all_ready_files:
-        # # TODO put magic numbers in config
-        valid_starfields = get_valid_starfields(session, f, timedelta_window=timedelta(days=14), file_type="CS")
+        valid_starfields = get_valid_starfields(session, f, timedelta_window=timedelta(days=starfield_window), file_type="CS")
 
-        if len(valid_starfields) >= 1:
+        if len(valid_starfields) >= 2:
             actually_ready_files.append(f)
             if len(actually_ready_files) >= max_n:
                 break
@@ -430,16 +438,21 @@ def level3_CTM_construct_flow_info(level2_files: list[File], level3_file: File,
     state = "planned"
     creation_time = datetime.now()
     priority = pipeline_config["flows"][flow_type]["priority"]["initial"]
+    starfield_window = pipeline_config["flows"][flow_type]["starfield_window"]
+
+    starfield_files = get_valid_starfields(session,
+                                            level2_files[0],
+                                            timedelta_window=timedelta(days=starfield_window),
+                                            file_type="CS")
+
+    before_starfield_path = get_closest_before_file(level2_files[0], starfield_files).filename()
+    after_starfield_path = get_closest_after_file(level2_files[0], starfield_files).filename()
 
-    starfield = get_closest_file(level2_files[0],
-                                 get_valid_starfields(session,
-                                                      level2_files[0],
-                                                      timedelta_window=timedelta(days=90),
-                                                      file_type="CS"))
     call_data = json.dumps(
         {
             "data_list": [level2_file.filename() for level2_file in level2_files],
-            "starfield_background_path": starfield.filename(),
+            "before_starfield_path": before_starfield_path,
+            "after_starfield_path": after_starfield_path,
         },
     )
     return Flow(
@@ -484,7 +497,7 @@ def level3_CTM_scheduler_flow(pipeline_config_path=None, session=None, reference
 
 
 def level3_CTM_call_data_processor(call_data: dict, pipeline_config, session=None) -> dict:
-    for key in ["data_list" , "starfield_background_path"]:
+    for key in ["data_list" , "before_starfield_path", "after_starfield_path"]:
         call_data[key] = file_name_to_full_path(call_data[key], pipeline_config["root"])
     return call_data
 

punchbowl/level3/stellar.py

@@ -14,6 +14,8 @@
 from prefect import get_run_logger
 from remove_starfield import ImageHolder, ImageProcessor, Starfield
 from remove_starfield.reducers import GaussianReducer
+from reproject import reproject_interp
+from reproject.mosaicking import find_optimal_celestial_wcs
 from scipy.ndimage import percentile_filter
 from scipy.stats import circmean
 from solpolpy import resolve
@@ -25,8 +27,9 @@
     calculate_helio_wcs_from_celestial,
     celestial_north_from_wcs,
 )
+from punchbowl.exceptions import InvalidDataError
 from punchbowl.prefect import punch_flow, punch_task
-from punchbowl.util import average_datetime
+from punchbowl.util import average_datetime, interpolate_data
 
 warnings.filterwarnings("ignore")
 
@@ -334,7 +337,8 @@ def generate_starfield_background(
 
 @punch_task
 def subtract_starfield_background_task(data_object: NDCube,
-                                       starfield_background_path: str | None,
+                                       before_starfield_path: str | None,
+                                       after_starfield_path: str | None,
                                        is_polarized: bool = False) -> NDCube:
     """
     Subtracts a background starfield from an input data frame.
@@ -346,8 +350,10 @@ def subtract_starfield_background_task(data_object: NDCube,
     ----------
     data_object : NDCube
         A NDCube data frame to be background subtracted
-    starfield_background_path : str
-        path to a NDCube background starfield map
+    before_starfield_path : str
+        path to a NDCube background starfield map centered before the observation
+    after_starfield_path : str
+        path to a NDCube background starfield map centered after the observation
     is_polarized : bool
         whether the data is polarized
 
@@ -360,17 +366,69 @@ def subtract_starfield_background_task(data_object: NDCube,
     logger = get_run_logger()
     logger.info("subtract_starfield_background started")
 
-    if starfield_background_path is None:
+    if before_starfield_path is None and after_starfield_path is None:
         output = data_object
         output.meta.history.add_now("LEVEL3-subtract_starfield_background",
                                            "starfield subtraction skipped since path is empty")
+    elif before_starfield_path is None or after_starfield_path is None:
+        raise InvalidDataError("subtract_starfield_background requires two input starfield models.")
     else:
-        star_datacube = load_ndcube_from_fits(starfield_background_path)
-        wcs_celestial = calculate_celestial_wcs_from_helio(star_datacube.wcs)
-        wcs_celestial.wcs.cdelt[0] = wcs_celestial.wcs.cdelt[0] * -1
+        star_datacube_before = load_ndcube_from_fits(before_starfield_path)
+        star_datacube_after = load_ndcube_from_fits(after_starfield_path)
+
+        shape_before = star_datacube_before.data.shape[-2:]
+        shape_after = star_datacube_after.data.shape[-2:]
+
+        wcs_celestial_before = calculate_celestial_wcs_from_helio(star_datacube_before.wcs)
+        wcs_celestial_before.wcs.cdelt[0] = wcs_celestial_before.wcs.cdelt[0] * -1
+
+        wcs_celestial_after = calculate_celestial_wcs_from_helio(star_datacube_after.wcs)
+        wcs_celestial_after.wcs.cdelt[0] = wcs_celestial_after.wcs.cdelt[0] * -1
+
+        # TODO - Test with polarized data...
+        union_wcs, union_shape = find_optimal_celestial_wcs(
+            [(shape_before, wcs_celestial_before),
+            (shape_after,  wcs_celestial_after)],
+            auto_rotate=False, projection="CAR")
+
+        starfield_reprojected_before = reproject_interp(
+            (np.stack([star_datacube_before.data, star_datacube_before.uncertainty.array], axis=0),
+            wcs_celestial_before),
+            union_wcs,
+            shape_out=union_shape,
+            return_footprint=False)
+
+        starfield_reprojected_after = reproject_interp(
+            (np.stack([star_datacube_after.data, star_datacube_after.uncertainty.array], axis=0),
+            wcs_celestial_after),
+            union_wcs,
+            shape_out=union_shape,
+            return_footprint=False)
+
+        starfield_before = NDCube(data=starfield_reprojected_before[0],
+                                uncertainty = StdDevUncertainty(starfield_reprojected_before[1]),
+                                wcs = union_wcs, meta=star_datacube_before.meta)
+        starfield_after = NDCube(data=starfield_reprojected_after[0],
+                                uncertainty = StdDevUncertainty(starfield_reprojected_after[1]),
+                                wcs = union_wcs, meta=star_datacube_after.meta)
+
+        starfield_data_interpolated, starfield_uncert_interpolated = interpolate_data(starfield_before,
+                                                        starfield_after,
+                                                        data_object.meta.datetime,
+                                                        allow_extrapolation=False,
+                                                        and_uncertainty=True,
+                                                        infill_nans=True)
+        # TODO - metadata...
+        star_datacube = NDCube(data=starfield_data_interpolated,
+                            uncertainty=StdDevUncertainty(starfield_uncert_interpolated),
+                            wcs = union_wcs,
+                            meta=star_datacube_before.meta)
+        wcs_celestial = union_wcs
 
         original_mask = data_object.data == 0
 
+        # TODO - Think about where to do the interpolation at this stage...
+        # Is this going to require a change in the subtraction code to avoid more reprojections back and forth?
         if is_polarized:
             starfield_model_m = Starfield(np.stack((star_datacube.data[0], star_datacube.uncertainty.array[0])),
                                           wcs_celestial[0])

punchbowl/util.py

@@ -270,7 +270,7 @@ def nan_gaussian(image: np.ndarray, sigma: float) -> np.ndarray:
 
 def interpolate_data(data_before: NDCube, data_after:NDCube, reference_time: datetime, time_key: str = "DATE-OBS",
                      allow_extrapolation: bool = False, and_uncertainty: bool = False,
-                     ) -> np.ndarray | tuple[np.ndarray, np.ndarray]:
+                     infill_nans: bool = False) -> np.ndarray | tuple[np.ndarray, np.ndarray]:
     """Interpolates between two data objects."""
     before_date = parse_datetime(data_before.meta[time_key].value + " UTC").timestamp()
     after_date = parse_datetime(data_after.meta[time_key].value + " UTC").timestamp()
@@ -296,10 +296,18 @@ def interpolate_data(data_before: NDCube, data_after:NDCube, reference_time: dat
         data_interpolated = ((data_after.data - data_before.data)
                               * (observation_date - before_date) / (after_date - before_date)
                               + data_before.data)
+        if infill_nans:
+            data_before_nan = np.isnan(data_before.data)
+            data_after_nan = np.isnan(data_after.data)
+            data_interpolated[data_before_nan] = data_after.data[data_before_nan]
+            data_interpolated[data_after_nan] = data_before.data[data_after_nan]
         if and_uncertainty:
             uncert_interpolated = ((data_after.uncertainty.array - data_before.uncertainty.array)
                                   * (observation_date - before_date) / (after_date - before_date)
                                   + data_before.uncertainty.array)
+            if infill_nans:
+                uncert_interpolated[data_before_nan] = data_after.uncertainty.array[data_before_nan]
+                uncert_interpolated[data_after_nan] = data_before.uncertainty.array[data_after_nan]
 
     if and_uncertainty:
         return data_interpolated, uncert_interpolated