2249 hi lo cg correction bin pset data 1 (IMAP-Science-Operations-Center#2308)

* Change how HiPointingSet implements ram/anti-ram filtering

* Revert changes in HiPointingSet that modify esa_energy_step coordinate

* Start integrating CG correction into Hi L2

* Fix test failing due to unspecified coordinate values

* Remove spin_phase argument to HiPointingSet

* Pre PR cleanup

* Pre PR cleanup

* Give a more descriptive name to TypeVar

* Fix test broken by reordering

* PR feedback

* Fix doc build

Author: tmplummer

imap_processing/ena_maps/ena_maps.py

@@ -647,28 +647,21 @@ class HiPointingSet(LoHiBasePointingSet):
     ----------
     dataset : xarray.Dataset | str | Path
         Hi L1C pointing set data loaded in a xarray.DataArray.
-    spin_phase : str
-        Include ENAs from "full", "ram" or "anti-ram" phases of the spin.
     """
 
-    def __init__(self, dataset: xr.Dataset | str | Path, spin_phase: str):
-        super().__init__(dataset, spice_reference_frame=geometry.SpiceFrame.ECLIPJ2000)
+    def __init__(self, dataset: xr.Dataset | str | Path):
+        super().__init__(dataset, spice_reference_frame=geometry.SpiceFrame.IMAP_HAE)
+
+        self.spatial_coords = ("spin_angle_bin",)
 
-        # Filter out ENAs from non-selected portions of the spin.
-        if spin_phase not in ["full", "ram", "anti"]:
-            raise ValueError(f"Unrecognized spin_phase value: {spin_phase}.")
+        # Naively generate the ram_mask variable assuming spacecraft frame
+        # binning. The ram_mask variable gets updated in the CG correction
+        # code if the CG correction is applied.
+        ram_mask = xr.zeros_like(self.data["spin_angle_bin"], dtype=bool)
         # ram only includes spin-phase interval [0, 0.5)
         # which is the first half of the spin_angle_bins
-        elif spin_phase == "ram":
-            self.data = self.data.isel(
-                spin_angle_bin=slice(0, self.data["spin_angle_bin"].data.size // 2)
-            )
-        # anti-ram includes spin-phase interval [0.5, 1)
-        # which is the second half of the spin_angle_bins
-        elif spin_phase == "anti":
-            self.data = self.data.isel(
-                spin_angle_bin=slice(self.data["spin_angle_bin"].data.size // 2, None)
-            )
+        ram_mask[slice(0, self.data["spin_angle_bin"].data.size // 2)] = True
+        self.data["ram_mask"] = ram_mask
 
         # Rename some PSET vars to match L2 variables
         self.data = self.data.rename(
@@ -684,8 +677,6 @@ def __init__(self, dataset: xr.Dataset | str | Path, spin_phase: str):
             self.data["exposure_factor"], self.data["epoch"].values[0]
         )
 
-        self.spatial_coords = ("spin_angle_bin",)
-
         # Update az_el_points using the base class method
         self.update_az_el_points()
 
@@ -810,12 +801,12 @@ def num_points(self) -> int:
         """
         return self.az_el_points.shape[0]
 
-    def project_pset_values_to_map(
+    def project_pset_values_to_map(  # noqa: PLR0912
         self,
         pointing_set: PointingSet,
         value_keys: list[str] | None = None,
         index_match_method: IndexMatchMethod = IndexMatchMethod.PUSH,
-        pset_valid_mask: NDArray | None = None,
+        pset_valid_mask: NDArray | xr.DataArray | None = None,
     ) -> None:
         """
         Project a pointing set's values to the map grid.
@@ -837,7 +828,7 @@ def project_pset_values_to_map(
         index_match_method : IndexMatchMethod, optional
             The method of index matching to use for all values.
             Default is IndexMatchMethod.PUSH.
-        pset_valid_mask : NDArray, optional
+        pset_valid_mask : xarray.DataArray or NDArray, optional
             A boolean mask of shape (number of pointing set pixels,) indicating
             which pixels in the pointing set should be considered valid for projection.
             If None, all pixels are considered valid. Default is None.
@@ -849,9 +840,9 @@ def project_pset_values_to_map(
         """
         if value_keys is None:
             value_keys = list(pointing_set.data.data_vars.keys())
-        for value_key in value_keys:
-            if value_key not in pointing_set.data.data_vars:
-                raise ValueError(f"Value key {value_key} not found in pointing set.")
+
+        if missing_keys := set(value_keys) - set(pointing_set.data.data_vars):
+            raise KeyError(f"Value keys not found in pointing set: {missing_keys}")
 
         if pset_valid_mask is None:
             pset_valid_mask = np.ones(pointing_set.num_points, dtype=bool)
@@ -876,9 +867,12 @@ def project_pset_values_to_map(
             )
 
         for value_key in value_keys:
+            if value_key not in pointing_set.data.data_vars:
+                raise ValueError(f"Value key {value_key} not found in pointing set.")
+
             # If multiple spatial axes present
             # (i.e (az, el) for rectangular coordinate PSET),
-            # flatten them in the values array to match the raveled indices
+            # stack them into a single coordinate to match the raveled indices
             raveled_pset_data = pointing_set.data[value_key].stack(
                 {CoordNames.GENERIC_PIXEL.value: pointing_set.spatial_coords}
             )
@@ -907,13 +901,22 @@ def project_pset_values_to_map(
                 data_bc, indices_bc = xr.broadcast(
                     raveled_pset_data, matched_indices_push
                 )
+                # If the valid mask is a xr.DataArray, broadcast it to the same shape
+                if isinstance(pset_valid_mask, xr.DataArray):
+                    stacked_valid_mask = pset_valid_mask.stack(
+                        {CoordNames.GENERIC_PIXEL.value: pointing_set.spatial_coords}
+                    )
+                    pset_valid_mask_bc, _ = xr.broadcast(data_bc, stacked_valid_mask)
+                    pset_valid_mask_values = pset_valid_mask_bc.values
+                else:
+                    pset_valid_mask_values = pset_valid_mask
 
                 # Extract numpy arrays for bincount operation
                 pointing_projected_values = map_utils.bin_single_array_at_indices(
                     value_array=data_bc.values,
                     projection_grid_shape=self.binning_grid_shape,
                     projection_indices=indices_bc.values,
-                    input_valid_mask=pset_valid_mask,
+                    input_valid_mask=pset_valid_mask_values,
                 )
                 # TODO: we may need to allow for unweighted/weighted means here by
                 # dividing pointing_projected_values by some binned weights.
@@ -934,10 +937,6 @@ def project_pset_values_to_map(
                 self.data_1d[value_key].values[..., valid_map_mask] += (
                     pointing_projected_values
                 )
-            else:
-                raise NotImplementedError(
-                    "Only PUSH and PULL index matching methods are supported."
-                )
 
         # TODO: The max epoch needs to include the pset duration. Right now it
         #     is just capturing the start epoch. See issue #1747

imap_processing/ena_maps/utils/corrections.py

@@ -1,18 +1,26 @@
 """L2 corrections common to multiple IMAP ENA instruments."""
 
 from pathlib import Path
+from typing import TypeVar
 
 import numpy as np
 import pandas as pd
 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
 
+# 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 +315,9 @@ class or child classes.
         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 +326,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 +357,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 +369,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 +398,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 +424,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 +489,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 +540,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 +570,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 +583,23 @@ 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