ENA - Add spacecraft position to Lo L1C and L2 datasets

imap_processing/cdf/config/imap_lo_l1c_variable_attrs.yaml

@@ -9,6 +9,21 @@ default_attrs: &default
   VALIDMIN: 0
   VAR_TYPE: data
 
+default_float32_attrs: &default_float32
+  <<: *default
+  FILLVAL: -1.0e31
+  FORMAT: F12.6
+  VALIDMAX: 3.4028235e+38
+  VALIDMIN: -3.4028235e+38
+  dtype: float32
+
+component:
+  CATDESC: Cartesian components (x,y,z)
+  FIELDNAM: component
+  FORMAT: A3
+  LABLAXIS: component
+  VAR_TYPE: metadata
+
 # Non-epoch Coordinates
 esa_energy_step_label:
   CATDESC: ESA Steps
@@ -268,3 +283,27 @@ o_background_rates_sys_err:
   LABL_PTR_1: esa_energy_step_label
   LABL_PTR_2: spin_angle_label
   LABL_PTR_3: off_angle_label
+
+sc_velocity:
+  <<: *default_float32
+  CATDESC: x,y,z-components in the HAE coordinate system of the mean velocity vector, in km/s
+  FIELDNAM: sc_velocity
+  FORMAT: F12.2
+  LABLAXIS: spacecraft velocity
+  LABL_PTR_1: label_vector_HAE
+  UNITS: "km/s"
+
+sc_position:
+  <<: *default_float32
+  CATDESC: x,y,z-components in the HAE coordinate system of the mean spacecraft position, in km
+  FIELDNAM: sc_position
+  FORMAT: F12.2
+  LABLAXIS: spacecraft position
+  LABL_PTR_1: label_vector_HAE
+  UNITS: "km"
+
+label_vector_HAE:
+  CATDESC: Cartesian components (x,y,z)
+  FIELDNAM: Cartesian components
+  FORMAT: A5
+  VAR_TYPE: metadata

imap_processing/ena_maps/utils/corrections.py

@@ -427,28 +427,30 @@ def apply_flux_correction(
         return corrected_flux_da, corrected_unc_da
 
 
-def add_spacecraft_velocity_to_pset(
+def add_spacecraft_position_and_velocity_to_pset(
     pset: xr.Dataset,
 ) -> xr.Dataset:
     """
-    Calculate and add spacecraft velocity data to pointing set dataset.
+    Calculate and add spacecraft position and velocity data to pointing set dataset.
 
     Parameters
     ----------
     pset : xr.Dataset
         Pointing set dataset to be updated. Must contain "epoch" coordinate
-        and "epoch_delta" data variable.
+        and "epoch_delta" data variable or "pointing_start_met" and
+        "pointing_end_met" data variables to compute.
 
     Returns
     -------
     pset_processed : xarray.Dataset
-        Pointing set dataset with spacecraft velocity data added.
+        Pointing set dataset with spacecraft position and velocity data added.
+        These values are calculated at the midpoint time of the pointing.
 
     Notes
     -----
     Adds the following DataArrays to input dataset:
     - "sc_velocity": Spacecraft velocity vector (km/s) with dims ["x_y_z"]
-    - "sc_direction_vector": Spacecraft velocity unit vector with dims ["x_y_z"]
+    - "sc_position": Spacecraft position vector (km) with dims ["x_y_z"]
     """
     # Hi and Lo need to use different methods for computing the Pointing
     # midpoint time.
@@ -465,7 +467,7 @@ def add_spacecraft_velocity_to_pset(
         ) * 1e9
     else:
         raise NotImplementedError(
-            f"add_spacecraft_velocity_to_pset does not support PSETs with "
+            f"add_spacecraft_position_and_velocity_to_pset does not support PSETs with "
             f"Logical_source: {pset.attrs['Logical_source']}"
         )
 
@@ -475,26 +477,29 @@ def add_spacecraft_velocity_to_pset(
     if pointing_duration_ns <= 0:
         logger.warning(
             "Pointing duration is zero or negative. "
-            "Setting spacecraft velocity to zero."
+            "Setting spacecraft position and velocity to zero."
         )
+        sc_position_vector = np.zeros(3)  # Zero position vector
         sc_velocity_vector = np.zeros(3)  # Zero velocity vector
     else:
         # Compute ephemeris time (J2000 seconds) of PSET midpoint
         et = ttj2000ns_to_et(pset["epoch"].values[0] + pointing_duration_ns / 2)
 
         # Get spacecraft state in HAE frame
         sc_state = geometry.imap_state(et, ref_frame=geometry.SpiceFrame.IMAP_HAE)
+        sc_position_vector = sc_state[0:3]
         sc_velocity_vector = sc_state[3:6]
 
+    # Store spacecraft position as DataArray
+    pset["sc_position"] = xr.DataArray(
+        sc_position_vector, dims=[CoordNames.CARTESIAN_VECTOR.value]
+    )
+
     # Store spacecraft velocity as DataArray
     pset["sc_velocity"] = xr.DataArray(
         sc_velocity_vector, dims=[CoordNames.CARTESIAN_VECTOR.value]
     )
 
-    # Calculate spacecraft speed and direction
-    sc_velocity_km_per_sec = np.linalg.norm(pset["sc_velocity"], axis=-1, keepdims=True)
-    pset["sc_direction_vector"] = pset["sc_velocity"] / sc_velocity_km_per_sec
-
     return pset
 
 
@@ -747,7 +752,7 @@ def apply_compton_getting_correction(
         Must contain the following variables:
           - 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.
+            `add_spacecraft_position_and_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

imap_processing/hi/hi_l2.py

@@ -14,7 +14,7 @@
 )
 from imap_processing.ena_maps.utils.corrections import (
     PowerLawFluxCorrector,
-    add_spacecraft_velocity_to_pset,
+    add_spacecraft_position_and_velocity_to_pset,
     apply_compton_getting_correction,
     calculate_ram_mask,
     get_pset_directional_mask,
@@ -276,8 +276,8 @@ def process_single_pset(
         pset_processed["exposure_factor"], float(mid_time)
     )
 
-    # Step 3: Add spacecraft velocity
-    pset_processed = add_spacecraft_velocity_to_pset(pset_processed)
+    # Step 3: Add spacecraft position and velocity
+    pset_processed = add_spacecraft_position_and_velocity_to_pset(pset_processed)
 
     # Step 4: Optionally apply Compton-Getting correction for heliocentric frame
     if descriptor.frame_descriptor == "hf":

imap_processing/lo/l1c/lo_l1c.py

@@ -9,7 +9,9 @@
 import xarray as xr
 
 from imap_processing.cdf.imap_cdf_manager import ImapCdfAttributes
-from imap_processing.ena_maps.utils.corrections import add_spacecraft_velocity_to_pset
+from imap_processing.ena_maps.utils.corrections import (
+    add_spacecraft_position_and_velocity_to_pset,
+)
 from imap_processing.lo import lo_ancillary
 from imap_processing.lo.l1b.lo_l1b import set_bad_or_goodtimes
 from imap_processing.spice.geometry import (
@@ -222,8 +224,18 @@ def lo_l1c(sci_dependencies: dict, anc_dependencies: list) -> list[xr.Dataset]:
         }
     )
 
-    # add the spacecraft velocity and direction
-    pset = add_spacecraft_velocity_to_pset(pset)
+    # Get the spacecraft position and velocity and direction
+    pset = add_spacecraft_position_and_velocity_to_pset(pset)
+
+    # Update the attributes for the spacecraft position and velocity variables
+    pset["sc_position"].attrs.update(attr_mgr.get_variable_attributes("sc_position"))
+    pset["sc_velocity"].attrs.update(attr_mgr.get_variable_attributes("sc_velocity"))
+    pset["label_vector_HAE"] = xr.DataArray(
+        np.array(["x HAE", "y HAE", "z HAE"], dtype=str),
+        name="label_vector_HAE",
+        dims=[" "],
+        attrs=attr_mgr.get_variable_attributes("label_vector_HAE", check_schema=False),
+    )
 
     return [pset]
 

imap_processing/lo/l2/lo_l2.py

@@ -2,6 +2,7 @@
 
 import logging
 from pathlib import Path
+from typing import cast
 
 import numpy as np
 import pandas as pd
@@ -12,7 +13,6 @@
 from imap_processing.ena_maps.ena_maps import AbstractSkyMap, RectangularSkyMap
 from imap_processing.ena_maps.utils.corrections import (
     PowerLawFluxCorrector,
-    add_spacecraft_velocity_to_pset,
     apply_compton_getting_correction,
     calculate_ram_mask,
     get_pset_directional_mask,
@@ -108,7 +108,7 @@ def lo_l2(
     )
 
     logger.info("Step 5: Finalizing dataset with attributes")
-    dataset = sky_map.build_cdf_dataset(  # type: ignore[attr-defined]
+    dataset = cast(RectangularSkyMap, sky_map).build_cdf_dataset(
         instrument="lo", level="l2", descriptor=descriptor, external_map_dataset=dataset
     )
 
@@ -371,10 +371,7 @@ def process_single_pset(
     # Step 3: Calculate efficiency-corrected quantities
     pset_processed = calculate_efficiency_corrected_quantities(pset_processed)
 
-    # Step 4: Add s/c velocity, optionally apply CG correction, and calculate
-    # ram-mask
-    pset_processed = add_spacecraft_velocity_to_pset(pset_processed)
-
+    # Step 4: Optionally apply CG correction and calculate ram-mask
     if cg_correct:
         # NOTE: Heliospheric frame energy selection for CG correction
         # The heliospheric (HF) energies passed to the CG correction algorithm

imap_processing/tests/ena_maps/test_corrections.py

@@ -13,7 +13,7 @@
     PowerLawFluxCorrector,
     _add_cartesian_look_direction,
     _calculate_compton_getting_transform,
-    add_spacecraft_velocity_to_pset,
+    add_spacecraft_position_and_velocity_to_pset,
     apply_compton_getting_correction,
     calculate_ram_mask,
     get_pset_directional_mask,
@@ -516,18 +516,18 @@ class TestComptonGettingCorrection:
 
     @mock.patch("imap_processing.ena_maps.utils.corrections.ttj2000ns_to_et")
     @mock.patch("imap_processing.ena_maps.utils.corrections.geometry.imap_state")
-    def test_add_spacecraft_velocity_to_pset(
+    def test_add_spacecraft_position_and_velocity_to_pset(
         self, mock_imap_state, mock_ttj2000_to_et, mock_hi_pset
     ):
-        """Test that spacecraft velocity is correctly added to pointing set."""
+        """Test that spacecraft position and velocity are correctly added to pset."""
         # Mock conversion from TTJ2000ns to ET
         et = 1000.0
         mock_ttj2000_to_et.return_value = et
         # Mock spacecraft state vector (position + velocity in HAE frame)
         mock_sc_state = np.array([1e8, 2e8, 3e8, 10.0, 20.0, 30.0])  # km and km/s
         mock_imap_state.return_value = mock_sc_state
 
-        mock_hi_pset = add_spacecraft_velocity_to_pset(mock_hi_pset)
+        mock_hi_pset = add_spacecraft_position_and_velocity_to_pset(mock_hi_pset)
 
         # Verify SPICE was called correctly
         mock_imap_state.assert_called_once_with(
@@ -541,20 +541,19 @@ def test_add_spacecraft_velocity_to_pset(
             mock_hi_pset["sc_velocity"].values, np.array([10.0, 20.0, 30.0])
         )
 
-        # Verify sc_direction_vector was added
-        assert "sc_direction_vector" in mock_hi_pset
-        expected_speed = np.sqrt(10**2 + 20**2 + 30**2)
-        expected_direction = np.array([10.0, 20.0, 30.0]) / expected_speed
-        np.testing.assert_allclose(
-            mock_hi_pset["sc_direction_vector"].values, expected_direction
+        # Verify sc_position was added
+        assert "sc_position" in mock_hi_pset
+        assert isinstance(mock_hi_pset["sc_position"], xr.DataArray)
+        np.testing.assert_array_equal(
+            mock_hi_pset["sc_position"].values, np.array([1e8, 2e8, 3e8])
         )
 
     @mock.patch("imap_processing.ena_maps.utils.corrections.ttj2000ns_to_et")
     @mock.patch("imap_processing.ena_maps.utils.corrections.geometry.imap_state")
-    def test_add_spacecraft_velocity_to_pset_lo(
+    def test_add_spacecraft_position_and_velocity_to_pset_lo(
         self, mock_imap_state, mock_ttj2000_to_et, mock_lo_pset
     ):
-        """Test that spacecraft velocity is correctly added to Lo pointing set."""
+        """Test that S/C position and velocity are correctly added to Lo pset."""
         # Mock conversion from TTJ2000ns to ET
         et = 1000.0
         mock_ttj2000_to_et.return_value = et
@@ -568,7 +567,7 @@ def test_add_spacecraft_velocity_to_pset_lo(
         # Midpoint: epoch + pointing_duration_ns / 2
         expected_midpoint_time_ns = mock_lo_pset["epoch"].values[0] + 1e11 / 2
 
-        mock_lo_pset = add_spacecraft_velocity_to_pset(mock_lo_pset)
+        mock_lo_pset = add_spacecraft_position_and_velocity_to_pset(mock_lo_pset)
 
         # Verify SPICE was called correctly
         mock_ttj2000_to_et.assert_called_once_with(expected_midpoint_time_ns)
@@ -583,15 +582,14 @@ def test_add_spacecraft_velocity_to_pset_lo(
             mock_lo_pset["sc_velocity"].values, np.array([15.0, 25.0, 35.0])
         )
 
-        # Verify sc_direction_vector was added
-        assert "sc_direction_vector" in mock_lo_pset
-        expected_speed = np.sqrt(15**2 + 25**2 + 35**2)
-        expected_direction = np.array([15.0, 25.0, 35.0]) / expected_speed
-        np.testing.assert_allclose(
-            mock_lo_pset["sc_direction_vector"].values, expected_direction
+        # Verify sc_position was added
+        assert "sc_position" in mock_lo_pset
+        assert isinstance(mock_lo_pset["sc_position"], xr.DataArray)
+        np.testing.assert_array_equal(
+            mock_lo_pset["sc_position"].values, np.array([1e8, 2e8, 3e8])
         )
 
-    def test_add_spacecraft_velocity_unsupported_instrument(self):
+    def test_add_spacecraft_position_and_velocity_unsupported_instrument(self):
         """Test that unsupported instrument raises NotImplementedError."""
         # Create a dataset with unsupported Logical_source
         unsupported_pset = xr.Dataset(
@@ -603,7 +601,18 @@ def test_add_spacecraft_velocity_unsupported_instrument(self):
         )
 
         with pytest.raises(NotImplementedError, match="does not support PSETs"):
-            add_spacecraft_velocity_to_pset(unsupported_pset)
+            add_spacecraft_position_and_velocity_to_pset(unsupported_pset)
+
+    def test_add_spacecraft_position_and_velocity_zero_duration(self, mock_hi_pset):
+        """Test that zero pointing duration sets pos and velocity to zero vectors."""
+        # Set epoch_delta to zero to simulate an empty/filtered pointing set
+        mock_hi_pset["epoch_delta"] = xr.DataArray(np.array([0.0]), dims=["epoch"])
+
+        result = add_spacecraft_position_and_velocity_to_pset(mock_hi_pset)
+
+        # Both sc_velocity and sc_position should be zero vectors
+        np.testing.assert_array_equal(result["sc_velocity"].values, np.zeros(3))
+        np.testing.assert_array_equal(result["sc_position"].values, np.zeros(3))
 
     def test_add_cartesian_look_direction(self, mock_hi_pset):
         """Test that look directions are correctly calculated and added."""
@@ -627,7 +636,7 @@ def test_calculate_compton_getting_transform(self, mock_imap_state, mock_hi_pset
         mock_sc_state = np.array([1e8, 2e8, 3e8, 10.0, 20.0, 30.0])
         mock_imap_state.return_value = mock_sc_state
 
-        mock_hi_pset = add_spacecraft_velocity_to_pset(mock_hi_pset)
+        mock_hi_pset = add_spacecraft_position_and_velocity_to_pset(mock_hi_pset)
         mock_hi_pset = _add_cartesian_look_direction(mock_hi_pset)
 
         # Create energy array
@@ -682,14 +691,13 @@ def test_apply_compton_getting_correction(self, mock_imap_state, mock_hi_pset):
         )
 
         # add the required sc_velocity to the pointing set
-        mock_hi_pset = add_spacecraft_velocity_to_pset(mock_hi_pset)
+        mock_hi_pset = add_spacecraft_position_and_velocity_to_pset(mock_hi_pset)
 
         # Apply the full correction
         mock_hi_pset = apply_compton_getting_correction(mock_hi_pset, energy_hf)
 
         # Verify all intermediate variables were added
         assert "sc_velocity" in mock_hi_pset
-        assert "sc_direction_vector" in mock_hi_pset
         assert "look_direction" in mock_hi_pset
         assert "energy_hf" in mock_hi_pset
         assert "energy_sc" in mock_hi_pset