imap-processing 0.17.0__py3-none-any.whl → 0.19.0__py3-none-any.whl

imap_processing/hi/utils.py

@@ -1,12 +1,12 @@
 """IMAP-Hi utils functions."""
 
 import re
-from collections.abc import Sequence
+from collections.abc import Iterable, Sequence
 from dataclasses import dataclass
 from enum import IntEnum
-from typing import Optional, Union
 
 import numpy as np
+import pandas as pd
 import xarray as xr
 
 from imap_processing.cdf.imap_cdf_manager import ImapCdfAttributes
@@ -15,11 +15,13 @@ from imap_processing.cdf.imap_cdf_manager import ImapCdfAttributes
 class HIAPID(IntEnum):
     """Create ENUM for apid."""
 
+    H45_MEMDMP = 740
     H45_APP_NHK = 754
     H45_SCI_CNT = 769
     H45_SCI_DE = 770
     H45_DIAG_FEE = 772
 
+    H90_MEMDMP = 804
     H90_APP_NHK = 818
     H90_SCI_CNT = 833
     H90_SCI_DE = 834
@@ -100,9 +102,9 @@ def parse_sensor_number(full_string: str) -> int:
 def full_dataarray(
     name: str,
     attrs: dict,
-    coords: Optional[dict[str, xr.DataArray]] = None,
-    shape: Optional[Union[int, Sequence[int]]] = None,
-    fill_value: Optional[float] = None,
+    coords: dict[str, xr.DataArray] | None = None,
+    shape: int | Sequence[int] | None = None,
+    fill_value: float | None = None,
 ) -> xr.DataArray:
     """
     Generate an empty xarray.DataArray with appropriate attributes.
@@ -158,9 +160,9 @@ def full_dataarray(
 
 def create_dataset_variables(
     variable_names: list[str],
-    variable_shape: Optional[Union[int, Sequence[int]]] = None,
-    coords: Optional[dict[str, xr.DataArray]] = None,
-    fill_value: Optional[float] = None,
+    variable_shape: int | Sequence[int] | None = None,
+    coords: dict[str, xr.DataArray] | None = None,
+    fill_value: float | None = None,
     att_manager_lookup_str: str = "{0}",
 ) -> dict[str, xr.DataArray]:
     """
@@ -247,3 +249,146 @@ class CoincidenceBitmap(IntEnum):
         matches = re.findall(pattern, detector_hit_str)
         # Sum the integer value assigned to the detector name for each match
         return sum(CoincidenceBitmap[m] for m in matches)
+
+
+class EsaEnergyStepLookupTable:
+    """Class for holding a esa_step to esa_energy lookup table."""
+
+    def __init__(self) -> None:
+        self.df = pd.DataFrame(
+            columns=["start_met", "end_met", "esa_step", "esa_energy_step"]
+        )
+        self._indexed = False
+
+        # Get the FILLVAL from the CDF attribute manager that will be returned
+        # for queries without matches
+        attr_mgr = ImapCdfAttributes()
+        attr_mgr.add_instrument_global_attrs("hi")
+        attr_mgr.add_instrument_variable_attrs(instrument="hi", level=None)
+        var_attrs = attr_mgr.get_variable_attributes(
+            "hi_de_esa_energy_step", check_schema=False
+        )
+        self._fillval = var_attrs["FILLVAL"]
+        self._esa_energy_step_dtype = var_attrs["dtype"]
+
+    def add_entry(
+        self, start_met: float, end_met: float, esa_step: int, esa_energy_step: int
+    ) -> None:
+        """
+        Add a single entry to the lookup table.
+
+        Parameters
+        ----------
+        start_met : float
+            Start mission elapsed time of the time range.
+        end_met : float
+            End mission elapsed time of the time range.
+        esa_step : int
+            ESA step value.
+        esa_energy_step : int
+            ESA energy step value to be stored.
+        """
+        new_row = pd.DataFrame(
+            {
+                "start_met": [start_met],
+                "end_met": [end_met],
+                "esa_step": [esa_step],
+                "esa_energy_step": [esa_energy_step],
+            }
+        )
+        self.df = pd.concat([self.df, new_row], ignore_index=True)
+        self._indexed = False
+
+    def _ensure_indexed(self) -> None:
+        """
+        Create index for faster queries if not already done.
+
+        Notes
+        -----
+        This method sorts the internal DataFrame by start_met and esa_step
+        for improved query performance.
+        """
+        if not self._indexed:
+            # Sort by start_met and esa_step for better query performance
+            self.df = self.df.sort_values(["start_met", "esa_step"]).reset_index(
+                drop=True
+            )
+            self._indexed = True
+
+    def query(
+        self,
+        query_met: float | Iterable[float],
+        esa_step: int | Iterable[float],
+    ) -> float | np.ndarray:
+        """
+        Query MET(s) and esa_step(s) to retrieve esa_energy_step(s).
+
+        Parameters
+        ----------
+        query_met : float or array_like
+            Mission elapsed time value(s) to query.
+            Can be a single float or array-like of floats.
+        esa_step : int or array_like
+            ESA step value(s) to match. Can be a single int or array-like of ints.
+            Must be same type (scalar or array-like) as query_met.
+
+        Returns
+        -------
+        float or numpy.ndarray
+            - If inputs are scalars: returns float (esa_energy_step)
+            - If inputs are array-like: returns numpy array of esa_energy_steps
+              with same length as inputs.
+              Contains FILLVAL for queries with no matches.
+
+        Raises
+        ------
+        ValueError
+            If one input is scalar and the other is array-like, or if both are
+            array-like but have different lengths.
+
+        Notes
+        -----
+        If multiple entries match a query, returns the first match found.
+        """
+        self._ensure_indexed()
+
+        # Check if inputs are scalars
+        is_scalar_met = np.isscalar(query_met)
+        is_scalar_step = np.isscalar(esa_step)
+
+        # Check for mismatched input types
+        if is_scalar_met != is_scalar_step:
+            raise ValueError(
+                "query_met and esa_step must both be scalars or both be array-like"
+            )
+
+        # Convert to arrays for uniform processing
+        query_mets = np.atleast_1d(query_met)
+        esa_steps = np.atleast_1d(esa_step)
+
+        # Ensure both arrays have the same shape
+        if query_mets.shape != esa_steps.shape:
+            raise ValueError(
+                "query_met and esa_step must have the same "
+                "length when both are array-like"
+            )
+
+        results = np.full_like(query_mets, self._fillval)
+
+        # Lookup esa_energy_steps for queries
+        for i, (qm, es) in enumerate(zip(query_mets, esa_steps, strict=False)):
+            mask = (
+                (self.df["start_met"] <= qm)
+                & (self.df["end_met"] >= qm)
+                & (self.df["esa_step"] == es)
+            )
+
+            matches = self.df[mask]
+            if not matches.empty:
+                results[i] = matches["esa_energy_step"].iloc[0]
+
+        # Return scalar for scalar inputs, array for array inputs
+        if is_scalar_met and is_scalar_step:
+            return results.astype(self._esa_energy_step_dtype)[0]
+        else:
+            return results.astype(self._esa_energy_step_dtype)

imap_processing/hit/l0/constants.py

@@ -114,6 +114,9 @@ FLAG_PATTERN = np.array([1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
 # Define size of science frame (num of packets)
 FRAME_SIZE = len(FLAG_PATTERN)
 
+# Mod 10 pattern
+MOD_10_PATTERN = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
+
 # Define the number of bits in the mantissa and exponent for
 # decompressing data
 MANTISSA_BITS = 12

imap_processing/hit/l0/decom_hit.py

@@ -124,9 +124,9 @@ def parse_count_rates(sci_dataset: xr.Dataset) -> None:
         for dim in dims:
             if dim not in sci_dataset.coords:
                 sci_dataset.coords[dim] = xr.DataArray(
-                    np.arange(sci_dataset.sizes[dim], dtype=np.int16)
+                    np.arange(sci_dataset.sizes[dim], dtype=np.uint16)
                     if dim == "gain"
-                    else np.arange(sci_dataset.sizes[dim], dtype=np.int32),
+                    else np.arange(sci_dataset.sizes[dim], dtype=np.uint32),
                     dims=[dim],
                     name=dim,
                 )
@@ -260,12 +260,9 @@ def assemble_science_frames(sci_dataset: xr.Dataset) -> xr.Dataset:
         height event data per valid science frame added as new
         data variables.
     """
-    # TODO: Figure out how to handle partial science frames at the
-    #  beginning and end of CCSDS files. These science frames are split
-    #  across CCSDS files and still need to be processed with packets
-    #  from the previous file. Only discard incomplete science frames
-    #  in the middle of the CCSDS file. The code currently skips all
-    #  incomplete science frames.
+    # TODO: The code currently skips all incomplete science frames.
+    #  Only discard incomplete science frames in the middle of the CCSDS file or
+    #  use fill values?
 
     # Convert sequence flags and counters to NumPy arrays for vectorized operations
     seq_flgs = sci_dataset.seq_flgs.values