fix _integrator_F as FCallable

Signed-off-by: nstarman <nstarman@users.noreply.github.com>
GalacticDynamics · Jan 12, 2024 · 187c80a · 187c80a
1 parent 8565d80
commit 187c80a
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/galax/potential/_potential/base.py b/src/galax/potential/_potential/base.py
@@ -290,8 +290,8 @@ def tidal_tensor(
     # Integrating orbits
 
     @partial_jit()
-    def _integrator_F(self, t: FloatScalar, xv: Vec6, args: tuple[Any, ...]) -> Vec6:
-        return xp.hstack([xv[3:6], self.acceleration(xv[0:3], t)])  # v, a
+    def _integrator_F(self, t: FloatScalar, qp: Vec6, args: tuple[Any, ...]) -> Vec6:
+        return xp.hstack([qp[3:6], self.acceleration(qp[0:3], t)])  # v, a
 
     @partial_jit(static_argnames=("integrator"))
     def integrate_orbit(
@@ -307,5 +307,5 @@ def integrate_orbit(
 
         integrator_ = default_integrator if integrator is None else replace(integrator)
 
-        ws = integrator_.run(self._integrator_F, qp0, t0, t1, ts)  # type: ignore[arg-type]
+        ws = integrator_.run(self._integrator_F, qp0, t0, t1, ts)
         return Orbit(q=ws[:, :3], p=ws[:, 3:-1], t=ws[:, -1], potential=self)