2341 hi l2 improve ram mask for spacecraft frame products (#2365)

* Add method to HiLoBasePointingSet for calculating the ram pixel mask

* Update CG correction code to call update_ram_mask method

* Move ram_mask calculation to corrections.py module

* Add calculation of ram_mask to hi l2 s/c frame workflow

* self review updates

* remove no longer relavent test

Author: tmplummer

imap_processing/ena_maps/ena_maps.py

@@ -615,9 +615,13 @@ def update_az_el_points(self) -> None:
         """
         Update the az_el_points instance variable with new az/el coordinates.
 
-        The values store in the "hae_longitude" and "hae_latitude" variables
+        The values stored in the "hae_longitude" and "hae_latitude" variables
         are used to construct the azimuth and elevation coordinates.
         """
+        logger.info(
+            "Updating az/el points based on data in hae_longitude and"
+            "hae_latitude variables."
+        )
         # Get lon/lat coordinates, squeeze the epoch dimension and stack along
         # the spatial dimensions. xarray.stack() takes possibly multiple spatial
         # dimensions and reshapes those into a single dimension.
@@ -654,15 +658,6 @@ def __init__(self, dataset: xr.Dataset | str | Path):
 
         self.spatial_coords = ("spin_angle_bin",)
 
-        # 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
-        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(
             {

imap_processing/ena_maps/utils/corrections.py

@@ -321,7 +321,7 @@ class or child classes.
         return corrected_flux, corrected_flux_stat_unc
 
 
-def _add_spacecraft_velocity_to_pset(
+def add_spacecraft_velocity_to_pset(
     pset: LoHiBasePsetSubclass,
 ) -> LoHiBasePsetSubclass:
     """
@@ -344,10 +344,9 @@ def _add_spacecraft_velocity_to_pset(
     - "sc_direction_vector": Spacecraft velocity unit vector with dims ["x_y_z"]
     """
     # Compute ephemeris time (J2000 seconds) of PSET midpoint time
-    # TODO: Use the Pointing midpoint time. Epoch should be start time
-    #     but use it until we can make Lo and Hi PSETs have a consistent
-    #     variable to hold the midpoint time.
-    et = ttj2000ns_to_et(pset.data["epoch"].values[0])
+    et = ttj2000ns_to_et(
+        pset.data["epoch"].values[0] + pset.data["epoch_delta"].values[0] / 2
+    )
     # Get spacecraft state in HAE frame
     sc_state = geometry.imap_state(et, ref_frame=geometry.SpiceFrame.IMAP_HAE)
     sc_velocity_vector = sc_state[3:6]
@@ -532,18 +531,64 @@ def _calculate_compton_getting_transform(
         ena_source_direction_helio[..., 2],
     )
 
+    # Update the PSET ram mask.
+    pset = calculate_ram_mask(pset)
+
+    return pset
+
+
+def calculate_ram_mask(pset: LoHiBasePsetSubclass) -> LoHiBasePsetSubclass:
+    """
+    Calculate the RAM mask using the input spacecraft velocity vector.
+
+    The RAM mask is a boolean array with the same dimensions as what is stored
+    in the "hae_longitude" and "hae_latitude" variables of the dataset.
+
+    Parameters
+    ----------
+    pset : LoHiBasePointingSet
+        Pointing set object. The pset dataset is assumed to have valid
+        "hae_longitude", "hae_latitude", and "sc_direction_vector" variables.
+
+    Returns
+    -------
+    pset : LoHiBasePointingSet
+        Pointing set object with ram_mask variable added.
+    """
+    logger.debug(
+        f"Calculating the RAM mask using input spacecraft direction"
+        f"vector: {pset.data['sc_direction_vector']} and hae coordinates in the"
+        f"dataset hae_longitude and hae_latitude variables."
+    )
+    longitude = pset.data["hae_longitude"]
+    latitude = pset.data["hae_latitude"]
+    spacecraft_direction_vec = pset.data["sc_direction_vector"].values
+    spherical_coords = np.stack(
+        [
+            np.ones_like(longitude.values),
+            longitude.values,
+            latitude.values,
+        ],
+        axis=-1,
+    )
+    cartesian_source_direction = xr.DataArray(
+        geometry.spherical_to_cartesian(spherical_coords),
+        dims=[*longitude.dims, CoordNames.CARTESIAN_VECTOR.value],
+    )
     # For ram/anti-ram filtering we can use the sign of the scalar projection
-    # of the ENA source direction onto the spacecraft velocity vector.
-    # ram_mask = (v⃗_helio · û_sc) >= 0
+    # of the ENA source direction vector (-v⃗_ena) onto the spacecraft velocity
+    # vector.
+    # ram_mask = (-v⃗_ena · û_sc) >= 0
+    # Use Einstein summation for efficient vectorized dot product
     ram_mask = (
         np.einsum(
-            "...i,...i->...", velocity_vector_helio, pset.data["sc_direction_vector"]
+            "...i,...i->...", spacecraft_direction_vec, cartesian_source_direction
         )
         >= 0
     )
     pset.data["ram_mask"] = xr.DataArray(
         ram_mask,
-        dims=velocity_vector_helio.dims[:-1],
+        dims=longitude.dims,
     )
 
     return pset
@@ -597,7 +642,7 @@ def apply_compton_getting_correction(
     which will be used for subsequent binning operations.
     """
     # Step 1: Add spacecraft velocity and direction to pset
-    pset = _add_spacecraft_velocity_to_pset(pset)
+    pset = add_spacecraft_velocity_to_pset(pset)
 
     # Step 2: Calculate and add look direction vectors to pset
     pset = _add_cartesian_look_direction(pset)

imap_processing/hi/hi_l2.py

@@ -13,7 +13,9 @@
 )
 from imap_processing.ena_maps.utils.corrections import (
     PowerLawFluxCorrector,
+    add_spacecraft_velocity_to_pset,
     apply_compton_getting_correction,
+    calculate_ram_mask,
     interpolate_map_flux_to_helio_frame,
 )
 from imap_processing.ena_maps.utils.naming import MapDescriptor
@@ -145,6 +147,9 @@ def generate_hi_map(
             # convert esa nominal central energy from keV to eV
             esa_energy_ev = energy_kev * 1000
             pset = apply_compton_getting_correction(pset, esa_energy_ev)
+        else:
+            pset = add_spacecraft_velocity_to_pset(pset)
+            pset = calculate_ram_mask(pset)
 
         # Multiply variables that need to be exposure time weighted average by
         # exposure factor.