SPICE geometry enums and IMAP state function (IMAP-Science-Operations…

…-Center#791)

* Add spice infrastructure code for spice bodies and frames
Add geometry functions for computing IMAP spacecraft state (position, velocity)
Add IMAP SPK kernel from Nick's 366 day simulation
Add test coverage for new geometry functions

* Remove _ensured_spkezr funciton.
Remove SpiceId(NamedTuple)
Parameterize test to cover single and iterable ET input
Change typing on ensure_spice to apease mypy
Fix up ensure_spice decorator factory typing

* Fix typo

* Use IntEnum

* Back out some of the typing changes not compatible with Python 3.9

* Use spice to define built-in bodies and frames

* Tell mypy not to check skip type check

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,)