diff --git a/imap_processing/spice/geometry.py b/imap_processing/spice/geometry.py index 0f04d2db6..ca7ab4fe3 100644 --- a/imap_processing/spice/geometry.py +++ b/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) diff --git a/imap_processing/spice/kernels.py b/imap_processing/spice/kernels.py index 6f7ce1492..0a3f18493 100644 --- a/imap_processing/spice/kernels.py +++ b/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 diff --git a/imap_processing/tests/spice/test_data/imap_spk_demo.bsp b/imap_processing/tests/spice/test_data/imap_spk_demo.bsp new file mode 100644 index 000000000..e8b7578f6 Binary files /dev/null and b/imap_processing/tests/spice/test_data/imap_spk_demo.bsp differ diff --git a/imap_processing/tests/spice/test_data/imap_test_metakernel.template b/imap_processing/tests/spice/test_data/imap_test_metakernel.template index b4a5d9f82..d727ba318 100644 --- a/imap_processing/tests/spice/test_data/imap_test_metakernel.template +++ b/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 \ No newline at end of file +{SPICE_TEST_DATA_PATH}/naif0012.tls +{SPICE_TEST_DATA_PATH}/imap_spk_demo.bsp \ No newline at end of file diff --git a/imap_processing/tests/spice/test_geometry.py b/imap_processing/tests/spice/test_geometry.py index 287b3eb28..f1a0ae448 100644 --- a/imap_processing/tests/spice/test_geometry.py +++ b/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,)