Source code for empyrean.states

"""Query body states from SPK ephemeris."""

import numpy as np

from empyrean._convert import frame_to_int, int_to_frame, naif_to_origin, origin_to_naif
from empyrean.coordinates.coordinates import CartesianCoordinates
from empyrean.coordinates.enums import Frame, Origin
from empyrean.coordinates.epoch import Epochs, _epochs_mjd_tdb




def get_states(
    target: Origin | str,
    center: Origin | str,
    epochs: Epochs,
    frame: Frame | str = Frame.ECLIPTICJ2000,
) -> CartesianCoordinates:
    """Query Cartesian states of a body from the SPK ephemeris.

    Parameters
    ----------
    target : Origin | str
        Target body. Pass an :class:`Origin` (e.g. ``Origin.EARTH``) or
        the canonical name (e.g. ``"Earth"``).
    center : Origin | str
        Center body. Same shape as ``target``.
    epochs : Epochs
        Epochs table, converted to TDB internally. Build one with
        ``Epochs.from_mjd(values, scale="tdb")`` (or ``scale="utc"``);
        a bare array is refused, as it carries no time scale.
    frame : Frame | str
        Reference frame (default: EclipticJ2000).

    Returns
    -------
    CartesianCoordinates
        Cartesian states (AU, AU/day) at each epoch.

    Raises
    ------
    TypeError
        ``epochs`` is not an :class:`Epochs` table.
    """
    from empyrean._empyrean_rs import _get_states

    target_naif = origin_to_naif(target)
    center_naif = origin_to_naif(center)
    frame_int = frame_to_int(frame)

    epochs_mjd = _epochs_mjd_tdb(epochs, "get_states() epochs")

    result = _get_states(target_naif, center_naif, epochs_mjd, frame_int)

    out_frame = int_to_frame(int(result["frame"][0])) if len(result["frame"]) > 0 else frame
    origin_strs = [naif_to_origin(int(o)) for o in result["origin"]]

    return CartesianCoordinates.from_kwargs(
        epoch=np.asarray(result["epoch"]),
        x=np.asarray(result["x"]),
        y=np.asarray(result["y"]),
        z=np.asarray(result["z"]),
        vx=np.asarray(result["vx"]),
        vy=np.asarray(result["vy"]),
        vz=np.asarray(result["vz"]),
        frame=out_frame.value if hasattr(out_frame, "value") else out_frame,
        origin=origin_strs,
    )