2434 hi lo l2 pull sc velocity and ram mask functions out of cg correction (#2468)

* Pull add_spacecraft_velocity_to_pset and calculate_ram_mask out of CG correction for better commonality between SF and HF map making

* Update Hi and Lo L2 pipelines to always call add_spacecraft_velocity_to_pset and calculate_ram_mask

* Remove duplicate calculation

Author: subagonsouth

imap_processing/ena_maps/utils/corrections.py

@@ -472,11 +472,12 @@ def _calculate_compton_getting_transform(
 
     # Calculate dot product between look directions and spacecraft direction vector
     # Use Einstein summation for efficient vectorized dot product
+    sc_direction_vector = pset["sc_velocity"] / sc_velocity_km_per_sec
     dot_product = xr.DataArray(
         np.einsum(
             "...i,...i->...",
             pset["look_direction"],
-            pset["sc_direction_vector"],
+            sc_direction_vector,
         ),
         dims=pset["look_direction"].dims[:-1],
     )
@@ -543,9 +544,6 @@ def _calculate_compton_getting_transform(
         ena_source_direction_helio[..., 2],
     )
 
-    # Update the PSET ram mask.
-    pset = calculate_ram_mask(pset)
-
     return pset
 
 
@@ -560,7 +558,7 @@ def calculate_ram_mask(pset: xr.Dataset) -> xr.Dataset:
     ----------
     pset : xarray.Dataset
         Pointing set dataset. The pset dataset is assumed to have valid
-        "hae_longitude", "hae_latitude", and "sc_direction_vector" variables.
+        "hae_longitude", "hae_latitude", and "sc_velocity" variables.
 
     Returns
     -------
@@ -569,12 +567,12 @@ def calculate_ram_mask(pset: xr.Dataset) -> xr.Dataset:
     """
     logger.debug(
         f"Calculating the RAM mask using input spacecraft direction "
-        f"vector: {pset['sc_direction_vector'].values} and hae coordinates in "
+        f"vector: {pset['sc_velocity'].values} and hae coordinates in "
         f"the dataset hae_longitude and hae_latitude variables."
     )
     longitude = pset["hae_longitude"]
     latitude = pset["hae_latitude"]
-    spacecraft_direction_vec = pset["sc_direction_vector"].values
+    spacecraft_velocity = pset["sc_velocity"].values
     spherical_coords = np.stack(
         [
             np.ones_like(longitude.values),
@@ -593,9 +591,7 @@ def calculate_ram_mask(pset: xr.Dataset) -> xr.Dataset:
     # ram_mask = (-v⃗_ena · û_sc) >= 0
     # Use Einstein summation for efficient vectorized dot product
     ram_mask = (
-        np.einsum(
-            "...i,...i->...", spacecraft_direction_vec, cartesian_source_direction
-        )
+        np.einsum("...i,...i->...", spacecraft_velocity, cartesian_source_direction)
         >= 0
     )
     pset["ram_mask"] = xr.DataArray(
@@ -630,7 +626,9 @@ def apply_compton_getting_correction(
         Pointing set dataset. Must contain the following coordinates:
           - epoch: start time of the pointing
         Must contain the following variables:
-          - epoch_delta: duration of the pointing in nanoseconds
+          - sc_velocity: velocity vector of the spacecraft in the HAE frame at
+            the midpoint time of the pointing [km/s]. See the
+            `add_spacecraft_velocity_to_pset` function.
           - hae_longitude: PSET bin longitudes in the HAE frame (degrees)
           - hae_latitude: PSET bin latitudes in the HAE frame (degrees)
     energy_hf : xr.DataArray
@@ -647,22 +645,17 @@ def apply_compton_getting_correction(
     Notes
     -----
     This function adds the following variables to the dataset:
-    - "sc_velocity": Spacecraft velocity vector (km/s)
-    - "sc_direction_vector": Spacecraft velocity unit vector
     - "look_direction": Cartesian unit vectors of observation directions
     - "energy_hf": ENA energies in heliosphere frame (eV)
     - "energy_sc": ENA energies in spacecraft frame (eV)
     This function modifies the following variables in the dataset:
     - "hae_longitude": ENA source longitudes in heliosphere frame (degrees)
     - "hae_latitude": ENA source latitudes in heliosphere frame (degrees)
     """
-    # Step 1: Add spacecraft velocity and direction to pset
-    processed_dataset = add_spacecraft_velocity_to_pset(pset)
-
-    # Step 2: Calculate and add look direction vectors to pset
-    processed_dataset = _add_cartesian_look_direction(processed_dataset)
+    # Step 1: Calculate and add look direction vectors to pset
+    processed_dataset = _add_cartesian_look_direction(pset)
 
-    # Step 3: Apply Compton-Getting transformation
+    # Step 2: Apply Compton-Getting transformation
     processed_dataset = _calculate_compton_getting_transform(
         processed_dataset, energy_hf
     )

imap_processing/hi/hi_l2.py

@@ -152,13 +152,14 @@ def generate_hi_map(
                 "All PSETs must have the same set of esa_energy_step values."
             )
 
+        pset_ds = add_spacecraft_velocity_to_pset(pset_ds)
+
         if descriptor.frame_descriptor == "hf":
             # convert esa nominal central energy from keV to eV
             esa_energy_ev = energy_kev * 1000
             pset_ds = apply_compton_getting_correction(pset_ds, esa_energy_ev)
-        else:
-            pset_ds = add_spacecraft_velocity_to_pset(pset_ds)
-            pset_ds = calculate_ram_mask(pset_ds)
+
+        pset_ds = calculate_ram_mask(pset_ds)
 
         # Multiply variables that need to be exposure time weighted average by
         # exposure factor.

imap_processing/lo/l2/lo_l2.py

@@ -371,7 +371,10 @@ def process_single_pset(
     # Step 3: Calculate efficiency-corrected quantities
     pset_processed = calculate_efficiency_corrected_quantities(pset_processed)
 
-    # Step 4: CG correction and compute ram mask
+    # Step 4: Add s/c velocity, optionally apply CG correction, and calculate
+    # ram-mask
+    pset_processed = add_spacecraft_velocity_to_pset(pset_processed)
+
     if cg_correct:
         # NOTE: Heliospheric frame energy selection for CG correction
         # The heliospheric (HF) energies passed to the CG correction algorithm
@@ -390,9 +393,9 @@ def process_single_pset(
         pset_processed = apply_compton_getting_correction(
             pset_processed, energy_values_ev
         )
-    else:
-        pset_processed = add_spacecraft_velocity_to_pset(pset_processed)
-        pset_processed = calculate_ram_mask(pset_processed)
+
+    # Always calculate ram-mask to identify ram/anti-ram bins
+    pset_processed = calculate_ram_mask(pset_processed)
 
     return pset_processed
 

imap_processing/tests/ena_maps/test_corrections.py

@@ -7,7 +7,8 @@
 import xarray as xr
 
 from imap_processing.cdf.utils import load_cdf
-from imap_processing.ena_maps import ena_maps
+from imap_processing.ena_maps.ena_maps import HiPointingSet
+from imap_processing.ena_maps.utils.coordinates import CoordNames
 from imap_processing.ena_maps.utils.corrections import (
     PowerLawFluxCorrector,
     _add_cartesian_look_direction,
@@ -439,7 +440,6 @@ def test_calculate_compton_getting_transform(self, mock_imap_state, mock_hi_pset
         assert "energy_sc" in mock_hi_pset
         assert "hae_longitude" in mock_hi_pset
         assert "hae_latitude" in mock_hi_pset
-        assert "ram_mask" in mock_hi_pset
 
         # Verify energy_hf matches input
         np.testing.assert_array_equal(
@@ -463,12 +463,6 @@ def test_calculate_compton_getting_transform(self, mock_imap_state, mock_hi_pset
         assert np.all(mock_hi_pset["hae_latitude"].values >= -90)
         assert np.all(mock_hi_pset["hae_latitude"].values <= 90)
 
-        # Verify ram_mask properties
-        ram_mask = mock_hi_pset["ram_mask"]
-        assert isinstance(ram_mask, xr.DataArray)
-        assert ram_mask.dtype == bool
-        assert ram_mask.shape == mock_hi_pset["energy_sc"].shape
-
     @mock.patch("imap_processing.ena_maps.utils.corrections.geometry.imap_state")
     def test_apply_compton_getting_correction(self, mock_imap_state, mock_hi_pset):
         """Test full Compton-Getting correction pipeline."""
@@ -483,6 +477,9 @@ def test_apply_compton_getting_correction(self, mock_imap_state, mock_hi_pset):
             coords={"esa_energy_step": [1, 2, 3]},
         )
 
+        # add the required sc_velocity to the pointing set
+        mock_hi_pset = add_spacecraft_velocity_to_pset(mock_hi_pset)
+
         # Apply the full correction
         mock_hi_pset = apply_compton_getting_correction(mock_hi_pset, energy_hf)
 
@@ -494,22 +491,21 @@ def test_apply_compton_getting_correction(self, mock_imap_state, mock_hi_pset):
         assert "energy_sc" in mock_hi_pset
         assert "hae_longitude" in mock_hi_pset
         assert "hae_latitude" in mock_hi_pset
-        assert "ram_mask" in mock_hi_pset
 
     @pytest.mark.external_test_data
-    @mock.patch("imap_processing.ena_maps.utils.corrections.geometry.imap_state")
-    def test_compton_getting_with_real_pset(self, mock_imap_state, hi_pset_cdf_path):
+    def test_compton_getting_with_real_pset(self, hi_pset_cdf_path):
         """Test Compton-Getting correction with real Hi PSET data."""
         # Load real pointing set
-        pset_ds = load_cdf(hi_pset_cdf_path)
-        hi_pset = ena_maps.HiPointingSet(pset_ds)
+        pset = load_cdf(hi_pset_cdf_path)
+        pset = pset.rename(HiPointingSet.l1c_to_l2_var_mapping)
 
         # Store original coordinates for comparison
-        original_lon = hi_pset.data["hae_longitude"].copy()
+        original_lon = pset["hae_longitude"].copy()
 
         # Mock spacecraft state
-        mock_sc_state = np.array([1e8, 2e8, 3e8, 12.0, -27.0, 0.02])  # km and km/s
-        mock_imap_state.return_value = mock_sc_state
+        pset["sc_velocity"] = xr.DataArray(
+            np.array([12.0, -27.0, 0.02]), dims=[CoordNames.CARTESIAN_VECTOR.value]
+        )
 
         # Create energy array (Hi has 9 energy steps)
         energy_hf = xr.DataArray(
@@ -521,11 +517,11 @@ def test_compton_getting_with_real_pset(self, mock_imap_state, hi_pset_cdf_path)
         )
 
         # Apply correction (pass the dataset, not the pointing set object)
-        hi_pset.data = apply_compton_getting_correction(hi_pset.data, energy_hf)
+        pset = apply_compton_getting_correction(pset, energy_hf)
 
         # Verify coordinates were modified
-        corrected_lon = hi_pset.data["hae_longitude"]
-        corrected_lat = hi_pset.data["hae_latitude"]
+        corrected_lon = pset["hae_longitude"]
+        corrected_lat = pset["hae_latitude"]
 
         # Shape should now include energy dimension
         assert "esa_energy_step" in corrected_lon.dims
@@ -542,17 +538,6 @@ def test_compton_getting_with_real_pset(self, mock_imap_state, hi_pset_cdf_path)
         assert np.all(corrected_lat.values >= -90)
         assert np.all(corrected_lat.values <= 90)
 
-        # Verify az_el_points was updated
-        assert hi_pset.az_el_points is not None
-        assert isinstance(hi_pset.az_el_points, xr.DataArray)
-
-        # Verify ram_mask was added and has correct properties
-        assert "ram_mask" in hi_pset.data
-        ram_mask = hi_pset.data["ram_mask"]
-        assert isinstance(ram_mask, xr.DataArray)
-        assert ram_mask.dtype == bool
-        assert ram_mask.shape == hi_pset.data["energy_sc"].shape
-
     def test_compton_getting_physical_consistency(self, mock_hi_pset):
         """Test physical consistency of Compton-Getting correction."""
         # Set up a known spacecraft velocity
@@ -586,6 +571,10 @@ def test_compton_getting_physical_consistency(self, mock_hi_pset):
         # 4. Energy variation should exist across different look directions
         assert energy_sc.values.std() > 0
 
+
+class TestRamMask:
+    """Test suite for calculate_ram_mask function."""
+
     def test_ram_mask_calculation(self):
         """Test ram_mask correctly identifies ram and anti-ram directions."""
         # Create a simple mock pset with specific look directions
@@ -613,9 +602,6 @@ def test_ram_mask_calculation(self):
         # Set up spacecraft velocity in +X direction
         sc_velocity = np.array([30.0, 0.0, 0.0])  # km/s
         dataset["sc_velocity"] = xr.DataArray(sc_velocity, dims=["x_y_z"])
-        dataset["sc_direction_vector"] = xr.DataArray(
-            sc_velocity / np.linalg.norm(sc_velocity), dims=["x_y_z"]
-        )
 
         # Add look directions
         dataset = _add_cartesian_look_direction(dataset)
@@ -626,6 +612,8 @@ def test_ram_mask_calculation(self):
         # Calculate CG transform
         dataset = _calculate_compton_getting_transform(dataset, energy_hf)
 
+        dataset = calculate_ram_mask(dataset)
+
         # Verify ram_mask exists
         assert "ram_mask" in dataset
         ram_mask = dataset["ram_mask"]
@@ -690,7 +678,7 @@ def test_update_ram_mask_plus_x_direction(self):
         pset = self.create_synthetic_pset_with_hae_coords()
 
         # Spacecraft velocity in +X direction (HAE frame)
-        pset["sc_direction_vector"] = np.array([1.0, 0.0, 0.0])
+        pset["sc_velocity"] = np.array([1.0, 0.0, 0.0])
         pset = calculate_ram_mask(pset)
 
         lon = pset["hae_longitude"].values
@@ -714,7 +702,7 @@ def test_update_ram_mask_minus_x_direction(self):
         pset = self.create_synthetic_pset_with_hae_coords()
 
         # Spacecraft velocity in -X direction (HAE frame)
-        pset["sc_direction_vector"] = np.array([-1.0, 0.0, 0.0])
+        pset["sc_velocity"] = np.array([-1.0, 0.0, 0.0])
         pset = calculate_ram_mask(pset)
 
         lon = pset["hae_longitude"].values
@@ -737,7 +725,7 @@ def test_update_ram_mask_plus_y_direction(self):
         pset = self.create_synthetic_pset_with_hae_coords()
 
         # Spacecraft velocity in +Y direction (HAE frame)
-        pset["sc_direction_vector"] = np.array([0.0, 1.0, 0.0])
+        pset["sc_velocity"] = np.array([0.0, 1.0, 0.0])
         pset = calculate_ram_mask(pset)
 
         lon = pset["hae_longitude"].values
@@ -758,11 +746,11 @@ def test_update_ram_mask_magnitude_invariance(self):
         pset = self.create_synthetic_pset_with_hae_coords()
 
         # Test with two different magnitudes in the same direction
-        pset["sc_direction_vector"] = np.array([1.0, 0.0, 0.0])
+        pset["sc_velocity"] = np.array([1.0, 0.0, 0.0])
         pset = calculate_ram_mask(pset)
         ram_mask_1 = pset["ram_mask"].values.copy()
 
-        pset["sc_direction_vector"] = np.array([100.0, 0.0, 0.0])
+        pset["sc_velocity"] = np.array([100.0, 0.0, 0.0])
         pset = calculate_ram_mask(pset)
         ram_mask_2 = pset["ram_mask"].values.copy()
 
@@ -774,7 +762,7 @@ def test_update_ram_mask_dot_product_correctness(self):
         pset = self.create_synthetic_pset_with_hae_coords()
 
         # Use a simple spacecraft velocity vector
-        pset["sc_direction_vector"] = np.array([1.0, 0.0, 0.0])
+        pset["sc_velocity"] = np.array([1.0, 0.0, 0.0])
         pset = calculate_ram_mask(pset)
 
         # Manually verify a few specific pixels
@@ -802,7 +790,7 @@ def test_update_ram_mask_dimensions_preserved(self):
         original_dims_2d = pset_2d["hae_longitude"].dims
 
         # Update ram_mask
-        pset_2d["sc_direction_vector"] = np.array([1.0, 1.0, 0.0])
+        pset_2d["sc_velocity"] = np.array([1.0, 1.0, 0.0])
         pset_2d = calculate_ram_mask(pset_2d)
 
         # Verify dimensions are preserved
@@ -833,7 +821,7 @@ def test_update_ram_mask_dimensions_preserved(self):
         original_dims_1d = dataset_1d["hae_longitude"].dims
 
         # Update ram_mask
-        dataset_1d["sc_direction_vector"] = np.array([1.0, 1.0, 0.0])
+        dataset_1d["sc_velocity"] = np.array([1.0, 1.0, 0.0])
         dataset_1d = calculate_ram_mask(dataset_1d)
 
         # Verify dimensions are preserved
@@ -844,12 +832,12 @@ def test_update_ram_mask_replaces_existing(self):
         pset = self.create_synthetic_pset_with_hae_coords()
 
         # Set initial mask with +X direction
-        pset["sc_direction_vector"] = np.array([1.0, 0.0, 0.0])
+        pset["sc_velocity"] = np.array([1.0, 0.0, 0.0])
         pset = calculate_ram_mask(pset)
         ram_mask_1 = pset["ram_mask"].values.copy()
 
         # Update mask with opposite direction
-        pset["sc_direction_vector"] = np.array([-1.0, 0.0, 0.0])
+        pset["sc_velocity"] = np.array([-1.0, 0.0, 0.0])
         pset = calculate_ram_mask(pset)
         ram_mask_2 = pset["ram_mask"].values.copy()
 
@@ -861,7 +849,7 @@ def test_update_ram_mask_arbitrary_direction(self):
         pset = self.create_synthetic_pset_with_hae_coords(shape=(36, 18))
 
         # Use an arbitrary direction (not aligned with axes)
-        pset["sc_direction_vector"] = np.array([1.0, 1.0, 0.5])
+        pset["sc_velocity"] = np.array([1.0, 1.0, 0.5])
         pset = calculate_ram_mask(pset)
 
         # Verify the mask was created