♻️ refactor(dynamics): evaluate_orbit uses DynamicsSolver

Author: nstarman

pyproject.toml

@@ -32,7 +32,7 @@
     "coordinax>=0.19.0",
     "dataclassish>=0.7.1",
     "diffrax>=0.6",
-    "diffraxtra>=1.0.2",
+    "diffraxtra>=1.1.0",
     "equinox>=0.11.8",
     "is-annotated>=1.0",
     "jax>=0.4.35",

src/galax/dynamics/_src/integrate/__init__.py

@@ -12,7 +12,6 @@
     "Interpolant",
 ]
 
-from .funcs import evaluate_orbit
 from .integrator import Integrator
 from .interp import Interpolant
 from .interp_psp import InterpolatedPhaseSpacePosition

src/galax/dynamics/_src/mockstream/mockstream_generator.py

@@ -20,13 +20,14 @@
 from .core import MockStream, MockStreamArm
 from .df import AbstractStreamDF, ProgenitorMassCallable
 from .utils import cond_reverse
-from galax.dynamics._src.integrate.funcs import _default_integrator, evaluate_orbit
-from galax.dynamics._src.integrate.integrator import Integrator
-from galax.dynamics._src.orbit import Orbit
+from galax.dynamics._src.dynamics import DynamicsSolver
+from galax.dynamics._src.orbit import Orbit, evaluate_orbit
 from galax.potential import AbstractPotential
 
 Carry: TypeAlias = tuple[gt.IntSz0, gt.SzN, gt.SzN]
 
+_default_solver: DynamicsSolver = DynamicsSolver()
+
 
 @final
 class MockStreamGenerator(eqx.Module):  # type: ignore[misc]
@@ -42,13 +43,11 @@ class MockStreamGenerator(eqx.Module):  # type: ignore[misc]
     """Potential in which the progenitor orbits and creates a stream."""
 
     _: KW_ONLY
-    progenitor_integrator: Integrator = eqx.field(
-        default=_default_integrator, static=True
-    )
-    """Integrator for the progenitor orbit."""
+    progenitor_solver: DynamicsSolver = eqx.field(default=_default_solver, static=True)
+    """Solver for the progenitor orbit."""
 
-    stream_integrator: Integrator = eqx.field(default=_default_integrator, static=True)
-    """Integrator for the stream."""
+    stream_solver: DynamicsSolver = eqx.field(default=_default_solver, static=True)
+    """Solver for the stream."""
 
     @property
     def units(self) -> u.AbstractUnitSystem:
@@ -63,9 +62,7 @@ def _progenitor_trajectory(
         """Integrate the progenitor orbit."""
         return cast(
             Orbit,
-            evaluate_orbit(
-                self.potential, w0, ts, integrator=self.progenitor_integrator
-            ),
+            evaluate_orbit(self.potential, w0, ts, integrator=self.progenitor_solver),
         )
 
     # ==========================================================================
@@ -103,7 +100,7 @@ def one_pt_intg(
             def integ_ics(ics: gt.Sz6) -> gt.SzN:
                 # TODO: only return the final state
                 return evaluate_orbit(
-                    self.potential, ics, tstep, integrator=self.stream_integrator
+                    self.potential, ics, tstep, integrator=self.stream_solver
                 ).w(units=self.units)[-1]
 
             # vmap integration over leading and trailing arm
@@ -138,10 +135,10 @@ def one_pt_intg(
         ) -> tuple[gt.Sz6, gt.Sz6]:
             tstep = jnp.asarray([ts[i], t_f])
             w_lead = evaluate_orbit(
-                self.potential, w0_l_i, tstep, integrator=self.stream_integrator
+                self.potential, w0_l_i, tstep, integrator=self.stream_solver
             ).w(units=self.potential.units)[-1]
             w_trail = evaluate_orbit(
-                self.potential, w0_t_i, tstep, integrator=self.stream_integrator
+                self.potential, w0_t_i, tstep, integrator=self.stream_solver
             ).w(units=self.potential.units)[-1]
             return w_lead, w_trail
 

src/galax/dynamics/_src/orbit/__init__.py

@@ -1,6 +1,7 @@
 """Orbits. Private module."""
 
-__all__ = ["Orbit", "plot_components"]
+__all__ = ["Orbit", "evaluate_orbit", "plot_components"]
 
+from .funcs import evaluate_orbit
 from .orbit import Orbit
 from .plot_helper import plot_components

src/galax/dynamics/_src/integrate/funcs.py renamed to src/galax/dynamics/_src/orbit/funcs.py

@@ -2,11 +2,8 @@
 
 __all__ = ["evaluate_orbit"]
 
-from collections.abc import Callable
 from typing import Any, Literal
 
-import jax
-from jaxtyping import Array
 from plum import dispatch
 
 import quaxed.numpy as jnp
@@ -16,17 +13,12 @@
 import galax.dynamics._src.custom_types as gdt
 import galax.potential as gp
 import galax.typing as gt
-from .integrator import Integrator
-from galax.dynamics._src.dynamics import HamiltonianField
-from galax.dynamics._src.orbit import Orbit
+from .interp import PhaseSpaceInterpolation
+from .orbit import Orbit
+from galax.dynamics._src.dynamics import DynamicsSolver, HamiltonianField
 
 # TODO: enable setting the default integrator
-_default_integrator: Integrator = Integrator()
-
-
-_select_w0: Callable[[Array, Array, Array], Array] = jax.numpy.vectorize(
-    jax.lax.select, signature="(),(6),(6)->(6)"
-)
+_default_solver: DynamicsSolver = DynamicsSolver()
 
 
 @dispatch
@@ -36,7 +28,7 @@ def evaluate_orbit(
     t: Any,
     /,
     *,
-    integrator: Integrator | None = None,
+    solver: DynamicsSolver | None = None,
     dense: Literal[True, False] = False,
 ) -> Orbit:
     """Compute an orbit in a potential.
@@ -119,14 +111,22 @@ def evaluate_orbit(
     ...                            t=u.Quantity(-100, "Myr"))
     >>> ts = u.Quantity(jnp.linspace(0, 1, 4), "Gyr")
     >>> orbit = gd.evaluate_orbit(potential, w0, ts)
-    >>> orbit
+    >>> print(orbit)
     Orbit(
-      q=CartesianPos3D(...), p=CartesianVel3D(...),
-      t=Quantity['time'](Array(..., dtype=float64), unit='Myr'),
-      frame=SimulationFrame(),
-      potential=KeplerPotential(...),
-      interpolant=None
-    )
+        q=<CartesianPos3D (x[kpc], y[kpc], z[kpc])
+            [[ 8.953 -1.324  0.   ]
+             [ 9.035  1.276  0.   ]
+             [ 2.644 -2.021  0.   ]
+             [ 9.827 -0.562  0.   ]]>,
+        p=<CartesianVel3D (x[kpc / Myr], y[kpc / Myr], z[kpc / Myr])
+            [[ 0.322  0.181  0.   ]
+             [-0.308  0.183  0.   ]
+             [ 1.336 -0.247  0.   ]
+             [ 0.126  0.201  0.   ]]>,
+        t=Quantity['time'](Array([...], dtype=float64), unit='Myr'),
+        frame=SimulationFrame(),
+        potential=KeplerPotential( ... ),
+        interpolant=None)
 
     Note how there are 4 points in the orbit, corresponding to the 4 requested
     return times. These are the times at which the orbit is evaluated, not the
@@ -166,7 +166,7 @@ def evaluate_orbit(
     >>> orbit
     Orbit(
       q=CartesianPos3D(
-        x=Quantity[PhysicalType('length')](value=f64[2,10], unit=Unit("kpc")),
+        x=Quantity[PhysicalType('length')](value=f64[10,2], unit=Unit("kpc")),
         ...
       ),
       p=CartesianVel3D(...),
@@ -199,7 +199,7 @@ def evaluate_orbit(
     """
     # Setup
     units = pot.units
-    integrator = _default_integrator if integrator is None else integrator
+    solver = _default_solver if solver is None else solver
     t = jnp.atleast_1d(FastQ.from_(t, units["time"]))  # ensure t units
 
     field = HamiltonianField(pot)
@@ -208,23 +208,29 @@ def evaluate_orbit(
     tw0 = w0.t if (isinstance(w0, gc.PhaseSpacePosition) and w0.t is not None) else t[0]
 
     # Initial integration `w0.t` to `t[0]`.
-    # TODO: get diffrax's `solver_state` to speed the second integration.
-    # TODO: get diffrax's `controller_state` to speed the second integration.
+    # TODO: diffrax's `solver_state`, `controller_state`
     # TODO: `max_steps` as kwarg.
-    qp0 = integrator(
-        field,
-        w0,
-        tw0,
-        jnp.full_like(tw0, fill_value=t[0]),
-        dense=False,
-    )
+    fullt0 = jnp.full_like(tw0, fill_value=t[0])
+    soln0 = solver.solve(field, w0, tw0, fullt0, dense=False, unbatch_time=True)
 
     # Orbit integration `t[0]` to `t[-1]`
     # TODO: `max_steps` as kwarg.
-    ws = integrator(field, qp0, t[0], t[-1], saveat=t, dense=dense)
+    ys = (FastQ(soln0.ys[0], units["length"]), FastQ(soln0.ys[1], units["speed"]))
+    soln = solver.solve(
+        field, ys, t[0], t[-1], saveat=t, dense=dense, vectorize_interpolation=True
+    )
 
     # Return the orbit object
-    return Orbit._from_psp(ws, t, pot)  # noqa: SLF001
+    return Orbit(
+        q=FastQ(soln.ys[0], units["length"]),
+        p=FastQ(soln.ys[1], units["speed"]),
+        t=t,
+        frame=getattr(w0, "frame", gc.frames.SimulationFrame()),
+        potential=pot,
+        interpolant=(
+            PhaseSpaceInterpolation(soln.interpolation, units=units) if dense else None
+        ),
+    )
 
 
 @dispatch

src/galax/dynamics/_src/orbit/interp.py

@@ -0,0 +1,61 @@
+"""galax: Galactic Dynamix in Jax."""
+
+__all__ = ["PhaseSpaceInterpolation"]
+
+from typing import Any, cast
+from typing_extensions import override
+
+import diffrax as dfx
+import equinox as eqx
+
+import diffraxtra as dfxtra
+import unxt as u
+
+import galax.coordinates as gc
+import galax.typing as gt
+
+
+class PhaseSpaceInterpolation(dfxtra.AbstractVectorizedDenseInterpolation):  # type: ignore[misc]
+    """Phase-space interpolation for orbit evaluation."""
+
+    #: The vectorized interpolation object.
+    interp: dfxtra.VectorizedDenseInterpolation = eqx.field(
+        converter=dfxtra.VectorizedDenseInterpolation.from_
+    )
+
+    #: The unit system for the interpolation.
+    units: u.AbstractUnitSystem = eqx.field(static=True, converter=u.unitsystem)
+
+    @override
+    @eqx.filter_jit  # type: ignore[misc]
+    def evaluate(self, ts: u.Quantity) -> gc.PhaseSpacePosition:
+        # Parse the time
+        t = u.Quantity.from_(ts, self.units["time"])
+
+        # Evaluate the interpolation
+        ys = self.interp.evaluate(t.ustrip(self.units["time"]))
+
+        # Return as a phase-space position
+        return gc.PhaseSpacePosition(
+            q=u.Quantity(ys[0], self.units["length"]),
+            p=u.Quantity(ys[1], self.units["speed"]),
+            t=t,
+        )
+
+    def __call__(self, *args: Any, **kwds: Any) -> gc.PhaseSpacePosition:
+        return cast(gc.PhaseSpacePosition, self.evaluate(*args, **kwds))
+
+    @property
+    def scalar_interpolation(self) -> dfx.DenseInterpolation:
+        """Return the scalar interpolation for the phase-space position."""
+        return cast(dfx.DenseInterpolation, self.interp.scalar_interpolation)
+
+    @property
+    def batch_shape(self) -> gt.Shape:
+        """Return the batch shape of the interpolation."""
+        return cast(gt.Shape, self.interp.batch_shape)
+
+    @property
+    def y0_shape(self) -> gt.Shape:
+        """Return the shape of the initial value."""
+        return cast(gt.Shape, self.interp.y0_shape)