fix mockstream shape (#88)

Signed-off-by: nstarman <nstarman@users.noreply.github.com>
GalacticDynamics · Jan 25, 2024 · 085228d · 085228d
1 parent 9382775
commit 085228d
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/galax/dynamics/mockstream/_df/base.py b/src/galax/dynamics/mockstream/_df/base.py
@@ -83,10 +83,10 @@ def scan_fn(carry: Carry, t: FloatScalar) -> tuple[Carry, Wif]:
             xp.asarray([0.0, 0.0, 0.0]),
             xp.asarray([0.0, 0.0, 0.0]),
         )
-        x_lead, x_trail, v_lead, v_trail = jax.lax.scan(scan_fn, init_carry, ts[1:])[1]
+        x_lead, x_trail, v_lead, v_trail = jax.lax.scan(scan_fn, init_carry, ts)[1]
 
-        mock_lead = MockStream(x_lead, v_lead, ts[1:])
-        mock_trail = MockStream(x_trail, v_trail, ts[1:])
+        mock_lead = MockStream(x_lead, v_lead, ts)
+        mock_trail = MockStream(x_trail, v_trail, ts)
 
         return mock_lead, mock_trail