ULTRA extended spin update variables to be ISTP compliant (#2826)

* updates to variables

* ISTP compliant arrays

* istp compliant arrays

* pr comments

* pr comments and logger statements

* fixes from running

* remove print statements

* remove print

Author: lacoak21

imap_processing/cdf/config/imap_ultra_l1b_variable_attrs.yaml

@@ -526,6 +526,7 @@ quality_scattering:
   LABLAXIS: quality_scattering
   # TODO: come back to format
   UNITS: " "
+  DEPEND_0: spin_number
 
 quality_low_voltage:
   <<: *default_uint16
@@ -534,6 +535,7 @@ quality_low_voltage:
   LABLAXIS: quality_low_voltage
   # TODO: come back to format
   UNITS: " "
+  DEPEND_0: spin_number
 
 quality_high_energy:
   <<: *default_uint16
@@ -542,6 +544,7 @@ quality_high_energy:
   LABLAXIS: quality_high_energy
   # TODO: come back to format
   UNITS: " "
+  DEPEND_0: spin_number
 
 quality_statistics:
   <<: *default_uint16
@@ -550,22 +553,30 @@ quality_statistics:
   LABLAXIS: quality_statistics
   # TODO: come back to format
   UNITS: " "
-
+  DEPEND_0: spin_number
 
 energy_range_edges:
-  <<: *default_float64
-  CATDESC: Energy range edges for culling data at l1b.
+  CATDESC: Energy range edges used for culling data.
   DISPLAY_TYPE: no_plot
   FIELDNAM: Energy Range Edges
+  FILLVAL: -1.0e+31
+  FORMAT: F12.4
   LABLAXIS: Energy Edges
   UNITS: keV
   VALIDMIN: 0.0
-  VALIDMAX: 9223372036854775807
+  VALIDMAX: 5000
+  VAR_TYPE: support_data
+  DEPEND_0: energy_range_edges_dim
 
 energy_range_flags:
-  <<: *default_float64
-  CATDESC: Bit flags for each culling energy range
+  CATDESC: Bit flags describing culling energy ranges.
   DISPLAY_TYPE: no_plot
   FIELDNAM: Energy Range Flags
+  FILLVAL: 0
+  FORMAT: I5
   LABLAXIS: Range Flags
-  UNITS: " "
+  UNITS: " "
+  VALIDMIN: 1
+  VALIDMAX: 65534
+  VAR_TYPE: support_data
+  DEPEND_0: energy_range_flags_dim

imap_processing/tests/ultra/unit/conftest.py

@@ -8,6 +8,7 @@
 import xarray as xr
 
 from imap_processing import imap_module_directory
+from imap_processing.ultra.constants import UltraConstants
 from imap_processing.ultra.l0.decom_ultra import (
     process_ultra_cmd_echo,
     process_ultra_energy_rates,
@@ -655,6 +656,15 @@ def mock_goodtimes_dataset():
     intervals, _, _ = build_energy_bins()
     energy_ranges = get_binned_energy_ranges(intervals)
     energy_flags = get_energy_range_flags(energy_ranges)
+
+    energy_flags_padded = np.zeros(UltraConstants.MAX_ENERGY_RANGES, dtype=np.uint16)
+    energy_flags_padded[: len(energy_flags)] = energy_flags
+
+    energy_ranges_padded = np.full(
+        UltraConstants.MAX_ENERGY_RANGE_EDGES, -1.0e31, dtype=np.float32
+    )
+    energy_ranges_padded[: len(energy_ranges)] = energy_ranges
+
     nspins = 100
     flags = 2 ** np.arange(9)
     quality = np.zeros(nspins, dtype=np.uint16)
@@ -664,11 +674,11 @@ def mock_goodtimes_dataset():
     return xr.Dataset(
         {
             "spin_number": ("epoch", np.zeros(nspins)),
-            "energy_range_flags": ("energy_flags", energy_flags),
+            "energy_range_flags": ("energy_flags", energy_flags_padded),
             "quality_low_voltage": ("spin_number", quality),
             "quality_high_energy": ("spin_number", np.zeros(nspins, dtype=np.uint16)),
             "quality_statistics": ("spin_number", np.zeros(nspins, dtype=np.uint16)),
-            "energy_range_edges": ("energy_ranges", energy_ranges),
+            "energy_range_edges": ("energy_ranges", energy_ranges_padded),
             "spin_period": (
                 "spin_number",
                 np.full(nspins, 15),

imap_processing/tests/ultra/unit/test_ultra_l1b.py

@@ -7,6 +7,7 @@
 from imap_processing import imap_module_directory
 from imap_processing.cdf.utils import load_cdf, write_cdf
 from imap_processing.quality_flags import ImapDEOutliersUltraFlags
+from imap_processing.ultra.constants import UltraConstants
 from imap_processing.ultra.l1b.de import FILLVAL_FLOAT32
 from imap_processing.ultra.l1b.ultra_l1b import ultra_l1b
 from imap_processing.ultra.utils.ultra_l1_utils import create_dataset
@@ -65,6 +66,13 @@ def mock_data_l1b_extendedspin_dict():
     quality = np.zeros((2, 3), dtype="uint16")
     # These should be shape: (3,)
     energy_dep_flags = np.zeros(len(spin), dtype="uint16")
+    energy_range_flags = np.zeros(UltraConstants.MAX_ENERGY_RANGES, dtype=np.uint16)
+    energy_range_flags[:5] = 1  # Set first 5 to 1 for testing
+    energy_range_edges = np.ones(
+        UltraConstants.MAX_ENERGY_RANGE_EDGES, dtype=np.float32
+    )
+    energy_range_edges[:4] = [3.0, 5.0, 7.0, 10.0]  # Example values
+    energy_range_edges[4:] = -1.0e31  # Fill remaining with fillval
     data_dict = {
         "epoch": epoch,
         "spin_number": spin,
@@ -74,8 +82,8 @@ def mock_data_l1b_extendedspin_dict():
         "quality_low_voltage": energy_dep_flags,
         "quality_high_energy": energy_dep_flags,
         "quality_statistics": energy_dep_flags,
-        "energy_range_flags": np.ones(5, dtype=np.uint16),
-        "energy_range_edges": np.ones(4, dtype=np.float64),
+        "energy_range_flags": energy_range_flags,
+        "energy_range_edges": energy_range_edges,
     }
     return data_dict
 

imap_processing/tests/ultra/unit/test_ultra_l1b_culling.py

@@ -415,16 +415,23 @@ def test_expand_bin_flags_to_spins(caplog):
 def test_get_energy_and_spin_dependent_rejection_mask():
     """Tests get_energy_and_spin_dependent_rejection_mask function."""
     n_spins = 10
+    energy_range_flags = np.zeros(16, dtype=np.uint16)
+    energy_range_flags[:3] = [2**1, 2**2, 2**3]  # Example flags for 3 energy bins
+    energy_range_edges = np.full(17, -1.0e31, dtype=np.float32)
+    energy_range_edges[:4] = [
+        3,
+        5,
+        7,
+        18,
+    ]  # Example energy bin edges (4 edges = 3 bins)
     goodtimes_dataset = xr.Dataset(
         data_vars={
             "spin_number": np.arange(n_spins),
             "quality_low_voltage": np.full(n_spins, 0),
             "quality_high_energy": np.full(n_spins, 0),
             "quality_statistics": np.full(n_spins, 0),
-            "energy_range_flags": np.array(
-                [2**1, 2**2, 2**3]
-            ),  # Example flags for energy bins
-            "energy_range_edges": np.array([3, 5, 7, 18]),  # Example energy bin edges
+            "energy_range_flags": energy_range_flags,
+            "energy_range_edges": energy_range_edges,
         }
     )
     # update quality flags to test that events get rejected

imap_processing/ultra/constants.py

@@ -216,3 +216,9 @@ class UltraConstants:
     STAT_CULLING_N_ITER = 5
     # Sigma threshold to use for statistical outlier culling.
     STAT_CULLING_STD_THRESHOLD = 0.05
+
+    # Set dimensions for extended spin/goodtime support variables
+    # ISTP requires fixed dimensions, so we set these to the maximum we expect to need
+    # and pad with fill values if we use fewer bins.
+    MAX_ENERGY_RANGES = 16
+    MAX_ENERGY_RANGE_EDGES = MAX_ENERGY_RANGES + 1

imap_processing/ultra/l1b/extendedspin.py

@@ -166,14 +166,28 @@ def calculate_extendedspin(
     extendedspin_dict["quality_hk"] = hk_qf
     extendedspin_dict["quality_instruments"] = inst_qf
     extendedspin_dict["quality_low_voltage"] = voltage_qf  # shape (nspin,)
-    # TODO calculate flags for high energy (SEPS) and statistics culling
-    # Initialize these flags to NONE for now.
     extendedspin_dict["quality_statistics"] = stat_outliers_qf  # shape (nspin,)
     extendedspin_dict["quality_high_energy"] = high_energy_qf  # shape (nspin,)
-    # Add an array of flags for each energy bin. Shape: (n_energy_bins)
-    extendedspin_dict["energy_range_flags"] = energy_bin_flags
-    # Add energy ranges  Shape: (n_energy_bins + 1)
-    extendedspin_dict["energy_range_edges"] = np.array(energy_ranges)
+    # ISTP requires stable dimension sizes, so this field must always remain size 16.
+    # If fewer bins are used, pad the remaining entries with 0.
+    energy_flags = np.full(UltraConstants.MAX_ENERGY_RANGES, 0, dtype=np.uint16)
+    energy_flags[: len(energy_bin_flags)] = energy_bin_flags
+    extendedspin_dict["energy_range_flags"] = energy_flags
+    extendedspin_dict["energy_range_flags_dim"] = np.arange(
+        UltraConstants.MAX_ENERGY_RANGES
+    )
+
+    # Initialize array of energy range edges with fill value, then fill in the valid
+    # energy ranges. Set the length to be the max number of energy bins we expect to
+    # use for culling. The number of edges is one more than the number of bins (17).
+    ranges = np.full(
+        (UltraConstants.MAX_ENERGY_RANGE_EDGES,), FILLVAL_FLOAT32, dtype=np.float32
+    )
+    ranges[: len(energy_ranges)] = energy_ranges
+    extendedspin_dict["energy_range_edges"] = ranges
+    extendedspin_dict["energy_range_edges_dim"] = np.arange(
+        UltraConstants.MAX_ENERGY_RANGE_EDGES
+    )
 
     extendedspin_dataset = create_dataset(extendedspin_dict, name, "l1b")
     return extendedspin_dataset

imap_processing/ultra/l1b/goodtimes.py

@@ -54,9 +54,7 @@ def calculate_goodtimes(extendedspin_dataset: xr.Dataset, name: str) -> xr.Datas
     )
 
     data_dict = extract_data_dict(filtered_dataset)
-
     goodtimes_dataset = create_dataset(data_dict, name, "l1b")
-
     if goodtimes_dataset["spin_number"].size == 0:
         goodtimes_dataset = goodtimes_dataset.drop_dims("spin_number")
         goodtimes_dataset = goodtimes_dataset.expand_dims(spin_number=[FILLVAL_UINT32])

imap_processing/ultra/l1b/ultra_l1b_culling.py

@@ -555,6 +555,8 @@ def get_energy_and_spin_dependent_rejection_mask(
     """
     # Get the ebin flags for each energy bin from the goodtimes dataset.
     energy_range_edges = goodtimes_dataset["energy_range_edges"].values
+    # Filter out fill values from energy_range_edges (negative or zero)
+    energy_range_edges = energy_range_edges[energy_range_edges > 0]
     # Get the quality flag arrays "turned on" for energy dependent culling from the
     # goodtimes dataset.
     flag_arrays = [
@@ -564,6 +566,8 @@ def get_energy_and_spin_dependent_rejection_mask(
     # Initialize all events to not rejected
     rejected = np.zeros_like(energy, dtype=bool)
     ebin_flags = goodtimes_dataset["energy_range_flags"].values
+    # Filter out fill values (0s) from energy_range_flags
+    ebin_flags = ebin_flags[ebin_flags > 0]
     # Get the index of the spin number in the goodtimes dataset for each event
     # all spin numbers should be present in the goodtimes dataset since we have already
     # filtered any events that are not
@@ -670,7 +674,11 @@ def flag_low_voltage(
     # For each low voltage ind, flag the corresponding flag
     quality_flags[lv_spin_inds] = True
 
-    #  TODO add log summary.
+    num_culled: int = np.sum(quality_flags)
+    logger.info(
+        f"Low voltage culling removed {num_culled} spin bins across all energy "
+        f"channels. Voltage threshold: {voltage_threshold} V."
+    )
 
     return quality_flags
 
@@ -748,7 +756,12 @@ def flag_high_energy(
         quality_flags[:, ~mask] = flagged[:, ~mask]
     else:
         quality_flags = flagged
-    # TODO add log summary. E.g Tim's hi goodtimes code
+
+    num_culled: int = np.sum(quality_flags)
+    logger.info(
+        f"High energy culling removed {num_culled} spin bins across {n_energy_bins} "
+        f"energy channels. Energy thresholds: {energy_thresholds.flatten()}, "
+    )
 
     return quality_flags
 
@@ -885,7 +898,7 @@ def flag_statistical_outliers(
                     convergence[e_idx] = True
 
     num_culled: int = np.sum(quality_stats)
-    logger.debug(
+    logger.info(
         f"Statistical culling removed {num_culled} spin bins across {n_energy_bins}"
         f" energy channels. Convergence: {convergence} after "
         f"{iterations} iterations."

imap_processing/ultra/l1c/ultra_l1c_pset_bins.py

@@ -480,14 +480,19 @@ def get_spacecraft_exposure_times(
     spin_periods = goodtimes_dataset["spin_period"].values
     energy_flags = goodtimes_dataset["energy_range_flags"].values
     # only get valid flags for the energy bins we are using at l1c
+    # Filter out fill values (0s) from energy_range_flags
     energy_flags = energy_flags[energy_flags > 0]
+    # Filter out fill values from energy_range_edges
+    energy_range_edges = goodtimes_dataset["energy_range_edges"].values
+    # Remove fill values (negative or zero)
+    energy_range_edges_valid = energy_range_edges[energy_range_edges > 0]
     # Get the quality flag arrays "turned on" for energy dependent culling from the
     # goodtimes dataset.
     flag_arrays = [
         goodtimes_dataset[flag_name].values
         for flag_name in ENERGY_DEPENDENT_SPIN_QUALITY_FLAG_FILTERS
     ]
-    bin_to_range = np.digitize(energy_bins, goodtimes_dataset.energy_range_edges)
+    bin_to_range = np.digitize(energy_bins, energy_range_edges_valid)
     valid_spins = (
         np.bitwise_or.reduce(flag_arrays)[np.newaxis, :] & energy_flags[:, np.newaxis]
     ) == 0

imap_processing/ultra/utils/ultra_l1_utils.py

@@ -97,7 +97,6 @@ def create_dataset(  # noqa: PLR0912
             # epoch coordinate already created with correct attrs
             continue
         elif key == "epoch_delta":
-            # Create epoch_delta variable
             dataset[key] = xr.DataArray(
                 data,
                 dims=["epoch"],
@@ -109,10 +108,16 @@ def create_dataset(  # noqa: PLR0912
             "pixel_index",
             "spin_phase_step",
         ]:
-            # update attrs
+            # update attrs on existing coords
             dataset[key].attrs = cdf_manager.get_variable_attributes(
                 key, check_schema=False
             )
+        elif key in ["energy_range_edges_dim", "energy_range_flags_dim"]:
+            dataset[key] = xr.DataArray(
+                data,
+                dims=[key],
+                attrs=cdf_manager.get_variable_attributes(key, check_schema=False),
+            )
         elif key in velocity_keys:
             dataset[key] = xr.DataArray(
                 data,
@@ -177,24 +182,22 @@ def create_dataset(  # noqa: PLR0912
                 dims=["energy_bin_geometric_mean", "pixel_index"],
                 attrs=cdf_manager.get_variable_attributes(key, check_schema=False),
             )
-        elif key in {
-            "dead_time_ratio",
-        }:
+        elif key in {"dead_time_ratio"}:
             dataset[key] = xr.DataArray(
                 data,
                 dims=["spin_phase_step"],
                 attrs=cdf_manager.get_variable_attributes(key, check_schema=False),
             )
-        elif key in {"energy_range_edges"}:
+        elif key == "energy_range_edges":
             dataset[key] = xr.DataArray(
                 data,
-                dims=["energy_range_edges"],
+                dims=["energy_range_edges_dim"],
                 attrs=cdf_manager.get_variable_attributes(key, check_schema=False),
             )
-        elif key in {"energy_range_flags"}:
+        elif key == "energy_range_flags":
             dataset[key] = xr.DataArray(
                 data,
-                dims=["energy_ranges"],
+                dims=["energy_range_flags_dim"],
                 attrs=cdf_manager.get_variable_attributes(key, check_schema=False),
             )
         else:
@@ -225,7 +228,15 @@ def extract_data_dict(dataset: xr.Dataset) -> dict:
     data_dict.update(
         {
             coord: dataset.coords[coord].values
-            for coord in ("spin_number", "energy_bin_geometric_mean", "epoch")
+            for coord in (
+                "spin_number",
+                "energy_bin_geometric_mean",
+                "epoch",
+                "energy_range_flags_dim",
+                "energy_range_edges_dim",
+                "energy_range_flags",
+                "energy_range_edges",
+            )
             if coord in dataset.coords
         }
     )