859 spice frame transform function

imap_processing/spice/geometry.py

@@ -28,7 +28,7 @@ class SpiceBody(IntEnum):
     # A subset of IMAP Specific bodies as defined in imap_wkcp.tf
     IMAP = -43
     IMAP_SPACECRAFT = -43000
-    # IMAP Pointing Frame (Despun) as defined in iamp_science_0001.tf
+    # IMAP Pointing Frame (Despun) as defined in imap_science_0001.tf
     IMAP_DPS = -43901
     # Standard NAIF bodies
     SOLAR_SYSTEM_BARYCENTER = spice.bodn2c("SOLAR_SYSTEM_BARYCENTER")
@@ -58,6 +58,8 @@ class SpiceFrame(IntEnum):
     IMAP_HIT = -43500
     IMAP_IDEX = -43700
     IMAP_GLOWS = -43750
+    # IMAP Pointing Frame (Despun) as defined in iamp_science_0001.tf
+    IMAP_DPS = -43901
 
 
 @typing.no_type_check
@@ -213,3 +215,108 @@ def get_spacecraft_spin_phase(
     if is_scalar:
         return spin_phases[0]
     return spin_phases
+
+
+@typing.no_type_check
+@ensure_spice
+def frame_transform(
+    et: Union[float, npt.NDArray],
+    position: npt.NDArray,
+    from_frame: SpiceFrame,
+    to_frame: SpiceFrame,
+) -> npt.NDArray:
+    """
+    Transform an <x, y, z> vector between reference frames (rotation only).
+
+    This function is a vectorized equivalent to performing the following SPICE
+    calls for each input time and position vector to perform the transform.
+    The matrix multiplication step is done using `numpy.matmul` rather than
+    `spice.mxv`.
+    >>> rotation_matrix = spice.pxform(from_frame, to_frame, et)
+    ... result = spice.mxv(rotation_matrix, position)
+
+    Parameters
+    ----------
+    et : float or npt.NDArray
+        Ephemeris time(s) corresponding to position(s).
+    position : npt.NDArray
+        <x, y, z> vector or array of vectors in reference frame `from_frame`.
+    from_frame : SpiceFrame
+        Reference frame of input vector(s).
+    to_frame : SpiceFrame
+        Reference frame of output vector(s).
+
+    Returns
+    -------
+    result : npt.NDArray
+        3d position vector(s) in reference frame `to_frame`.
+    """
+    if position.ndim == 1:
+        if not len(position) == 3:
+            raise ValueError(
+                "Position vectors with one dimension must have 3 elements."
+            )
+        if not isinstance(et, float):
+            raise ValueError(
+                "Ephemeris time must be float when single position vector is provided."
+            )
+    elif position.ndim == 2:
+        if not position.shape[1] == 3:
+            raise ValueError(
+                f"Invalid position shape: {position.shape}. "
+                f"Each input position vector must have 3 elements."
+            )
+        if not len(position) == len(et):
+            raise ValueError(
+                "Mismatch in number of position vectors and Ephemeris times provided."
+                f"Position has {len(position)} elements and et has {len(et)} elements."
+            )
+
+    # rotate will have shape = (3, 3) or (n, 3, 3)
+    # position will have shape = (3,) or (n, 3)
+    rotate = get_rotation_matrix(et, from_frame, to_frame)
+    # adding a dimension to position results in the following input and output
+    # shapes from matrix multiplication
+    # Single et/position:      (3, 3),(3, 1) -> (3, 1)
+    # Multiple et/positions :  (n, 3, 3),(n, 3, 1) -> (n, 3, 1)
+    result = np.squeeze(rotate @ position[..., np.newaxis])
+
+    return result
+
+
+def get_rotation_matrix(
+    et: Union[float, npt.NDArray],
+    from_frame: SpiceFrame,
+    to_frame: SpiceFrame,
+) -> npt.NDArray:
+    """
+    Get the rotation matrix/matrices that can be used to transform between frames.
+
+    This is a vectorized wrapper around `spiceypy.pxform`
+    "Return the matrix that transforms position vectors from one specified frame
+    to another at a specified epoch."
+    https://naif.jpl.nasa.gov/pub/naif/toolkit_docs/C/cspice/pxform_c.html
+
+    Parameters
+    ----------
+    et : float or npt.NDArray
+        Ephemeris time(s) for which to get the rotation matrices.
+    from_frame : SpiceFrame
+        Reference frame to transform from.
+    to_frame : SpiceFrame
+        Reference frame to transform to.
+
+    Returns
+    -------
+    rotation : npt.NDArray
+        If et is a float, the returned rotation matrix is of shape (3, 3). If
+        et is a np.ndarray, the returned rotation matrix is of shape (n, 3, 3)
+        where n matches the number of elements in et.
+    """
+    vec_pxform = np.vectorize(
+        spice.pxform,
+        excluded=["fromstr", "tostr"],
+        signature="(),(),()->(3,3)",
+        otypes=[np.float64],
+    )
+    return vec_pxform(from_frame.name, to_frame.name, et)

imap_processing/tests/conftest.py

@@ -4,6 +4,8 @@
 import os
 import re
 import time
+from contextlib import contextmanager
+from pathlib import Path
 from typing import Optional
 
 import cdflib
@@ -123,6 +125,18 @@ def furnish_sclk(spice_test_data_path):
     spice.kclear()
 
 
+@pytest.fixture()
+def furnish_kernels(spice_test_data_path):
+    """Return a function that will furnish an arbitrary list of kernels."""
+
+    @contextmanager
+    def furnish_kernels(kernels: list[Path]):
+        with spice.KernelPool([str(spice_test_data_path / k) for k in kernels]) as pool:
+            yield pool
+
+    return furnish_kernels
+
+
 @pytest.fixture(scope="session")
 def monkeypatch_session():
     from _pytest.monkeypatch import MonkeyPatch

imap_processing/tests/spice/test_geometry.py

@@ -3,9 +3,13 @@
 import numpy as np
 import pandas as pd
 import pytest
+import spiceypy as spice
 
 from imap_processing.spice.geometry import (
     SpiceBody,
+    SpiceFrame,
+    frame_transform,
+    get_rotation_matrix,
     get_spacecraft_spin_phase,
     get_spin_data,
     imap_state,
@@ -106,3 +110,105 @@ def test_get_spin_data():
         "thruster_firing",
         "spin_start_time",
     }, "Spin data must have the specified fields."
+
+
+def test_frame_transform(furnish_kernels):
+    """Test transformation of vectors from one frame to another, with the option
+    to normalize the result."""
+    # This test requires an IMAP attitude kernel and pointing (despun) kernel
+    kernels = [
+        "naif0012.tls",
+        "imap_wkcp.tf",
+        "imap_science_0001.tf",
+        "sim_1yr_imap_attitude.bc",
+        "sim_1yr_imap_pointing_frame.bc",
+    ]
+    with furnish_kernels(kernels):
+        # Test single et and position calculation
+        et_0 = spice.utc2et("2025-04-30T12:00:00.000")
+        position = np.arange(3) + 1
+        result_0 = frame_transform(
+            et_0, position, SpiceFrame.IMAP_ULTRA_45, SpiceFrame.IMAP_DPS
+        )
+        # compare against pure SPICE calculation
+        rotation_matrix = spice.pxform(
+            SpiceFrame.IMAP_ULTRA_45.name, SpiceFrame.IMAP_DPS.name, et_0
+        )
+        spice_result = spice.mxv(rotation_matrix, position)
+        np.testing.assert_allclose(result_0, spice_result, atol=1e-12)
+
+        # test multiple et and position calculation
+        ets = np.array([et_0, et_0 + 10])
+        positions = np.array([[1, 1, 1], [1, 2, 3]])
+        vec_result = frame_transform(
+            ets, positions, SpiceFrame.IMAP_HI_90, SpiceFrame.IMAP_DPS
+        )
+
+        assert vec_result.shape == (2, 3)
+        # compare with direct spice calculations
+        for et, pos, result in zip(ets, positions, vec_result):
+            rotation_matrix = spice.pxform(
+                SpiceFrame.IMAP_HI_90.name, SpiceFrame.IMAP_DPS.name, et
+            )
+            spice_result = spice.mxv(rotation_matrix, pos)
+            np.testing.assert_allclose(result, spice_result, atol=1e-12)
+
+
+def test_frame_transform_exceptions():
+    """Test that the proper exceptions get raised when input arguments are invalid."""
+    with pytest.raises(
+        ValueError, match="Position vectors with one dimension must have 3 elements."
+    ):
+        frame_transform(
+            0, np.arange(4), SpiceFrame.IMAP_SPACECRAFT, SpiceFrame.IMAP_CODICE
+        )
+    with pytest.raises(
+        ValueError,
+        match="Ephemeris time must be float when single position vector is provided.",
+    ):
+        frame_transform(
+            np.asarray(0),
+            np.array([1, 0, 0]),
+            SpiceFrame.ECLIPJ2000,
+            SpiceFrame.IMAP_HIT,
+        )
+    with pytest.raises(ValueError, match="Invalid position shape: "):
+        frame_transform(
+            np.arange(2),
+            np.arange(4).reshape((2, 2)),
+            SpiceFrame.ECLIPJ2000,
+            SpiceFrame.IMAP_HIT,
+        )
+    with pytest.raises(
+        ValueError,
+        match="Mismatch in number of position vectors and Ephemeris times provided.",
+    ):
+        frame_transform(
+            np.arange(2),
+            np.arange(9).reshape((3, 3)),
+            SpiceFrame.ECLIPJ2000,
+            SpiceFrame.IMAP_HIT,
+        )
+
+
+def test_get_rotation_matrix(furnish_kernels):
+    """Test coverage for get_rotation_matrix()."""
+    kernels = [
+        "naif0012.tls",
+        "imap_wkcp.tf",
+        "imap_science_0001.tf",
+        "sim_1yr_imap_attitude.bc",
+        "sim_1yr_imap_pointing_frame.bc",
+    ]
+    with furnish_kernels(kernels):
+        et = spice.utc2et("2025-09-30T12:00:00.000")
+        # test input of float
+        rotation = get_rotation_matrix(
+            et, SpiceFrame.IMAP_IDEX, SpiceFrame.IMAP_SPACECRAFT
+        )
+        assert rotation.shape == (3, 3)
+        # test array of et input
+        rotation = get_rotation_matrix(
+            np.arange(10) + et, SpiceFrame.IMAP_IDEX, SpiceFrame.IMAP_SPACECRAFT
+        )
+        assert rotation.shape == (10, 3, 3)