imap-processing 0.11.0__py3-none-any.whl → 0.13.0__py3-none-any.whl

imap_processing/hi/l1c/hi_l1c.py

@@ -4,25 +4,45 @@ from __future__ import annotations
 
 import logging
 from pathlib import Path
+from typing import NamedTuple
 
 import numpy as np
 import pandas as pd
 import xarray as xr
+from numpy import typing as npt
+from numpy._typing import NDArray
 
 from imap_processing.cdf.imap_cdf_manager import ImapCdfAttributes
 from imap_processing.cdf.utils import parse_filename_like
-from imap_processing.hi.utils import create_dataset_variables, full_dataarray
+from imap_processing.hi.l1a.science_direct_event import (
+    DE_CLOCK_TICK_S,
+    HALF_CLOCK_TICK_S,
+)
+from imap_processing.hi.utils import (
+    CoincidenceBitmap,
+    create_dataset_variables,
+    full_dataarray,
+    parse_sensor_number,
+)
 from imap_processing.spice.geometry import (
     SpiceFrame,
     frame_transform,
     frame_transform_az_el,
 )
+from imap_processing.spice.spin import (
+    get_instrument_spin_phase,
+    get_spin_data,
+)
 from imap_processing.spice.time import ttj2000ns_to_et
 
+N_SPIN_BINS = 3600
+SPIN_PHASE_BIN_EDGES = np.linspace(0, 1, N_SPIN_BINS + 1)
+SPIN_PHASE_BIN_CENTERS = (SPIN_PHASE_BIN_EDGES[:-1] + SPIN_PHASE_BIN_EDGES[1:]) / 2
+
 logger = logging.getLogger(__name__)
 
 
-def hi_l1c(dependencies: list, data_version: str) -> xr.Dataset:
+def hi_l1c(dependencies: list) -> list[xr.Dataset]:
     """
     High level IMAP-Hi l1c processing function.
 
@@ -36,10 +56,6 @@ def hi_l1c(dependencies: list, data_version: str) -> xr.Dataset:
     dependencies : list
         Input dependencies needed for l1c processing.
 
-    data_version : str
-        Data version to write to CDF files and the Data_version CDF attribute.
-        Should be in the format Vxxx.
-
     Returns
     -------
     l1c_dataset : xarray.Dataset
@@ -56,9 +72,7 @@ def hi_l1c(dependencies: list, data_version: str) -> xr.Dataset:
             "Input dependencies not recognized for l1c pset processing."
         )
 
-    # TODO: revisit this
-    l1c_dataset.attrs["Data_version"] = data_version
-    return l1c_dataset
+    return [l1c_dataset]
 
 
 def generate_pset_dataset(
@@ -88,18 +102,19 @@ def generate_pset_dataset(
     config_df = CalibrationProductConfig.from_csv(calibration_prod_config_path)
 
     pset_dataset = empty_pset_dataset(
+        de_dataset.epoch.data[0],
         de_dataset.esa_energy_step.data,
         config_df.cal_prod_config.number_of_products,
         logical_source_parts["sensor"],
     )
-    # For ISTP, epoch should be the center of the time bin.
-    pset_dataset.epoch.data[0] = np.mean(de_dataset.epoch.data[[0, -1]]).astype(
-        np.int64
-    )
     pset_et = ttj2000ns_to_et(pset_dataset.epoch.data[0])
     # Calculate and add despun_z, hae_latitude, and hae_longitude variables to
     # the pset_dataset
     pset_dataset.update(pset_geometry(pset_et, logical_source_parts["sensor"]))
+    # Bin the counts into the spin-bins
+    pset_dataset.update(pset_counts(pset_dataset.coords, config_df, de_dataset))
+    # Calculate and add the exposure time to the pset_dataset
+    pset_dataset.update(pset_exposure(pset_dataset.coords, de_dataset))
 
     # TODO: The following section will go away as PSET algorithms to populate
     #    these variables are written.
@@ -107,8 +122,6 @@ def generate_pset_dataset(
     attr_mgr.add_instrument_global_attrs("hi")
     attr_mgr.add_instrument_variable_attrs(instrument="hi", level=None)
     for var_name in [
-        "counts",
-        "exposure_times",
         "background_rates",
         "background_rates_uncertainty",
     ]:
@@ -122,13 +135,15 @@ def generate_pset_dataset(
 
 
 def empty_pset_dataset(
-    l1b_energy_steps: np.ndarray, n_cal_prods: int, sensor_str: str
+    epoch_val: int, l1b_energy_steps: np.ndarray, n_cal_prods: int, sensor_str: str
 ) -> xr.Dataset:
     """
     Allocate an empty xarray.Dataset with appropriate pset coordinates.
 
     Parameters
     ----------
+    epoch_val : int
+        The starting epoch in J2000 TT nanoseconds for data in the PSET.
     l1b_energy_steps : np.ndarray
         The array of esa_energy_step data from the L1B DE product.
     n_cal_prods : int
@@ -148,12 +163,12 @@ def empty_pset_dataset(
     # preallocate coordinates xr.DataArrays
     coords = dict()
     # epoch coordinate has only 1 entry for pointing set
-    epoch_attrs = attr_mgr.get_variable_attributes("epoch")
+    epoch_attrs = attr_mgr.get_variable_attributes("epoch", check_schema=False)
     epoch_attrs.update(
         attr_mgr.get_variable_attributes("hi_pset_epoch", check_schema=False)
     )
     coords["epoch"] = xr.DataArray(
-        np.empty(1, dtype=np.int64),  # TODO: get dtype from cdf attrs?
+        np.array([epoch_val], dtype=np.int64),  # TODO: get dtype from cdf attrs?
         name="epoch",
         dims=["epoch"],
         attrs=epoch_attrs,
@@ -188,7 +203,7 @@ def empty_pset_dataset(
     ).copy()
     dtype = attrs.pop("dtype")
     coords["spin_angle_bin"] = xr.DataArray(
-        np.arange(int(360 / 0.1), dtype=dtype),
+        np.arange(N_SPIN_BINS, dtype=dtype),
         name="spin_angle_bin",
         dims=["spin_angle_bin"],
         attrs=attrs,
@@ -271,8 +286,8 @@ def pset_geometry(pset_et: float, sensor_str: str) -> dict[str, xr.DataArray]:
     el = 0 if "90" in sensor_str else -45
     dps_az_el = np.array(
         [
-            np.arange(0.05, 360, 0.1),
-            np.full(3600, el),
+            SPIN_PHASE_BIN_CENTERS * 360,
+            np.full(N_SPIN_BINS, el),
         ]
     ).T
     hae_az_el = frame_transform_az_el(
@@ -282,7 +297,7 @@ def pset_geometry(pset_et: float, sensor_str: str) -> dict[str, xr.DataArray]:
     geometry_vars.update(
         create_dataset_variables(
             ["hae_latitude", "hae_longitude"],
-            (1, 3600),
+            (1, N_SPIN_BINS),
             att_manager_lookup_str="hi_pset_{0}",
         )
     )
@@ -295,6 +310,361 @@ def pset_geometry(pset_et: float, sensor_str: str) -> dict[str, xr.DataArray]:
     return geometry_vars
 
 
+def pset_counts(
+    pset_coords: dict[str, xr.DataArray],
+    config_df: pd.DataFrame,
+    l1b_de_dataset: xr.Dataset,
+) -> dict[str, xr.DataArray]:
+    """
+    Bin direct events into PSET spin-bins.
+
+    Parameters
+    ----------
+    pset_coords : dict[str, xr.DataArray]
+        The PSET coordinates from the xr.Dataset.
+    config_df : pd.DataFrame
+        The calibration product configuration dataframe.
+    l1b_de_dataset : xr.Dataset
+        The L1B dataset for the pointing being processed.
+
+    Returns
+    -------
+    dict[str, xr.DataArray]
+        Dictionary containing new exposure_times DataArray to be added to the PSET
+        dataset.
+    """
+    # Generate exposure time variable filled with zeros
+    counts_var = create_dataset_variables(
+        ["counts"],
+        coords=pset_coords,
+        att_manager_lookup_str="hi_pset_{0}",
+        fill_value=0,
+    )
+
+    # Convert list of DEs to pandas dataframe for ease indexing/filtering
+    de_df = l1b_de_dataset.drop_dims("epoch").to_pandas()
+
+    # Remove DEs not in Goodtimes/angles
+    good_mask = good_time_and_phase_mask(
+        l1b_de_dataset.event_met.values, l1b_de_dataset.spin_phase.values
+    )
+    de_df = de_df[good_mask]
+
+    # The calibration product configuration potentially has different coincidence
+    # types for each ESA and different TOF windows for each calibration product,
+    # esa energy step combination. Because of this we need to filter DEs that
+    # belong to each combo individually.
+    # Loop over the esa_energy_step values first
+    for esa_energy, esa_df in config_df.groupby(level="esa_energy_step"):
+        # Create a mask for all DEs at the current esa_energy_step.
+        # esa_energy_step is recorded for each packet rather than for each DE,
+        # so we use ccsds_index to get the esa_energy_step for each DE
+        esa_mask = (
+            l1b_de_dataset["esa_energy_step"].data[de_df["ccsds_index"].to_numpy()]
+            == esa_energy
+        )
+        # Now loop over the calibration products for the current ESA energy
+        for config_row in esa_df.itertuples():
+            # Remove DEs that are not at the current ESA energy and in the list
+            # of coincidence types for the current calibration product
+            type_mask = de_df["coincidence_type"].isin(
+                config_row.coincidence_type_values
+            )
+            filtered_de_df = de_df[(esa_mask & type_mask)]
+
+            # Use the TOF window mask to remove DEs with TOFs outside the allowed range
+            tof_fill_vals = {
+                f"tof_{detector_pair}": l1b_de_dataset[f"tof_{detector_pair}"].attrs[
+                    "FILLVAL"
+                ]
+                for detector_pair in CalibrationProductConfig.tof_detector_pairs
+            }
+            tof_in_window_mask = get_tof_window_mask(
+                filtered_de_df, config_row, tof_fill_vals
+            )
+            filtered_de_df = filtered_de_df[tof_in_window_mask]
+
+            # Bin remaining DEs into spin-bins
+            i_esa = np.flatnonzero(pset_coords["esa_energy_step"].data == esa_energy)[0]
+            # spin_phase is in the range [0, 1). Multiplying by N_SPIN_BINS and
+            # truncating to an integer gives the correct bin index
+            spin_bin_indices = (
+                filtered_de_df["spin_phase"].to_numpy() * N_SPIN_BINS
+            ).astype(int)
+            # When iterating over rows of a dataframe, the names of the multi-index
+            # are not preserved. Below, `config_row.Index[0]` gets the cal_prod_num
+            # value from the namedtuple representing the dataframe row.
+            np.add.at(
+                counts_var["counts"].data[0, i_esa, config_row.Index[0]],
+                spin_bin_indices,
+                1,
+            )
+    return counts_var
+
+
+def get_tof_window_mask(
+    de_df: pd.DataFrame, prod_config_row: NamedTuple, fill_vals: dict
+) -> NDArray[bool]:
+    """
+    Generate a mask indicating which DEs to keep based on TOF windows.
+
+    Parameters
+    ----------
+    de_df : pd.DataFrame
+        The Direct Event dataframe for the DEs to filter based on the TOF
+        windows.
+    prod_config_row : namedtuple
+        A single row of the prod config dataframe represented as a named tuple.
+    fill_vals : dict
+        A dictionary containing the fill values used in the input DE TOF
+        dataframe values. This value should be derived from the L1B DE CDF
+        TOF variable attributes.
+
+    Returns
+    -------
+    window_mask : np.ndarray
+        A mask with one entry per DE in the input `de_df` indicating which DEs
+        contain TOF values within the windows specified by `prod_config_row`.
+        The mask is intended to directly filter the DE dataframe.
+    """
+    detector_pairs = CalibrationProductConfig.tof_detector_pairs
+    tof_in_window_mask = np.empty((len(detector_pairs), len(de_df)), dtype=bool)
+    for i_pair, detector_pair in enumerate(detector_pairs):
+        low_limit = getattr(prod_config_row, f"tof_{detector_pair}_low")
+        high_limit = getattr(prod_config_row, f"tof_{detector_pair}_high")
+        tof_array = de_df[f"tof_{detector_pair}"].to_numpy()
+        # The TOF in window mask contains True wherever the TOF is within
+        # the configuration low/high bounds OR the FILLVAL is present. The
+        # FILLVAL indicates that the detector pair was not hit. DEs with
+        # the incorrect coincidence_type are already filtered out and this
+        # implementation simplifies combining the tof_in_window_masks in
+        # the next step.
+        tof_in_window_mask[i_pair] = np.logical_or(
+            np.logical_and(low_limit <= tof_array, tof_array <= high_limit),
+            tof_array == fill_vals[f"tof_{detector_pair}"],
+        )
+    return np.all(tof_in_window_mask, axis=0)
+
+
+def pset_exposure(
+    pset_coords: dict[str, xr.DataArray], l1b_de_dataset: xr.Dataset
+) -> dict[str, xr.DataArray]:
+    """
+    Calculate PSET exposure time.
+
+    Parameters
+    ----------
+    pset_coords : dict[str, xr.DataArray]
+        The PSET coordinates from the xr.Dataset.
+    l1b_de_dataset : xr.Dataset
+        The L1B dataset for the pointing being processed.
+
+    Returns
+    -------
+    dict[str, xr.DataArray]
+        Dictionary containing new exposure_times DataArray to be added to the PSET
+        dataset.
+    """
+    # Extract the sensor number (45 or 90) for computing spin phase
+    sensor_number = parse_sensor_number(l1b_de_dataset.attrs["Logical_source"])
+
+    # Generate exposure time variable filled with zeros
+    exposure_var = create_dataset_variables(
+        ["exposure_times"],
+        coords=pset_coords,
+        att_manager_lookup_str="hi_pset_{0}",
+        fill_value=0,
+    )
+
+    # Get a subset of the l1b_de_dataset that contains only the second
+    # of each pair of packets at an ESA step.
+    data_subset = find_second_de_packet_data(l1b_de_dataset)
+
+    # Get the pandas dataframe with spin data
+    spin_df = get_spin_data()
+
+    # Loop over each of the CCSDS data rows that have been identified as the second
+    # packet at an ESA step.
+    # When implementing this, the memory needed to avoid this for loop was computed
+    # and determined to be so large that the for loop is warranted.
+    for _, packet_row in data_subset.groupby("epoch"):
+        clock_tick_mets, clock_tick_weights = get_de_clock_ticks_for_esa_step(
+            packet_row["ccsds_met"].values, spin_df
+        )
+
+        # Clock tick MET times are accumulation "edges". To get the mean spin-phase
+        # for a given clock tick, add 1/2 clock tick and compute spin-phase.
+        spin_phases = np.atleast_1d(
+            get_instrument_spin_phase(
+                clock_tick_mets + HALF_CLOCK_TICK_S,
+                SpiceFrame[f"IMAP_HI_{sensor_number}"],
+            )
+        )
+
+        # Remove ticks not in good times/angles
+        good_mask = good_time_and_phase_mask(clock_tick_mets, spin_phases)
+        spin_phases = spin_phases[good_mask]
+        clock_tick_weights = clock_tick_weights[good_mask]
+
+        # TODO: Account for flyback time. See alg doc section 2.3.5
+
+        # Bin exposure times into spin-phase bins
+        new_exposure_times, _ = np.histogram(
+            spin_phases, bins=SPIN_PHASE_BIN_EDGES, weights=clock_tick_weights
+        )
+        # Accumulate the new exposure times for current esa_step
+        i_esa = np.flatnonzero(
+            pset_coords["esa_energy_step"].values
+            == packet_row["esa_energy_step"].values
+        )[0]
+        exposure_var["exposure_times"].values[:, i_esa] += new_exposure_times
+
+    # Convert exposure clock ticks to seconds
+    exposure_var["exposure_times"].values *= DE_CLOCK_TICK_S
+
+    return exposure_var
+
+
+def find_second_de_packet_data(l1b_dataset: xr.Dataset) -> xr.Dataset:
+    """
+    Find the telemetry entries for the second packet at an ESA step.
+
+    Parameters
+    ----------
+    l1b_dataset : xr.Dataset
+        The L1B Direct Event Dataset for the current pointing.
+
+    Returns
+    -------
+    reduced_dataset : xr.Dataset
+        A dataset containing only the entries for the second packet at an ESA step.
+    """
+    epoch_dataset = l1b_dataset.drop_dims("event_met")
+    # We should get two CCSDS packets per 8-spin ESA step.
+    # Get the indices of the packet before each ESA change.
+    esa_step = epoch_dataset["esa_step"].values
+    second_esa_packet_idx = np.append(
+        np.flatnonzero(np.diff(esa_step) != 0), len(esa_step) - 1
+    )
+    # Remove esa steps at 0 - these are calibrations
+    second_esa_packet_idx = second_esa_packet_idx[esa_step[second_esa_packet_idx] != 0]
+    # Remove indices where we don't have two consecutive packets at the same ESA
+    if second_esa_packet_idx[0] == 0:
+        logger.warning(
+            f"Removing packet 0 with ESA step: {esa_step[0]} from"
+            f"calculation of exposure time due to missing matched pair."
+        )
+        second_esa_packet_idx = second_esa_packet_idx[1:]
+    missing_esa_pair_mask = (
+        esa_step[second_esa_packet_idx - 1] != esa_step[second_esa_packet_idx]
+    )
+    if missing_esa_pair_mask.any():
+        logger.warning(
+            f"Removing {missing_esa_pair_mask.sum()} packets from exposure "
+            f"time calculation due to missing ESA step DE packet pairs."
+        )
+    second_esa_packet_idx = second_esa_packet_idx[~missing_esa_pair_mask]
+    # Reduce the dataset to just the second packet entries
+    data_subset = epoch_dataset.isel(epoch=second_esa_packet_idx)
+    return data_subset
+
+
+def get_de_clock_ticks_for_esa_step(
+    ccsds_met: float, spin_df: pd.DataFrame
+) -> tuple[np.ndarray, np.ndarray]:
+    """
+    Generate an array of clock tick MET times for an 8-spin ESA step.
+
+    Find the closest spin start time in the input spin dataframe to the packet
+    creation time (`ccsds_met`) and generate an array of clock tick MET times
+    for the period covered by the previous 8-spin group and an array of weights
+    that represent the fraction of each clock tick that occurred in the 8-spin
+    group.
+
+    Parameters
+    ----------
+    ccsds_met : float
+        The CCSDS MET of the second packet in a DE packet pair.
+    spin_df : pd.DataFrame
+        Universal spin table dataframe.
+
+    Returns
+    -------
+    clock_tick_mets : np.ndarray
+        Array of MET times that a clock tick occurred in an 8-spin group of spins
+        during which the ESA step was constant.
+    clock_tick_weights : np.ndarray
+        Array of weights to use when binning the clock tick MET times into spin-bins.
+    """
+    # Find the last spin_table entry with the start less than the CCSDS MET.
+    # The CCSDS packet gets created just AFTER the final spin in the 8-spin
+    # ESA step group so this match is the end time. The start time is
+    # 8-spins earlier.
+    spin_start_mets = spin_df.spin_start_met.to_numpy()
+    # CCSDS MET has one second resolution, add one to it to make sure it is
+    # greater than the spin start time it ended on.
+    end_time_ind = np.flatnonzero(ccsds_met + 1 >= spin_start_mets).max()
+
+    # If the minimum absolute difference is greater than 1/2 the spin-phase
+    # we have a problem.
+    if (
+        ccsds_met - spin_start_mets[end_time_ind]
+        > spin_df.iloc[end_time_ind].spin_period_sec / 2
+    ):
+        raise ValueError(
+            "The difference between ccsds_met and spin_start_met, "
+            f"{ccsds_met - spin_start_mets[end_time_ind]} seconds, "
+            f"is too large. Check the spin table loaded for this pointing."
+        )
+    # If the end time index less than 8, we don't have enough spins in the
+    # spin table to get a start time, so raise an error.
+    if end_time_ind < 8:
+        raise ValueError(
+            "Error determining start/end time for exposure time. "
+            f"The CCSDS MET time {ccsds_met} "
+            "is less than 8 spins from the loaded spin table data."
+        )
+    clock_tick_mets = np.arange(
+        spin_start_mets[end_time_ind - 8],
+        spin_start_mets[end_time_ind],
+        DE_CLOCK_TICK_S,
+        dtype=float,
+    )
+    # The final clock-tick bin has less exposure time because the next spin
+    # will trigger FSW to change ESA steps part way through that time. To
+    # account for this in exposure time calculation, assign an array of
+    # weights to use when binnig the clock-ticks to spin-bins. Weights are
+    # fractional clock ticks. All weights are 1 except for the last one in
+    # the array.
+    clock_tick_weights = np.ones_like(clock_tick_mets, dtype=float)
+    clock_tick_weights[-1] = (
+        spin_start_mets[end_time_ind] - clock_tick_mets[-1]
+    ) / DE_CLOCK_TICK_S
+    return clock_tick_mets, clock_tick_weights
+
+
+def good_time_and_phase_mask(
+    tick_mets: np.ndarray, spin_phases: np.ndarray
+) -> npt.NDArray:
+    """
+    Filter out the clock tick times that are not in good times and angles.
+
+    Parameters
+    ----------
+    tick_mets : np.ndarray
+        Clock-tick MET times.
+    spin_phases : np.ndarray
+        Spin phases for each clock tick.
+
+    Returns
+    -------
+    keep_mask : np.ndarray
+        Boolean mask indicating which clock ticks are in good times/phases.
+    """
+    # TODO: Implement this once we have Goodtimes data product defined.
+    return np.full_like(tick_mets, True, dtype=bool)
+
+
 @pd.api.extensions.register_dataframe_accessor("cal_prod_config")
 class CalibrationProductConfig:
     """
@@ -310,21 +680,20 @@ class CalibrationProductConfig:
         "cal_prod_num",
         "esa_energy_step",
     )
+    tof_detector_pairs = ("ab", "ac1", "bc1", "c1c2")
     required_columns = (
         "coincidence_type_list",
-        "tof_ab_low",
-        "tof_ab_high",
-        "tof_ac1_low",
-        "tof_ac1_high",
-        "tof_bc1_low",
-        "tof_bc1_high",
-        "tof_c1c2_low",
-        "tof_c1c2_high",
+        *[
+            f"tof_{det_pair}_{limit}"
+            for det_pair in tof_detector_pairs
+            for limit in ["low", "high"]
+        ],
     )
 
     def __init__(self, pandas_obj: pd.DataFrame) -> None:
         self._validate(pandas_obj)
         self._obj = pandas_obj
+        self._add_coincidence_values_column()
 
     def _validate(self, df: pd.DataFrame) -> None:
         """
@@ -351,6 +720,18 @@ class CalibrationProductConfig:
         # TODO: Verify that the same ESA energy steps exist in all unique calibration
         #   product numbers
 
+    def _add_coincidence_values_column(self) -> None:
+        """Generate and add the coincidence_type_values column to the dataframe."""
+        # Add a column that consists of the coincidence type strings converted
+        # to integer values
+        self._obj["coincidence_type_values"] = self._obj.apply(
+            lambda row: tuple(
+                CoincidenceBitmap.detector_hit_str_to_int(entry)
+                for entry in row["coincidence_type_list"]
+            ),
+            axis=1,
+        )
+
     @classmethod
     def from_csv(cls, path: Path) -> pd.DataFrame:
         """
@@ -366,12 +747,15 @@ class CalibrationProductConfig:
         dataframe : pandas.DataFrame
             Validated calibration product configuration data frame.
         """
-        return pd.read_csv(
+        df = pd.read_csv(
             path,
             index_col=cls.index_columns,
-            converters={"coincidence_type_list": lambda s: s.split("|")},
+            converters={"coincidence_type_list": lambda s: tuple(s.split("|"))},
             comment="#",
         )
+        # Force the _init_ method to run by using the namespace
+        _ = df.cal_prod_config.number_of_products
+        return df
 
     @property
     def number_of_products(self) -> int: