SPICE geometry enums and IMAP state function

imap_processing/spice/geometry.py

@@ -1 +1,89 @@
-"""Functions for computing geometry using SPICE."""
+"""
+Functions for computing geometry, many of which use SPICE.
+
+Paradigms for developing this module:
+* Use @ensure_spice decorator on functions that directly wrap spiceypy functions
+* Vectorize everything at the lowest level possible (e.g. the decorated spiceypy
+  wrapper function)
+* Always return numpy arrays for vectorized calls.
+"""
+
+import typing
+from enum import IntEnum
+from typing import Union
+
+import numpy as np
+import spiceypy as spice
+
+from imap_processing.spice.kernels import ensure_spice
+
+
+class SpiceBody(IntEnum):
+    """Enum containing SPICE IDs for bodies that we use."""
+
+    # 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_DPS = -43901
+    # Standard NAIF bodies
+    SOLAR_SYSTEM_BARYCENTER = spice.bodn2c("SOLAR_SYSTEM_BARYCENTER")
+    SUN = spice.bodn2c("SUN")
+    EARTH = spice.bodn2c("EARTH")
+
+
+class SpiceFrame(IntEnum):
+    """Enum containing SPICE IDs for reference frames, defined in imap_wkcp.tf."""
+
+    # Standard SPICE Frames
+    J2000 = spice.irfnum("J2000")
+    ECLIPJ2000 = spice.irfnum("ECLIPJ2000")
+    # IMAP specific as defined in imap_wkcp.tf
+    IMAP_SPACECRAFT = -43000
+    IMAP_LO_BASE = -43100
+    IMAP_LO_STAR_SENSOR = -43103
+    IMAP_LO = -43105
+    IMAP_HI_45 = -43150
+    IMAP_HI_90 = -43160
+    IMAP_ULTRA_45 = -43200
+    IMAP_ULTRA_90 = -43210
+    IMAP_MAG = -43250
+    IMAP_SWE = -43300
+    IMAP_SWAPI = -43350
+    IMAP_CODICE = -43400
+    IMAP_HIT = -43500
+    IMAP_IDEX = -43700
+    IMAP_GLOWS = -43750
+
+
+@typing.no_type_check
+@ensure_spice
+def imap_state(
+    et: Union[np.ndarray, float],
+    ref_frame: SpiceFrame = SpiceFrame.ECLIPJ2000,
+    observer: SpiceBody = SpiceBody.SUN,
+) -> np.ndarray:
+    """
+    Get the state (position and velocity) of the IMAP spacecraft.
+
+    By default, the state is returned in the ECLIPJ2000 frame as observed by the Sun.
+
+    Parameters
+    ----------
+    et : np.ndarray or float
+        Epoch time(s) [J2000 seconds] to get the IMAP state for.
+    ref_frame : SpiceFrame, optional
+        Reference frame which the IMAP state is expressed in.
+    observer : SpiceBody, optional
+        Observing body.
+
+    Returns
+    -------
+    state : np.ndarray
+     The Cartesian state vector representing the position and velocity of the
+     IMAP spacecraft.
+    """
+    state, _ = spice.spkezr(
+        SpiceBody.IMAP.name, et, ref_frame.name, "NONE", observer.name
+    )
+    return np.asarray(state)

imap_processing/spice/kernels.py

@@ -6,7 +6,7 @@
 from collections.abc import Generator
 from contextlib import contextmanager
 from pathlib import Path
-from typing import Any, Callable, Optional
+from typing import Any, Callable, Optional, Union, overload
 
 import numpy as np
 import spiceypy as spice
@@ -16,15 +16,29 @@
 logger = logging.getLogger(__name__)
 
 
+# Declarations to help with typing. Taken from mypy documentation on
+# decorator-factories:
+# https://mypy.readthedocs.io/en/stable/generics.html#decorator-factories
+# Bare decorator usage
+@overload
 def ensure_spice(
-    f_py: Optional[Callable] = None, time_kernels_only: bool = False
-) -> Callable:
+    __func: Callable[..., Any],
+) -> Callable[..., Any]: ...  # numpydoc ignore=GL08
+# Decorator with arguments
+@overload
+def ensure_spice(
+    *, time_kernels_only: bool = False
+) -> Callable[[Callable[..., Any]], Callable[..., Any]]: ...  # numpydoc ignore=GL08
+# Implementation
+def ensure_spice(
+    __func: Optional[Callable[..., Any]] = None, *, time_kernels_only: bool = False
+) -> Union[Callable[..., Any], Callable[[Callable[..., Any]], Callable[..., Any]]]:
     """
     Decorator/wrapper that automatically furnishes SPICE kernels.
 
     Parameters
     ----------
-    f_py : Callable
+    __func : Callable
         The function requiring SPICE that we are going to wrap if being used
         explicitly, otherwise None, in which case ensure_spice is being used,
         not as a function wrapper (see l2a_processing.py) but as a true
@@ -82,11 +96,6 @@ def ensure_spice(
         >>> wrapped = ensure_spice(spicey_func, time_kernels_only=True)
         ... result = wrapped(*args, **kwargs)
     """
-    if f_py and not callable(f_py):
-        raise ValueError(
-            f"Received a non-callable object {f_py} as the f_py argument to"
-            f"ensure_spice.  f_py must be a callable object."
-        )
 
     def _decorator(func: Callable[..., Callable]) -> Callable:
         """
@@ -157,8 +166,8 @@ def wrapper_ensure_spice(*args: Any, **kwargs: Any) -> Any:
     # Note: This return was originally implemented as a ternary operator, but
     # this caused mypy to fail due to this bug:
     # https://github.com/python/mypy/issues/4134
-    if callable(f_py):
-        return _decorator(f_py)
+    if callable(__func):
+        return _decorator(__func)
     else:
         return _decorator
 

imap_processing/tests/spice/test_data/imap_test_metakernel.template

@@ -1,2 +1,3 @@
 {SPICE_TEST_DATA_PATH}/imap_sclk_0000.tsc
-{SPICE_TEST_DATA_PATH}/naif0012.tls
+{SPICE_TEST_DATA_PATH}/naif0012.tls
+{SPICE_TEST_DATA_PATH}/imap_spk_demo.bsp

imap_processing/tests/spice/test_geometry.py

@@ -1 +1,25 @@
 """Tests coverage for imap_processing/spice/geometry.py"""
+
+import numpy as np
+import pytest
+
+from imap_processing.spice.geometry import (
+    SpiceBody,
+    imap_state,
+)
+
+
+@pytest.mark.parametrize(
+    "et",
+    [
+        798033670,
+        np.linspace(798033670, 798033770),
+    ],
+)
+def test_imap_state(et, use_test_metakernel):
+    """Test coverage for imap_state()"""
+    state = imap_state(et, observer=SpiceBody.EARTH)
+    if hasattr(et, "__len__"):
+        np.testing.assert_array_equal(state.shape, (len(et), 6))
+    else:
+        assert state.shape == (6,)