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

imap_processing/ena_maps/ena_maps.py

@@ -73,7 +73,7 @@ def match_coords_to_indices(
     input_object: PointingSet | AbstractSkyMap,
     output_object: PointingSet | AbstractSkyMap,
     event_et: float | None = None,
-) -> NDArray:
+) -> NDArray | xr.DataArray:
     """
     Find the output indices corresponding to each input coord between 2 spatial objects.
 
@@ -87,7 +87,7 @@ def match_coords_to_indices(
     This function always "pushes" the pixels of the input object to corresponding pixels
     in the output object's unwrapped rectangular grid or healpix tessellation;
     however, by swapping the input and output objects, one can apply the "pull" method
-    of index  matching.
+    of index matching.
 
     At present, the allowable inputs are either:
     - A PointingSet object and a SkyMap object, in either order of input/output.
@@ -97,7 +97,7 @@ def match_coords_to_indices(
     Parameters
     ----------
     input_object : PointingSet | AbstractSkyMap
-        An object containing 1D spatial pixel centers in azimuth and elevation,
+        An object containing spatial pixel centers in azimuth and elevation,
         which will be matched to 1D indices of spatial pixels in the output frame.
         Must contain the Spice frame in which the pixel centers are defined.
     output_object : PointingSet | AbstractSkyMap
@@ -114,11 +114,13 @@ def match_coords_to_indices(
 
     Returns
     -------
-    flat_indices_input_grid_output_frame : NDArray
-        1D array of pixel indices of the output object corresponding to each pixel in
-        the input object. The length of the array is equal to the number of pixels in
-        the input object, and may contain 0, 1, or multiple occurrences of the same
-        output index.
+    flat_indices_input_grid_output_frame : xr.DataArray
+        Array of pixel indices mapping each input object pixel center to a pixel
+        in the output object. The output xr.DataArray will have the same leading
+        dimension labels preserved. The shape of the output array is (..., n)
+        where ... matches the non-spatial dimensions of the input object, and n
+        is the number of spatial pixels in the input object. Output indices may
+        contain 0, 1, or multiple occurrences of the same output index.
 
     Raises
     ------
@@ -166,14 +168,14 @@ def match_coords_to_indices(
         # use ravel_multi_index to get the 1D indices of the pixels in the output frame.
         az_indices = (
             np.digitize(
-                input_obj_az_el_output_frame[:, 0],
+                input_obj_az_el_output_frame[..., 0],
                 output_object.sky_grid.az_bin_edges,
             )
             - 1
         )
         el_indices = (
             np.digitize(
-                input_obj_az_el_output_frame[:, 1],
+                input_obj_az_el_output_frame[..., 1],
                 output_object.sky_grid.el_bin_edges,
             )
             - 1
@@ -191,8 +193,8 @@ def match_coords_to_indices(
         # which directly returns the index on the output frame's Healpix tessellation.
         flat_indices_input_grid_output_frame = hp.ang2pix(
             nside=output_object.nside,
-            theta=input_obj_az_el_output_frame[:, 0],  # Lon in degrees
-            phi=input_obj_az_el_output_frame[:, 1],  # Lat in degrees
+            theta=input_obj_az_el_output_frame[..., 0],  # Lon in degrees
+            phi=input_obj_az_el_output_frame[..., 1],  # Lat in degrees
             nest=output_object.nested,
             lonlat=True,
         )
@@ -202,6 +204,14 @@ def match_coords_to_indices(
             f"Received: {output_object.tiling_type}"
         )
 
+    # Wrap the output indices in a DataArray with the same leading dimensions as
+    # the input object az_el_points to preserve broadcasting information
+    input_dims = input_obj_az_el_input_frame.dims[:-1]
+    flat_indices_input_grid_output_frame = xr.DataArray(
+        flat_indices_input_grid_output_frame,
+        dims=input_dims,
+    )
+
     return flat_indices_input_grid_output_frame
 
 
@@ -236,9 +246,10 @@ class PointingSet(ABC):
     spice_reference_frame: geometry.SpiceFrame
 
     # ======== Attributes required to be set in a subclass ========
-    # Azimuth and elevation coordinates of each spatial pixel. The ndarray should
-    # have the shape (n, 2) where n is the number of spatial pixels
-    az_el_points: np.ndarray
+    # Azimuth and elevation coordinates of each spatial pixel. Must be an
+    # xr.DataArray with dimensions (..., spatial_dim, az_el_coord) to preserve
+    # dimension labels
+    az_el_points: xr.DataArray
     # Tuple containing the names of each spatial coordinate of the xarray.Dataset
     # stored in the data attribute
     spatial_coords: tuple[str, ...]
@@ -274,7 +285,9 @@ class PointingSet(ABC):
         num_points: int
             The number of spatial pixels in the pointing set.
         """
-        return self.az_el_points.shape[0]
+        # Last dimension is az/el vector, the second to last dimension is
+        # the number of pixels.
+        return self.az_el_points.shape[-2]
 
     @property
     def epoch(self) -> int:
@@ -430,11 +443,12 @@ class RectangularPointingSet(PointingSet):
         # into shape (number of points in tiling of the sky, 2) where
         # column 0 (az_el_points[:, 0]) is the azimuth of that point and
         # column 1 (az_el_points[:, 1]) is the elevation of that point.
-        self.az_el_points = np.column_stack(
-            (
-                self.sky_grid.az_grid.ravel(),
-                self.sky_grid.el_grid.ravel(),
-            )
+        self.az_el_points = xr.DataArray(
+            np.stack(
+                (self.sky_grid.az_grid.ravel(), self.sky_grid.el_grid.ravel()),
+                axis=-1,
+            ),
+            dims=[CoordNames.GENERIC_PIXEL.value, CoordNames.AZ_EL_VECTOR.value],
         )
 
 
@@ -540,8 +554,9 @@ class UltraPointingSet(HealpixPointingSet):
         # The coordinates of the healpix pixel centers are stored as a 2D array
         # of shape (num_points, 2) where column 0 is the lon/az
         # and column 1 is the lat/el.
-        self.az_el_points = np.column_stack(
-            (azimuth_pixel_center, elevation_pixel_center)
+        self.az_el_points = xr.DataArray(
+            np.stack((azimuth_pixel_center, elevation_pixel_center), axis=-1),
+            dims=[CoordNames.GENERIC_PIXEL.value, CoordNames.AZ_EL_VECTOR.value],
         )
 
     @property
@@ -586,7 +601,45 @@ class UltraPointingSet(HealpixPointingSet):
         )
 
 
-class HiPointingSet(PointingSet):
+class LoHiBasePointingSet(PointingSet):
+    """
+    Base class for Lo and Hi pointing sets with HAE coordinate data.
+
+    This class provides common functionality for pointing sets that contain
+    hae_longitude and hae_latitude coordinates in the dataset.
+    """
+
+    tiling_type: SkyTilingType = SkyTilingType.RECTANGULAR
+
+    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
+        are used to construct the azimuth and elevation coordinates.
+        """
+        # 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.
+        az_stacked = (
+            self.data["hae_longitude"]
+            .squeeze("epoch")
+            .stack({CoordNames.GENERIC_PIXEL.value: self.spatial_coords})
+        )
+        el_stacked = (
+            self.data["hae_latitude"]
+            .squeeze("epoch")
+            .stack({CoordNames.GENERIC_PIXEL.value: self.spatial_coords})
+        )
+
+        # Stack lon/lat along last axis to create shape (..., 2)
+        self.az_el_points = xr.DataArray(
+            np.stack([az_stacked.values, el_stacked.values], axis=-1),
+            dims=[*az_stacked.dims, CoordNames.AZ_EL_VECTOR.value],
+        )
+
+
+class HiPointingSet(LoHiBasePointingSet):
     """
     PointingSet object specific to Hi L1C PSet data.
 
@@ -602,7 +655,7 @@ class HiPointingSet(PointingSet):
         super().__init__(dataset, spice_reference_frame=geometry.SpiceFrame.ECLIPJ2000)
 
         # Filter out ENAs from non-selected portions of the spin.
-        if spin_phase not in ["full", "ram", "anti-ram"]:
+        if spin_phase not in ["full", "ram", "anti"]:
             raise ValueError(f"Unrecognized spin_phase value: {spin_phase}.")
         # ram only includes spin-phase interval [0, 0.5)
         # which is the first half of the spin_angle_bins
@@ -612,7 +665,7 @@ class HiPointingSet(PointingSet):
             )
         # anti-ram includes spin-phase interval [0.5, 1)
         # which is the second half of the spin_angle_bins
-        elif spin_phase == "anti-ram":
+        elif spin_phase == "anti":
             self.data = self.data.isel(
                 spin_angle_bin=slice(self.data["spin_angle_bin"].data.size // 2, None)
             )
@@ -631,16 +684,13 @@ class HiPointingSet(PointingSet):
             self.data["exposure_factor"], self.data["epoch"].values[0]
         )
 
-        self.az_el_points = np.column_stack(
-            (
-                np.squeeze(self.data["hae_longitude"]),
-                np.squeeze(self.data["hae_latitude"]),
-            )
-        )
         self.spatial_coords = ("spin_angle_bin",)
 
+        # Update az_el_points using the base class method
+        self.update_az_el_points()
+
 
-class LoPointingSet(PointingSet):
+class LoPointingSet(LoHiBasePointingSet):
     """
     PointingSet object specific to Lo L1C PSet data.
 
@@ -653,15 +703,11 @@ class LoPointingSet(PointingSet):
     def __init__(self, dataset: xr.Dataset):
         super().__init__(dataset, spice_reference_frame=geometry.SpiceFrame.IMAP_HAE)
 
-        # The HAE centers are stored in the pset as (1, 3600, 40) arrays
-        self.az_el_points = np.column_stack(
-            (
-                np.squeeze(self.data["hae_longitude"]).values.ravel(),
-                np.squeeze(self.data["hae_latitude"]).values.ravel(),
-            )
-        )
         self.spatial_coords = ("spin_angle", "off_angle")
 
+        # Update az_el_points using the base class method
+        self.update_az_el_points()
+
 
 # Define the Map classes
 class AbstractSkyMap(ABC):
@@ -694,9 +740,10 @@ class AbstractSkyMap(ABC):
     max_epoch: int
 
     # ======== Attributes required to be set in a subclass ========
-    # Azimuth and elevation coordinates of each spatial pixel. The ndarray should
-    # have the shape (n, 2) where n is the number of spatial pixels
-    az_el_points: np.ndarray
+    # Azimuth and elevation coordinates of each spatial pixel. The xarray.DataArray
+    # should have the shape (n, 2) where n is the number of spatial pixels.
+    # Always a simple numpy array for maps (no need for multi-dimensional coords).
+    az_el_points: xr.DataArray
     # Type of sky tiling
     tiling_type: SkyTilingType
     # Dictionary of xr.DataArray objects for each non-spatial coordinate in the SkyMap
@@ -829,19 +876,11 @@ class AbstractSkyMap(ABC):
             )
 
         for value_key in value_keys:
-            pset_values = pointing_set.data[value_key]
-
             # If multiple spatial axes present
             # (i.e (az, el) for rectangular coordinate PSET),
             # flatten them in the values array to match the raveled indices
-            non_spatial_axes_shape = tuple(
-                size
-                for key, size in pset_values.sizes.items()
-                if key not in pointing_set.spatial_coords
-            )
-            raveled_pset_data = pset_values.data.reshape(
-                *non_spatial_axes_shape,
-                pointing_set.num_points,
+            raveled_pset_data = pointing_set.data[value_key].stack(
+                {CoordNames.GENERIC_PIXEL.value: pointing_set.spatial_coords}
             )
 
             if value_key not in self.data_1d.data_vars:
@@ -864,10 +903,16 @@ class AbstractSkyMap(ABC):
             if index_match_method is IndexMatchMethod.PUSH:
                 # Bin the values at the matched indices. There may be multiple
                 # pointing set pixels that correspond to the same sky map pixel.
+                # Broadcast all arrays together using xarray dimension alignment
+                data_bc, indices_bc = xr.broadcast(
+                    raveled_pset_data, matched_indices_push
+                )
+
+                # Extract numpy arrays for bincount operation
                 pointing_projected_values = map_utils.bin_single_array_at_indices(
-                    value_array=raveled_pset_data,
+                    value_array=data_bc.values,
                     projection_grid_shape=self.binning_grid_shape,
-                    projection_indices=matched_indices_push,
+                    projection_indices=indices_bc.values,
                     input_valid_mask=pset_valid_mask,
                 )
                 # TODO: we may need to allow for unweighted/weighted means here by
@@ -879,7 +924,7 @@ class AbstractSkyMap(ABC):
                 valid_map_mask = pset_valid_mask[matched_indices_pull]
                 # We know that there will only be one value per sky map pixel,
                 # so we can use the matched indices directly
-                pointing_projected_values = raveled_pset_data[
+                pointing_projected_values = raveled_pset_data.values[
                     ..., matched_indices_pull[valid_map_mask]
                 ]
                 # TODO: we may need to allow for unweighted/weighted means here by
@@ -1120,7 +1165,10 @@ class RectangularSkyMap(AbstractSkyMap):
         el_points = self.sky_grid.el_grid.ravel()
 
         # Stack so axis 0 is different pixels, and axis 1 is (az, el) of the pixel
-        self.az_el_points = np.column_stack((az_points, el_points))
+        self.az_el_points = xr.DataArray(
+            np.column_stack((az_points, el_points)),
+            dims=[CoordNames.GENERIC_PIXEL.value, CoordNames.AZ_EL_VECTOR.value],
+        )
 
         # Calculate solid angles of each pixel in the map grid in units of steradians
         self.solid_angle_grid = spatial_utils.build_solid_angle_map(
@@ -1267,7 +1315,7 @@ class RectangularSkyMap(AbstractSkyMap):
             if ("L2" in name)
         ]
         l2_coords.append(CoordNames.TIME.value)
-        for map_coord in cdf_ds.dims.keys():
+        for map_coord in cdf_ds.dims:
             if map_coord not in l2_coords:
                 cdf_ds = cdf_ds.drop_dims(map_coord)
 
@@ -1331,25 +1379,28 @@ class RectangularSkyMap(AbstractSkyMap):
         )
         cdf_ds.attrs.update(map_attrs)
 
-        # Set the variable attributes
-        for var in [*cdf_ds.data_vars, *cdf_ds.coords]:
+        # Set the variable and coordinate attributes
+        for name, data_array in {**cdf_ds.data_vars, **cdf_ds.coords}.items():
             try:
-                # Don't check schema on label or delta variables
-                ignore_schema_substrings = ["_label", "_delta"]
-                check_schema = (
-                    False if any(s in var for s in ignore_schema_substrings) else True
-                )
+                # We only check the schema on data variables that include "epoch"
+                # in their list of dimensions (But not epoch itself).
+                check_schema = name != "epoch" and "epoch" in data_array.dims
                 var_attrs = cdf_attrs.get_variable_attributes(
-                    variable_name=var,
+                    variable_name=name,
                     check_schema=check_schema,
                 )
             except KeyError as e:
                 raise KeyError(
-                    f"Attributes for variable {var} not found in "
+                    f"Attributes for variable {name} not found in "
                     f"loaded variable attributes."
                 ) from e
 
-            cdf_ds[var].attrs.update(var_attrs)
+            cdf_ds[name].attrs.update(var_attrs)
+
+        # Manually adjust epoch attributes
+        cdf_ds["epoch"].attrs.update(
+            {"DELTA_PLUS_VAR": "epoch_delta", "BIN_LOCATION": 0}
+        )
 
         return cdf_ds
 
@@ -1428,7 +1479,10 @@ class HealpixSkyMap(AbstractSkyMap):
             nside=nside, ipix=np.arange(hp.nside2npix(nside)), nest=nested, lonlat=True
         )
         # Stack so axis 0 is different pixels, and axis 1 is (az, el) of the pixel
-        self.az_el_points = np.column_stack((pixel_az, pixel_el))
+        self.az_el_points = xr.DataArray(
+            np.column_stack((pixel_az, pixel_el)),
+            dims=[CoordNames.GENERIC_PIXEL.value, CoordNames.AZ_EL_VECTOR.value],
+        )
 
         self.spatial_coords = {
             CoordNames.HEALPIX_INDEX.value: xr.DataArray(
@@ -1764,7 +1818,7 @@ class HealpixSkyMap(AbstractSkyMap):
                     value_array=healpix_values_array,
                     max_subdivision_depth=max_subdivision_depth,
                 )
-                for lon_lat in rect_map.az_el_points
+                for lon_lat in rect_map.az_el_points.values
             ]
 
             # Separate the best value and the recursion depth for each pixel

imap_processing/ena_maps/utils/coordinates.py

@@ -18,3 +18,8 @@ class CoordNames(Enum):
     ELEVATION_L1C = "latitude"
     AZIMUTH_L2 = "longitude"
     ELEVATION_L2 = "latitude"
+
+    # Common name for dimension along azimuth/elevation vector
+    AZ_EL_VECTOR = "az_el"
+    # Commoon name for dimension along Cartesian vector
+    CARTESIAN_VECTOR = "x_y_z"

imap_processing/ena_maps/utils/corrections.py

@@ -4,7 +4,23 @@ from pathlib import Path
 
 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.utils.coordinates import CoordNames
+from imap_processing.spice import geometry
+from imap_processing.spice.time import ttj2000ns_to_et
+
+# Physical constants for Compton-Getting correction
+# Units: electron_volt = [J / eV]
+#        erg = [J / erg]
+# To get [erg / eV], => electron_volt [J / eV] / erg [J / erg] = erg_per_ev [erg / eV]
+ERG_PER_EV = electron_volt / erg  # erg per eV - unit conversion factor
+# Units: proton_mass = [kg]
+# Here, we convert proton_mass to grams
+PROTON_MASS_GRAMS = proton_mass * 1e3  # proton mass in grams
 
 
 class PowerLawFluxCorrector:
@@ -289,3 +305,255 @@ class PowerLawFluxCorrector:
             )
 
         return corrected_flux, corrected_flux_stat_unc
+
+
+def _add_spacecraft_velocity_to_pset(pset: LoHiBasePointingSet) -> None:
+    """
+    Calculate and add spacecraft velocity data to pointing set.
+
+    Parameters
+    ----------
+    pset : LoHiBasePointingSet
+        Pointing set object to be updated.
+
+    Notes
+    -----
+    Adds the following DataArrays to pset.data:
+    - "sc_velocity": Spacecraft velocity vector (km/s) with dims ["x_y_z"]
+    - "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])
+    # 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]
+
+    # Store spacecraft velocity as DataArray
+    pset.data["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.data["sc_velocity"], axis=-1, keepdims=True
+    )
+    pset.data["sc_direction_vector"] = pset.data["sc_velocity"] / sc_velocity_km_per_sec
+
+
+def _add_cartesian_look_direction(pset: LoHiBasePointingSet) -> None:
+    """
+    Calculate and add look direction vectors to pointing set.
+
+    Parameters
+    ----------
+    pset : LoHiBasePointingSet
+        Pointing set object to be updated.
+
+    Notes
+    -----
+    Adds the following DataArray to pset.data:
+    - "look_direction": Cartesian unit vectors with dims [...spatial_dims, "x_y_z"]
+    """
+    longitudes = pset.data["hae_longitude"]
+    latitudes = pset.data["hae_latitude"]
+
+    # Stack spherical coordinates (r=1 for unit vectors, azimuth, elevation)
+    spherical_coords = np.stack(
+        [
+            np.ones_like(longitudes),  # r = 1 for unit vectors
+            longitudes,  # azimuth = longitude
+            latitudes,  # elevation = latitude
+        ],
+        axis=-1,
+    )
+
+    # Convert to Cartesian coordinates and store as DataArray
+    pset.data["look_direction"] = xr.DataArray(
+        geometry.spherical_to_cartesian(spherical_coords),
+        dims=[*longitudes.dims, CoordNames.CARTESIAN_VECTOR.value],
+    )
+
+
+def _calculate_compton_getting_transform(
+    pset: LoHiBasePointingSet,
+    energy_hf: xr.DataArray,
+) -> None:
+    """
+    Apply Compton-Getting transformation to compute ENA source directions.
+
+    This implements the Compton-Getting velocity transformation to correct
+    for the motion of the spacecraft through the heliosphere. The transformation
+    accounts for the Doppler shift of ENA energies and the aberration of
+    arrival directions.
+
+    All calculations are performed using xarray DataArrays to preserve
+    dimension information throughout the computation.
+
+    Parameters
+    ----------
+    pset : LoHiBasePointingSet
+        Pointing set object with sc_velocity, sc_direction_vector, and
+        look_direction already added.
+    energy_hf : xr.DataArray
+        ENA energies in the heliosphere frame in eV.
+
+    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)
+    """
+    # Store heliosphere frame energies
+    pset.data["energy_hf"] = energy_hf
+
+    # Calculate spacecraft speed
+    sc_velocity_km_per_sec = np.linalg.norm(
+        pset.data["sc_velocity"], axis=-1, keepdims=True
+    )
+
+    # Calculate dot product between look directions and spacecraft direction vector
+    # Use Einstein summation for efficient vectorized dot product
+    dot_product = xr.DataArray(
+        np.einsum(
+            "...i,...i->...",
+            pset.data["look_direction"],
+            pset.data["sc_direction_vector"],
+        ),
+        dims=pset.data["look_direction"].dims[:-1],
+    )
+
+    # Calculate the kinetic energy of a hydrogen ENA traveling at spacecraft velocity
+    # E_u = (1/2) * m * U_sc^2 (convert km/s to cm/s with 1.0e5 factor)
+    energy_u = (
+        0.5 * PROTON_MASS_GRAMS * (sc_velocity_km_per_sec * 1e5) ** 2 / ERG_PER_EV
+    )
+
+    # Note: Tim thinks that this approach seems backwards. Here, we are assuming
+    #     that ENAs are observed in the heliosphere frame at the ESA energy levels.
+    #     We then calculate the velocity that said ENAs would have in the spacecraft
+    #     frame as well as the CG corrected energy level in the spacecraft frame.
+    #     We then use this velocity to calculate and the velocity of the spacecraft
+    #     to do the vector math which determines the ENA source direction in the
+    #     heliosphere frame.
+    #     The ENAs are in fact observed in the spacecraft frame at a known energy
+    #     level in the spacecraft frame. Why don't we use that energy level to
+    #     calculate the source direction in the spacecraft frame and then do the
+    #     vector math to find the source direction in the heliosphere frame? We
+    #     would also need to calculate the CG corrected ENA energy in the heliosphere
+    #     frame and keep track of that when binning.
+
+    # Calculate y values for each energy level (Equation 61)
+    # y_k = sqrt(E^h_k / E^u)
+    y = np.sqrt(pset.data["energy_hf"] / energy_u)
+
+    # Velocity magnitude factor calculation (Equation 62)
+    # x_k = (êₛ · û_sc) + sqrt(y² + (êₛ · û_sc)² - 1)
+    x = dot_product + np.sqrt(y**2 + dot_product**2 - 1)
+
+    # Calculate ENA speed in the spacecraft frame
+    # |v⃗_sc| = x_k * U_sc
+    velocity_sc = x * sc_velocity_km_per_sec
+
+    # Calculate the kinetic energy in the spacecraft frame
+    # E_sc = (1/2) * M_p * v_sc² (convert km/s to cm/s with 1.0e5 factor)
+    pset.data["energy_sc"] = (
+        0.5 * PROTON_MASS_GRAMS * (velocity_sc * 1e5) ** 2 / ERG_PER_EV
+    )
+
+    # Calculate the velocity vector in the spacecraft frame
+    # v⃗_sc = |v_sc| * êₛ (velocity direction follows look direction)
+    velocity_vector_sc = velocity_sc * pset.data["look_direction"]
+
+    # Calculate the ENA velocity vector in the heliosphere frame
+    # v⃗_helio = v⃗_sc - U⃗_sc (simple velocity addition)
+    velocity_vector_helio = velocity_vector_sc - pset.data["sc_velocity"]
+
+    # Convert to spherical coordinates to get ENA source directions
+    ena_source_direction_helio = geometry.cartesian_to_spherical(
+        velocity_vector_helio.data
+    )
+
+    # Update the PSET hae_longitude and hae_latitude variables with the new
+    # energy-dependent values.
+    pset.data["hae_longitude"] = (
+        pset.data["energy_sc"].dims,
+        ena_source_direction_helio[..., 1],
+    )
+    pset.data["hae_latitude"] = (
+        pset.data["energy_sc"].dims,
+        ena_source_direction_helio[..., 2],
+    )
+
+    # 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
+    ram_mask = (
+        np.einsum(
+            "...i,...i->...", velocity_vector_helio, pset.data["sc_direction_vector"]
+        )
+        >= 0
+    )
+    pset.data["ram_mask"] = xr.DataArray(
+        ram_mask,
+        dims=velocity_vector_helio.dims[:-1],
+    )
+
+
+def apply_compton_getting_correction(
+    pset: LoHiBasePointingSet,
+    energy_hf: xr.DataArray,
+) -> None:
+    """
+    Apply Compton-Getting correction to a pointing set and update coordinates.
+
+    This function performs the Compton-Getting velocity transformation to correct
+    ENA observations for the motion of the spacecraft through the heliosphere.
+    The corrected coordinates represent the true source directions of the ENAs
+    in the heliosphere frame.
+
+    The pointing set is modified in-place: new variables are added to the dataset
+    for the corrected coordinates and energies, and the az_el_points attribute
+    is updated to use the corrected coordinates for binning.
+
+    All calculations are performed using xarray DataArrays to preserve dimension
+    information throughout the computation.
+
+    Parameters
+    ----------
+    pset : LoHiBasePointingSet
+        Pointing set object containing HAE longitude/latitude coordinates.
+    energy_hf : xr.DataArray
+        ENA energies in the heliosphere frame in eV. Must be 1D with an
+        energy dimension.
+
+    Notes
+    -----
+    This function adds the following variables to the pointing set 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)
+    - "ena_source_hae_longitude": ENA source longitudes in heliosphere frame (degrees)
+    - "ena_source_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)
+
+    # Step 2: Calculate and add look direction vectors to pset
+    _add_cartesian_look_direction(pset)
+
+    # Step 3: Apply Compton-Getting transformation
+    _calculate_compton_getting_transform(pset, energy_hf)
+
+    # Step 4: Update az_el_points to use the corrected coordinates
+    pset.update_az_el_points()