feat: interpolated integration

Signed-off-by: nstarman <nstarman@users.noreply.github.com>

pyproject.toml

@@ -238,6 +238,7 @@ messages_control.disable = [
   "protected-access",              # ruff SLF001
   "unnecessary-ellipsis",
   "unnecessary-lambda-assignment", # ruff E731
+  "unnecessary-pass",              # handled by ruff
   "wrong-import-position",
   "wrong-import-order",            # handled by ruff
 ]

src/galax/coordinates/_psp/interp.py

@@ -22,9 +22,22 @@ class PhaseSpacePositionInterpolant(Protocol):
     """Protocol for interpolating phase-space positions."""
 
     units: AbstractUnitSystem
+    """The unit system for the interpolation."""
 
-    def __call__(self, t: gt.VecTime) -> gt.BatchVecTime6:
-        pass
+    def __call__(self, t: gt.QVecTime) -> PhaseSpacePosition:
+        """Evaluate the interpolation.
+
+        Parameters
+        ----------
+        t : Quantity[float, (time,), 'time']
+            The times at which to evaluate the interpolation.
+
+        Returns
+        -------
+        :class:`galax.coordinates.PhaseSpacePosition`
+            The interpolated phase-space positions.
+        """
+        ...
 
 
 @final
@@ -53,7 +66,7 @@ class InterpolatedPhaseSpacePosition(AbstractPhaseSpacePosition):
     as `q` and `p`.
     """
 
-    interpolation: PhaseSpacePositionInterpolant
+    interpolant: PhaseSpacePositionInterpolant
     """The interpolation function."""
 
     def __post_init__(self) -> None:
@@ -65,13 +78,7 @@ def __post_init__(self) -> None:
 
     def __call__(self, t: gt.BatchFloatQScalar) -> PhaseSpacePosition:
         """Call the interpolation."""
-        qp = self.interpolation(t)
-        units = self.interpolation.units
-        return PhaseSpacePosition(
-            q=Quantity(qp[..., 0:3], units["length"]),
-            p=Quantity(qp[..., 3:6], units["speed"]),
-            t=t,
-        )
+        return self.interpolant(t)
 
     # ==========================================================================
     # Array properties

src/galax/dynamics/_dynamics/integrate/_api.py

@@ -1,11 +1,11 @@
 __all__ = ["Integrator"]
 
-from typing import Any, Protocol, TypeAlias, runtime_checkable
+from typing import Any, Literal, Protocol, TypeAlias, runtime_checkable
 
 from unxt import AbstractUnitSystem
 
+import galax.coordinates as gc
 import galax.typing as gt
-from galax.coordinates import AbstractPhaseSpacePosition, PhaseSpacePosition
 from galax.utils.dataclasses import _DataclassInstance
 
 SaveT: TypeAlias = gt.BatchQVecTime | gt.QVecTime | gt.BatchVecTime | gt.VecTime
@@ -52,14 +52,15 @@ class Integrator(_DataclassInstance, Protocol):
     def __call__(
         self,
         F: FCallable,
-        w0: AbstractPhaseSpacePosition | gt.BatchVec6,
+        w0: gc.AbstractPhaseSpacePosition | gt.BatchVec6,
         t0: gt.FloatQScalar | gt.FloatScalar,
         t1: gt.FloatQScalar | gt.FloatScalar,
         /,
         savet: SaveT | None = None,
         *,
         units: AbstractUnitSystem,
-    ) -> PhaseSpacePosition:
+        interpolated: Literal[False, True] = False,
+    ) -> gc.PhaseSpacePosition | gc.InterpolatedPhaseSpacePosition:
         """Integrate.
 
         Parameters

src/galax/dynamics/_dynamics/integrate/_base.py

@@ -1,14 +1,15 @@
 __all__ = ["AbstractIntegrator"]
 
 import abc
+from typing import Literal
 
 import equinox as eqx
 
 from unxt import AbstractUnitSystem
 
+import galax.coordinates as gc
 import galax.typing as gt
 from ._api import FCallable
-from galax.coordinates import AbstractPhaseSpacePosition, PhaseSpacePosition
 
 
 class AbstractIntegrator(eqx.Module, strict=True):  # type: ignore[call-arg, misc]
@@ -28,7 +29,7 @@ class AbstractIntegrator(eqx.Module, strict=True):  # type: ignore[call-arg, mis
     def __call__(
         self,
         F: FCallable,
-        w0: AbstractPhaseSpacePosition | gt.BatchVec6,
+        w0: gc.AbstractPhaseSpacePosition | gt.BatchVec6,
         t0: gt.FloatQScalar | gt.FloatScalar,
         t1: gt.FloatQScalar | gt.FloatScalar,
         /,
@@ -37,7 +38,8 @@ def __call__(
         ) = None,
         *,
         units: AbstractUnitSystem,
-    ) -> PhaseSpacePosition:
+        interpolated: Literal[False, True] = False,
+    ) -> gc.PhaseSpacePosition | gc.InterpolatedPhaseSpacePosition:
         """Run the integrator.
 
         Parameters