imap-processing 1.0.1__py3-none-any.whl → 1.0.3__py3-none-any.whl

imap_processing/ena_maps/utils/corrections.py

@@ -1,6 +1,8 @@
 """L2 corrections common to multiple IMAP ENA instruments."""
 
+import logging
 from pathlib import Path
+from typing import TypeVar
 
 import numpy as np
 import pandas as pd
@@ -8,11 +10,23 @@ import xarray as xr
 from numpy.polynomial import Polynomial
 from scipy.constants import electron_volt, erg, proton_mass
 
-from imap_processing.ena_maps.ena_maps import LoHiBasePointingSet
+from imap_processing.ena_maps.ena_maps import (
+    LoHiBasePointingSet,
+)
 from imap_processing.ena_maps.utils.coordinates import CoordNames
 from imap_processing.spice import geometry
 from imap_processing.spice.time import ttj2000ns_to_et
 
+logger = logging.getLogger(__name__)
+
+# Tell ruff to ignore ambiguous Greek letters in formulas in this file
+# ruff: noqa: RUF003
+
+# Create a TypeVar to represent the specific class being passed in
+# Bound to LoHiBasePointingSet, meaning it must be LoHiBasePointingSet
+# or a subclass of it
+LoHiBasePsetSubclass = TypeVar("LoHiBasePsetSubclass", bound=LoHiBasePointingSet)
+
 # Physical constants for Compton-Getting correction
 # Units: electron_volt = [J / eV]
 #        erg = [J / erg]
@@ -307,7 +321,9 @@ class PowerLawFluxCorrector:
         return corrected_flux, corrected_flux_stat_unc
 
 
-def _add_spacecraft_velocity_to_pset(pset: LoHiBasePointingSet) -> None:
+def _add_spacecraft_velocity_to_pset(
+    pset: LoHiBasePsetSubclass,
+) -> LoHiBasePsetSubclass:
     """
     Calculate and add spacecraft velocity data to pointing set.
 
@@ -316,6 +332,11 @@ def _add_spacecraft_velocity_to_pset(pset: LoHiBasePointingSet) -> None:
     pset : LoHiBasePointingSet
         Pointing set object to be updated.
 
+    Returns
+    -------
+    pset : LoHiBasePointingSet
+        Pointing set object with spacecraft velocity data added.
+
     Notes
     -----
     Adds the following DataArrays to pset.data:
@@ -342,8 +363,10 @@ def _add_spacecraft_velocity_to_pset(pset: LoHiBasePointingSet) -> None:
     )
     pset.data["sc_direction_vector"] = pset.data["sc_velocity"] / sc_velocity_km_per_sec
 
+    return pset
+
 
-def _add_cartesian_look_direction(pset: LoHiBasePointingSet) -> None:
+def _add_cartesian_look_direction(pset: LoHiBasePsetSubclass) -> LoHiBasePsetSubclass:
     """
     Calculate and add look direction vectors to pointing set.
 
@@ -352,6 +375,11 @@ def _add_cartesian_look_direction(pset: LoHiBasePointingSet) -> None:
     pset : LoHiBasePointingSet
         Pointing set object to be updated.
 
+    Returns
+    -------
+    pset : LoHiBasePointingSet
+        Pointing set object with look direction vectors added.
+
     Notes
     -----
     Adds the following DataArray to pset.data:
@@ -376,11 +404,13 @@ def _add_cartesian_look_direction(pset: LoHiBasePointingSet) -> None:
         dims=[*longitudes.dims, CoordNames.CARTESIAN_VECTOR.value],
     )
 
+    return pset
+
 
 def _calculate_compton_getting_transform(
-    pset: LoHiBasePointingSet,
+    pset: LoHiBasePsetSubclass,
     energy_hf: xr.DataArray,
-) -> None:
+) -> LoHiBasePsetSubclass:
     """
     Apply Compton-Getting transformation to compute ENA source directions.
 
@@ -400,14 +430,24 @@ def _calculate_compton_getting_transform(
     energy_hf : xr.DataArray
         ENA energies in the heliosphere frame in eV.
 
+    Returns
+    -------
+    pset : LoHiBasePointingSet
+        Pointing set object with Compton-Getting related variables added and
+        updated az_el_points.
+
     Notes
     -----
     The algorithm is based on the "Appendix A. The IMAP-Lo Mapping Algorithms"
     document.
     Adds the following DataArrays to pset.data:
     - "energy_sc": ENA energies in spacecraft frame (eV)
-    - "ena_source_hae_longitude": ENA source longitudes in heliosphere frame (degrees)
-    - "ena_source_hae_latitude": ENA source latitudes in heliosphere frame (degrees)
+    - "energy_hf": ENA energies in the heliosphere frame (eV)
+    - "ram_mask": Mask indicating whether ENA source direction is from the ram
+      direction.
+    Updates the following DataArrays in pset.data:
+    - "hae_longitude": ENA source longitudes in heliosphere frame (degrees)
+    - "hae_latitude": ENA source latitudes in heliosphere frame (degrees)
     """
     # Store heliosphere frame energies
     pset.data["energy_hf"] = energy_hf
@@ -455,6 +495,8 @@ def _calculate_compton_getting_transform(
     # Velocity magnitude factor calculation (Equation 62)
     # x_k = (êₛ · û_sc) + sqrt(y² + (êₛ · û_sc)² - 1)
     x = dot_product + np.sqrt(y**2 + dot_product**2 - 1)
+    # Get the dimensions in the right order so that spatial is last
+    x = x.transpose(dot_product.dims[0], y.dims[0], dot_product.dims[1])
 
     # Calculate ENA speed in the spacecraft frame
     # |v⃗_sc| = x_k * U_sc
@@ -504,11 +546,13 @@ def _calculate_compton_getting_transform(
         dims=velocity_vector_helio.dims[:-1],
     )
 
+    return pset
+
 
 def apply_compton_getting_correction(
-    pset: LoHiBasePointingSet,
+    pset: LoHiBasePsetSubclass,
     energy_hf: xr.DataArray,
-) -> None:
+) -> LoHiBasePsetSubclass:
     """
     Apply Compton-Getting correction to a pointing set and update coordinates.
 
@@ -532,6 +576,11 @@ def apply_compton_getting_correction(
         ENA energies in the heliosphere frame in eV. Must be 1D with an
         energy dimension.
 
+    Returns
+    -------
+    pset : LoHiBasePointingSet
+        Updated pointing set object with Compton-Getting related variables added.
+
     Notes
     -----
     This function adds the following variables to the pointing set dataset:
@@ -540,20 +589,153 @@ def apply_compton_getting_correction(
     - "look_direction": Cartesian unit vectors of observation directions
     - "energy_hf": ENA energies in heliosphere frame (eV)
     - "energy_sc": ENA energies in spacecraft frame (eV)
-    - "ena_source_hae_longitude": ENA source longitudes in heliosphere frame (degrees)
-    - "ena_source_hae_latitude": ENA source latitudes in heliosphere frame (degrees)
+    This function modifies the following variables in the pointing set dataset:
+    - "hae_longitude": ENA source longitudes in heliosphere frame (degrees)
+    - "hae_latitude": ENA source latitudes in heliosphere frame (degrees)
 
     The az_el_points attribute is updated to use the corrected coordinates,
     which will be used for subsequent binning operations.
     """
     # Step 1: Add spacecraft velocity and direction to pset
-    _add_spacecraft_velocity_to_pset(pset)
+    pset = _add_spacecraft_velocity_to_pset(pset)
 
     # Step 2: Calculate and add look direction vectors to pset
-    _add_cartesian_look_direction(pset)
+    pset = _add_cartesian_look_direction(pset)
 
     # Step 3: Apply Compton-Getting transformation
-    _calculate_compton_getting_transform(pset, energy_hf)
+    pset = _calculate_compton_getting_transform(pset, energy_hf)
 
     # Step 4: Update az_el_points to use the corrected coordinates
     pset.update_az_el_points()
+
+    return pset
+
+
+def interpolate_map_flux_to_helio_frame(
+    map_ds: xr.Dataset,
+    esa_energies_ev: xr.DataArray,
+    helio_energies_ev: xr.DataArray,
+) -> xr.Dataset:
+    """
+    Interpolate flux from spacecraft frame to heliocentric frame energies.
+
+    This implements the Compton-Getting interpolation step that transforms
+    flux measurements from the spacecraft frame to the heliocentric frame.
+    The algorithm follows these steps:
+    1. For each spatial pixel and energy step, get the spacecraft energy
+    2. Find bounding ESA energy channels for interpolation
+    3. Perform power-law interpolation between bounding channels to spacecraft energy
+    4. Apply energy scaling transformation to heliocentric frame
+
+    Parameters
+    ----------
+    map_ds : xarray.Dataset
+        Map dataset with `energy_sc` data variable containing the spacecraft
+        frame energies for each spatial pixel and ESA energy step.
+    esa_energies_ev : xarray.DataArray
+        The ESA nominal central energies (in eV).
+    helio_energies_ev : xarray.DataArray
+        The heliocentric frame energies to interpolate to (in eV).
+        In practice, these are the same as esa_energies_ev.
+
+    Returns
+    -------
+    map_ds : xarray.Dataset
+        Updated map dataset with interpolated heliocentric frame fluxes.
+    """
+    logger.info("Performing Compton-Getting interpolation to heliocentric frame")
+
+    # Work with xarray DataArrays to handle arbitrary spatial dimensions
+    energy_sc = map_ds["energy_sc"]
+    intensity = map_ds["ena_intensity"]
+    stat_unc = map_ds["ena_intensity_stat_uncert"]
+    sys_err = map_ds["ena_intensity_sys_err"]
+
+    # Step 1: Find bounding ESA energy indices for each position
+    # Use np.searchsorted on flattened array, then reshape back
+    esa_energy_vals = esa_energies_ev.values
+    energy_sc_flat = energy_sc.values.ravel()
+
+    # Find right bound index for each element (vectorized)
+    right_idx_flat = np.searchsorted(esa_energy_vals, energy_sc_flat, side="right")
+    right_idx_flat = np.clip(right_idx_flat, 1, len(esa_energy_vals) - 1)
+    left_idx_flat = right_idx_flat - 1
+
+    # Reshape indices back to match energy_sc shape
+    right_idx = right_idx_flat.reshape(energy_sc.shape)
+    left_idx = left_idx_flat.reshape(energy_sc.shape)
+
+    # Create DataArrays for indices with same dims as energy_sc
+    # Note: we need to avoid coordinate name conflicts when using isel()
+    # The energy dimension should be present in dims but not as a coordinate
+    # since we're using these as indices into the energy dimension
+    # Create coordinates dict without the energy coordinate
+    coords_without_energy = {k: v for k, v in energy_sc.coords.items() if k != "energy"}
+
+    right_idx_da = xr.DataArray(
+        right_idx, dims=energy_sc.dims, coords=coords_without_energy
+    )
+    left_idx_da = xr.DataArray(
+        left_idx, dims=energy_sc.dims, coords=coords_without_energy
+    )
+
+    # Step 2: Extract flux values at bounding energy channels
+    # Use xarray's advanced indexing to get fluxes at left and right indices
+    flux_left = intensity.isel({"energy": left_idx_da})
+    flux_right = intensity.isel({"energy": right_idx_da})
+    stat_unc_left = stat_unc.isel({"energy": left_idx_da})
+    stat_unc_right = stat_unc.isel({"energy": right_idx_da})
+    sys_err_left = sys_err.isel({"energy": left_idx_da})
+
+    # Get energy values at boundaries - select from esa_energies_ev using indices
+    energy_left = esa_energies_ev.isel({"energy": left_idx_da})
+    energy_right = esa_energies_ev.isel({"energy": right_idx_da})
+
+    # Step 3: Perform power-law interpolation to spacecraft energy
+    # slope = log(f_right/f_left) / log(e_right/e_left)
+    # flux_sc = f_left * (energy_sc / e_left)^slope
+    with np.errstate(divide="ignore", invalid="ignore"):
+        # Calculate slope for power-law interpolation
+        slope = np.log(flux_right / flux_left) / np.log(energy_right / energy_left)
+
+        # Interpolate flux using power-law
+        flux_sc = flux_left * ((energy_sc / energy_left) ** slope)
+
+        # Interpolation factor for uncertainty propagation (Equations 75 & 76)
+        unc_factor = np.log(energy_sc / energy_left) / np.log(
+            energy_right / energy_left
+        )
+
+        # Statistical uncertainty propagation (Equation 75):
+        # δJ = J * sqrt((δJ_left/J_left)^2 * (1 + unc_factor^2) + (δJ_right/J_right)^2)
+        stat_unc_sc = flux_sc * np.sqrt(
+            (stat_unc_left / flux_left) ** 2 * (1.0 + unc_factor**2)
+            + (stat_unc_right / flux_right) ** 2
+        )
+
+        # Systematic uncertainty propagation (Equation 76):
+        # σJ^g = σJ^src_kref * (⟨E^s_kref⟩ / E^ESA_kref)^γ_kref * (E^h / ⟨E^s_kref⟩)
+        # Systematic error scales proportionally with flux during power-law
+        # interpolation
+        sys_err_sc = sys_err_left * ((energy_sc / energy_left) ** slope)
+
+    # Step 4: Energy scaling transformation (Liouville theorem)
+    # flux_helio = flux_sc * (helio_energy / energy_sc)
+    # Use xarray broadcasting - helio_energies_ev will broadcast along esa_energy_step
+    with np.errstate(divide="ignore", invalid="ignore"):
+        energy_ratio = helio_energies_ev / energy_sc
+        flux_helio = flux_sc * energy_ratio
+        stat_unc_helio = stat_unc_sc * energy_ratio
+        sys_err_helio = sys_err_sc * energy_ratio
+
+    # Set any location where the value is not finite to NaN (converts +/-inf to NaN)
+    flux_helio = flux_helio.where(np.isfinite(flux_helio), np.nan)
+    stat_unc_helio = stat_unc_helio.where(np.isfinite(stat_unc_helio), np.nan)
+    sys_err_helio = sys_err_helio.where(np.isfinite(sys_err_helio), np.nan)
+
+    # Update the dataset with interpolated values
+    map_ds["ena_intensity"] = flux_helio
+    map_ds["ena_intensity_stat_uncert"] = stat_unc_helio
+    map_ds["ena_intensity_sys_err"] = sys_err_helio
+
+    return map_ds

imap_processing/ena_maps/utils/naming.py

@@ -342,7 +342,9 @@ class MapDescriptor:
         elif frame_str == "gcs":
             return SpiceFrame.IMAP_GCS
         else:
-            raise NotImplementedError("Coordinate frame is not yet implemented.")
+            raise NotImplementedError(
+                f"Coordinate frame {frame_str} is not yet implemented."
+            )
 
     def to_empty_map(
         self,

imap_processing/hi/hi_l1c.py

@@ -29,11 +29,12 @@ from imap_processing.spice.geometry import (
     frame_transform,
     frame_transform_az_el,
 )
+from imap_processing.spice.repoint import get_pointing_times
 from imap_processing.spice.spin import (
     get_instrument_spin_phase,
     get_spin_data,
 )
-from imap_processing.spice.time import ttj2000ns_to_et
+from imap_processing.spice.time import met_to_ttj2000ns, ttj2000ns_to_et
 
 N_SPIN_BINS = 3600
 SPIN_PHASE_BIN_EDGES = np.linspace(0, 1, N_SPIN_BINS + 1)
@@ -101,14 +102,14 @@ 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,
+        de_dataset.ccsds_met.data.mean(),
+        de_dataset.esa_energy_step,
         config_df.cal_prod_config.number_of_products,
         logical_source_parts["sensor"],
     )
-    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_et = ttj2000ns_to_et(pset_dataset.epoch.data[0])
     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))
@@ -121,16 +122,17 @@ def generate_pset_dataset(
 
 
 def empty_pset_dataset(
-    epoch_val: int, l1b_energy_steps: np.ndarray, n_cal_prods: int, sensor_str: str
+    l1b_met: float, l1b_energy_steps: xr.DataArray, 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
+    l1b_met : float
+        Any met from the input L1B DE dataset. This is used to query the
+        repoint-table data to get the start and end times of the pointing.
+    l1b_energy_steps : xarray.DataArray
         The array of esa_energy_step data from the L1B DE product.
     n_cal_prods : int
         Number of calibration products to allocate.
@@ -148,13 +150,18 @@ def empty_pset_dataset(
 
     # preallocate coordinates xr.DataArrays
     coords = dict()
+
+    # Get the Pointing start and end times
+    pointing_mets = get_pointing_times(l1b_met)
+    epochs = met_to_ttj2000ns(np.asarray(pointing_mets))
+
     # epoch coordinate has only 1 entry for pointing set
     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.array([epoch_val], dtype=np.int64),  # TODO: get dtype from cdf attrs?
+        np.array([epochs[0]], dtype=np.int64),
         name="epoch",
         dims=["epoch"],
         attrs=epoch_attrs,
@@ -164,8 +171,12 @@ def empty_pset_dataset(
         "hi_pset_esa_energy_step", check_schema=False
     ).copy()
     dtype = attrs.pop("dtype")
-    # Find the unique, non-zero esa_energy_steps from the L1B data
-    esa_energy_steps = np.array(sorted(set(l1b_energy_steps) - {0}), dtype=dtype)
+    # Find the unique esa_energy_steps from the L1B data
+    # Exclude 0 and FILLVAL
+    esa_energy_steps = np.array(
+        sorted(set(l1b_energy_steps.values) - {0, l1b_energy_steps.attrs["FILLVAL"]}),
+        dtype=dtype,
+    )
     coords["esa_energy_step"] = xr.DataArray(
         esa_energy_steps,
         name="esa_energy_step",
@@ -197,6 +208,15 @@ def empty_pset_dataset(
 
     # Allocate the coordinate label variables
     data_vars = dict()
+    # Generate the epoch_delta variable
+    data_vars["epoch_delta"] = xr.DataArray(
+        np.diff(epochs),
+        name="epoch_delta",
+        dims=["epoch"],
+        attrs=attr_mgr.get_variable_attributes(
+            "hi_pset_epoch_delta", check_schema=False
+        ),
+    )
     # Generate label variables
     data_vars["esa_energy_step_label"] = xr.DataArray(
         coords["esa_energy_step"].values.astype(str),
@@ -253,7 +273,7 @@ def pset_geometry(pset_et: float, sensor_str: str) -> dict[str, xr.DataArray]:
     Returns
     -------
     geometry_vars : dict[str, xarray.DataArray]
-        Keys are variable names and values are data arrays.
+        Keys are variable names, and values are data arrays.
     """
     geometry_vars = create_dataset_variables(
         ["despun_z"], (1, 3), att_manager_lookup_str="hi_pset_{0}"
@@ -571,11 +591,26 @@ def find_second_de_packet_data(l1b_dataset: xr.Dataset) -> xr.Dataset:
     # 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
+    esa_energy_step = epoch_dataset["esa_energy_step"].values
+    # A change in esa_step should indicate the location of the second packet in
+    # each pair of DE packets at an esa_energy_step. In practice, during some
+    # calibration activities, it was observed that the esa_energy_step can change
+    # when the esa_step did not. So, we look for either to change and use the
+    # indices of those changes to identify the second packet in each pair. We
+    # also need to add the last packet index and assume an energy step change
+    # occurs after the last packet.
     second_esa_packet_idx = np.append(
-        np.flatnonzero(np.diff(esa_step) != 0), len(esa_step) - 1
+        np.flatnonzero((np.diff(esa_step) != 0) | (np.diff(esa_energy_step) != 0)),
+        len(esa_step) - 1,
+    )
+    # Remove esa energy steps at 0 - these are calibrations
+    keep_mask = esa_energy_step[second_esa_packet_idx] != 0
+    # Remove esa energy steps at FILLVAL - these are unidentified
+    keep_mask &= (
+        esa_energy_step[second_esa_packet_idx]
+        != l1b_dataset["esa_energy_step"].attrs["FILLVAL"]
     )
-    # Remove esa steps at 0 - these are calibrations
-    second_esa_packet_idx = second_esa_packet_idx[esa_step[second_esa_packet_idx] != 0]
+    second_esa_packet_idx = second_esa_packet_idx[keep_mask]
     # Remove indices where we don't have two consecutive packets at the same ESA
     if second_esa_packet_idx[0] == 0:
         logger.warning(
@@ -584,7 +619,8 @@ def find_second_de_packet_data(l1b_dataset: xr.Dataset) -> xr.Dataset:
         )
         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]
+        esa_energy_step[second_esa_packet_idx - 1]
+        != esa_energy_step[second_esa_packet_idx]
     )
     if missing_esa_pair_mask.any():
         logger.warning(
@@ -629,9 +665,11 @@ def get_de_clock_ticks_for_esa_step(
     # 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()
+    # CCSDS MET has one second resolution, add two to it to make sure it is
+    # greater than the spin start time it ended on. Theotretically, adding
+    # one second should be sufficeint, but in practice, with flight data, adding
+    # two seconds was found to be necessary.
+    end_time_ind = np.flatnonzero(ccsds_met + 2 >= spin_start_mets).max()
 
     # If the minimum absolute difference is greater than 1/2 the spin-phase
     # we have a problem.