feat: return PSP, not full orbit (GalacticDynamics#133)

* feat: return PSP, not full orbit
* fix: tests

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

docs/getting_started.rst

@@ -67,9 +67,9 @@ position near the Sun::
 
     >>> xyz = [-8., 0, 0] * u.kpc
     >>> mw.potential_energy(xyz, t=0)
-    Quantity(value=f64[], unit=Unit("kpc2 / Myr2"))
+    Quantity[PhysicalType('specific energy')](value=f64[], unit=Unit("kpc2 / Myr2"))
     >>> mw.acceleration(xyz, t=0)
-    Quantity(value=f64[3], unit=Unit("kpc / Myr2"))
+    Quantity[PhysicalType('acceleration')](value=f64[3], unit=Unit("kpc / Myr2"))
 
 The values that are returned by most methods in :mod:`galax` are provided as
 Astropy :class:`~astropy.units.Quantity` objects, which represent numerical data
@@ -78,9 +78,9 @@ re-represented in any equivalent units, so, for example, we could display the
 energy or acceleration in other units::
 
     >>> mw.potential_energy(xyz, t=0)
-    Quantity(value=f64[], unit=Unit("kpc2 / Myr2"))
+    Quantity[PhysicalType('specific energy')](value=f64[], unit=Unit("kpc2 / Myr2"))
     >>> mw.acceleration(xyz, t=0)
-    Quantity(value=f64[3], unit=Unit("kpc / Myr2"))
+    Quantity[PhysicalType('acceleration')](value=f64[3], unit=Unit("kpc / Myr2"))
 
 Now that we have a potential model, if we want to compute an orbit, we need to
 specify a set of initial conditions to initialize the numerical orbit

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

@@ -15,7 +15,7 @@
 from galax.coordinates import PhaseSpacePosition, PhaseSpaceTimePosition
 from galax.dynamics._dynamics.integrate._api import Integrator
 from galax.dynamics._dynamics.integrate._builtin import DiffraxIntegrator
-from galax.dynamics._dynamics.orbit import Orbit, evaluate_orbit, integrate_orbit
+from galax.dynamics._dynamics.orbit import evaluate_orbit, integrate_orbit
 from galax.potential._potential.base import AbstractPotentialBase
 from galax.typing import BatchVec6, FloatScalar, IntScalar, Vec6, VecN, VecTime
 
@@ -132,7 +132,7 @@ def run(
         *,
         seed_num: int,
         vmapped: bool | None = None,
-    ) -> tuple[MockStream, Orbit]:
+    ) -> tuple[MockStream, PhaseSpaceTimePosition]:
         """Generate mock stellar stream.
 
         Parameters
@@ -158,10 +158,10 @@ def run(
 
         Returns
         -------
-        mockstream : MockStream
+        mockstream : :class:`galax.dynamcis.MockStream`
             Leading and/or trailing arms of the mock stream.
-        prog_o : Orbit
-            Orbit of the progenitor.
+        prog_o : :class:`galax.coordinates.PhaseSpaceTimePosition`
+            The final phase-space(+time) position of the progenitor.
         """
         # TODO: ꜛ a discussion about the stripping times
         # Parse vmapped
@@ -186,8 +186,8 @@ def run(
             self.potential, w0, ts, integrator=self.progenitor_integrator
         )
 
-        # Generate stream initial conditions along the integrated progenitor
-        # orbit. The release times are stripping times.
+        # Generate initial conditions from the DF, along the integrated
+        # progenitor orbit. The release times are the stripping times.
         mock0_lead, mock0_trail = self.df.sample(
             self.potential, prog_o, prog_mass, seed_num=seed_num
         )
@@ -222,4 +222,4 @@ def run(
 
         mockstream = MockStream(q=q, p=p, t=t, release_time=release_time)
 
-        return mockstream, prog_o
+        return mockstream, prog_o[-1]

tests/unit/dynamics/mockstream/test_mockstreamgenerator.py

@@ -94,7 +94,7 @@ def test_run_scan(
 
         # TODO: more rigorous tests
         assert mock.q.shape == (2 * len(t_stripping), 3)
-        assert prog_o.q.shape == (len(t_stripping), 3)
+        assert prog_o.q.shape == (3,)
 
         # Test that the positions and momenta are finite
         assert jnp.isfinite(mock.q).all()

tests/unit/potential/test_base.py

@@ -3,6 +3,7 @@
 from typing import Any
 
 import array_api_jax_compat as xp
+import astropy.units as u
 import equinox as eqx
 import jax
 import jax.numpy as jnp
@@ -111,7 +112,7 @@ def _potential_energy(self, q: Vec3, /, t: FloatOrIntScalar) -> FloatScalar:
 
     def test_potential_energy(self, pot: AbstractPotentialBase, x: Vec3) -> None:
         """Test the `AbstractPotentialBase.potential_energy` method."""
-        assert pot.potential_energy(x, t=0) == 6
+        assert pot.potential_energy(x, t=0) == Quantity(6, u.kpc**2 / u.Myr**2)
 
     # ---------------------------------