Extract reusable compute_pointing_directions from set_pointing_directions

Generalize the Lo L1C pointing-direction geometry so it can be reused by
other products (e.g. the l1c quickmap) without duplicating the DPS az/el ->
sky transform.

Author: vineetbansal

imap_processing/lo/l1c/lo_l1c.py

@@ -780,6 +780,63 @@ def set_background_rates(
     return bg_rates_data, bg_stat_uncert_data, bg_sys_err_data
 
 
+def compute_pointing_directions(
+    epoch: float,
+    pivot_angle: float,
+    spin_angles: np.ndarray | None = None,
+    off_angles: np.ndarray | None = None,
+    to_frame: SpiceFrame = SpiceFrame.IMAP_HAE,
+) -> np.ndarray:
+    """
+    Transform a DPS spin/off-angle grid to sky longitude/latitude at an epoch.
+
+    For the given ``epoch`` the instrument despun (IMAP_DPS) spin-angle / off-angle
+    grid is transformed into ``to_frame`` longitude/latitude using SPICE. Off-angles
+    are measured relative to ``pivot_angle`` (elevation = 90 - pivot_angle + off).
+
+    This is the reusable geometry core shared by pointing-direction products: pass
+    the desired spin/off-angle bin centers and destination frame.
+
+    Parameters
+    ----------
+    epoch : float
+        The epoch time in TTJ2000ns.
+    pivot_angle : float
+        The pivot angle in degrees.
+        Off-angles are adjusted relative to this pivot angle before transformation.
+    spin_angles : numpy.ndarray, optional
+        Spin-angle bin centers in degrees. Defaults to ``SPIN_ANGLE_BIN_CENTERS``
+        (the 3600-bin PSET grid).
+    off_angles : numpy.ndarray, optional
+        Off-angle bin centers in degrees. Defaults to ``OFF_ANGLE_BIN_CENTERS``
+        (the 40-bin PSET grid).
+    to_frame : SpiceFrame, optional
+        Destination reference frame. Defaults to ``SpiceFrame.IMAP_HAE``.
+
+    Returns
+    -------
+    numpy.ndarray
+        Array of shape ``(n_spin, n_off, 2)`` where ``[..., 0]`` is longitude and
+        ``[..., 1]`` is latitude, both in degrees in ``to_frame``.
+    """
+    if spin_angles is None:
+        spin_angles = SPIN_ANGLE_BIN_CENTERS
+    if off_angles is None:
+        off_angles = OFF_ANGLE_BIN_CENTERS
+
+    et = ttj2000ns_to_et(epoch)
+    # create a meshgrid of spin and off angles using the bin centers
+    spin, off = np.meshgrid(spin_angles, off_angles, indexing="ij")
+    # off_angles need to account for the pivot_angle
+    off = off + (90 - pivot_angle)
+    dps_az_el = np.stack([spin, off], axis=-1)
+
+    # Transform from DPS Az/El to the destination frame's lon/lat
+    return frame_transform_az_el(
+        et, dps_az_el, SpiceFrame.IMAP_DPS, to_frame, degrees=True
+    )
+
+
 def set_pointing_directions(
     epoch: float,
     attr_mgr: ImapCdfAttributes,
@@ -809,19 +866,7 @@ def set_pointing_directions(
     hae_latitude : xr.DataArray
         The HAE latitude for each spin and off angle bin.
     """
-    et = ttj2000ns_to_et(epoch)
-    # create a meshgrid of spin and off angles using the bin centers
-    spin, off = np.meshgrid(
-        SPIN_ANGLE_BIN_CENTERS, OFF_ANGLE_BIN_CENTERS, indexing="ij"
-    )
-    # off_angles need to account for the pivot_angle
-    off += 90 - pivot_angle
-    dps_az_el = np.stack([spin, off], axis=-1)
-
-    # Transform from DPS Az/El to HAE lon/lat
-    hae_az_el = frame_transform_az_el(
-        et, dps_az_el, SpiceFrame.IMAP_DPS, SpiceFrame.IMAP_HAE, degrees=True
-    )
+    hae_az_el = compute_pointing_directions(epoch, pivot_angle)
 
     return xr.DataArray(
         data=hae_az_el[np.newaxis, :, :, 0].astype(np.float64),

imap_processing/tests/lo/test_lo_l1c.py

@@ -14,12 +14,14 @@
     PSET_SHAPE,
     FilterType,
     calculate_exposure_times,
+    compute_pointing_directions,
     create_pset_counts,
     filter_goodtimes,
     lo_l1c,
     set_background_rates,
     set_pointing_directions,
 )
+from imap_processing.spice.geometry import SpiceFrame
 from imap_processing.spice.time import met_to_ttj2000ns
 
 
@@ -729,3 +731,80 @@ def test_set_pointing_directions_pivot_angle(attr_mgr, pivot_angle):
         # dps_az_el[:, :, 1] should have the adjusted off angles repeated across spin
         actual_off_angles = dps_az_el[0, :, 1]  # Take first spin angle
         np.testing.assert_allclose(actual_off_angles, expected_off_angles, rtol=1e-10)
+
+
+def test_compute_pointing_directions_defaults():
+    """Default grid/frame reproduce the PSET (3600 x 40) IMAP_DPS->IMAP_HAE case."""
+    mock_et = 123456789.0
+    mock_az_el = np.stack(
+        np.meshgrid(np.arange(3600), np.arange(40), indexing="ij"), axis=-1
+    )
+    with (
+        patch("imap_processing.lo.l1c.lo_l1c.ttj2000ns_to_et") as mock_ttj2000ns_to_et,
+        patch(
+            "imap_processing.lo.l1c.lo_l1c.frame_transform_az_el"
+        ) as mock_frame_transform,
+    ):
+        mock_ttj2000ns_to_et.return_value = mock_et
+        mock_frame_transform.return_value = mock_az_el
+
+        result = compute_pointing_directions(1000000000.0, 90)
+
+        # Returns the raw (n_spin, n_off, 2) array, not a DataArray.
+        assert result.shape == (3600, 40, 2)
+        call_args = mock_frame_transform.call_args
+        assert call_args[0][1].shape == (3600, 40, 2)  # dps_az_el grid
+        assert call_args[0][2] == SpiceFrame.IMAP_DPS  # from_frame
+        assert call_args[0][3] == SpiceFrame.IMAP_HAE  # default to_frame
+
+
+def test_compute_pointing_directions_custom_grid_and_frame():
+    """Custom spin/off angles and destination frame are honored."""
+    spin_angles = np.arange(3.0, 360.0, 6.0)  # 60 bins
+    off_angles = np.array([0.0])  # single pivot-cone off-angle
+    pivot_angle = 75.0
+    with (
+        patch("imap_processing.lo.l1c.lo_l1c.ttj2000ns_to_et") as mock_ttj2000ns_to_et,
+        patch(
+            "imap_processing.lo.l1c.lo_l1c.frame_transform_az_el"
+        ) as mock_frame_transform,
+    ):
+        mock_ttj2000ns_to_et.return_value = 123456789.0
+        mock_frame_transform.side_effect = lambda et, az_el, *a, **k: az_el
+
+        result = compute_pointing_directions(
+            1000000000.0,
+            pivot_angle,
+            spin_angles=spin_angles,
+            off_angles=off_angles,
+            to_frame=SpiceFrame.ECLIPJ2000,
+        )
+
+        assert result.shape == (60, 1, 2)
+        # Spin component matches the requested spin angles.
+        np.testing.assert_allclose(result[:, 0, 0], spin_angles)
+        # Off component is the single off-angle offset by (90 - pivot_angle).
+        np.testing.assert_allclose(result[:, 0, 1], 90 - pivot_angle)
+        # Destination frame is forwarded.
+        assert mock_frame_transform.call_args[0][3] == SpiceFrame.ECLIPJ2000
+
+
+def test_set_pointing_directions_delegates(attr_mgr):
+    """set_pointing_directions wraps compute_pointing_directions output unchanged."""
+    mock_az_el = np.stack(
+        np.meshgrid(np.arange(3600), np.arange(40), indexing="ij"), axis=-1
+    ).astype(float)
+    with patch(
+        "imap_processing.lo.l1c.lo_l1c.compute_pointing_directions"
+    ) as mock_compute:
+        mock_compute.return_value = mock_az_el
+
+        hae_longitude, hae_latitude = set_pointing_directions(
+            1000000000.0, attr_mgr, 90
+        )
+
+        mock_compute.assert_called_once_with(1000000000.0, 90)
+        assert hae_longitude.dims == ("epoch", "spin_angle", "off_angle")
+        assert hae_longitude.shape == (1, 3600, 40)
+        np.testing.assert_array_equal(hae_longitude.values[0], mock_az_el[:, :, 0])
+        np.testing.assert_array_equal(hae_latitude.values[0], mock_az_el[:, :, 1])