07_17_2025: added option to pass gradient (default is jax.grad)

Author: i-a-morozov

sympint/integrators.py

@@ -6,16 +6,20 @@
 
 """
 from typing import Callable
+from typing import Optional
 
 import jax
 from jax import Array
+from jax import grad
 
 from sympint.functional import nest
 from sympint.functional import fold
 
 from sympint.yoshida import sequence
 
-def midpoint(H:Callable[..., Array], ns:int=1) -> Callable[..., Array]:
+def midpoint(H:Callable[..., Array],
+            ns:int=1,
+            gradient:Optional[Callable[..., Array]] = None) -> Callable[..., Array]:
     """
     Generate implicit midpoint integrator
 
@@ -25,15 +29,18 @@ def midpoint(H:Callable[..., Array], ns:int=1) -> Callable[..., Array]:
         Hamiltonian function H(q, p, dt, t, *args)
     ns: int, default=1
         number of Newton iteration steps
+    gradient: Optional[Callable[..., Array]], default=None
+        gradient function (defaults to jax.grad)
 
     Returns
     -------
     Callable[[Array, *Any], Array]
         integrator(qp, dt, t, *args)
 
     """
-    dHdq = jax.grad(H, argnums=0)
-    dHdp = jax.grad(H, argnums=1)
+    gradient = grad if gradient is None else gradient
+    dHdq = gradient(H, argnums=0)
+    dHdp = gradient(H, argnums=1)
     def integrator(state: Array, dt: Array, t: Array, *args: Array) -> Array:
         q, p = jax.numpy.reshape(state, (2, -1))
         t_m = t + 0.5*dt
@@ -53,7 +60,9 @@ def newton(state: Array) -> Array:
     return integrator
 
 
-def tao(H:Callable[..., Array], binding:float=0.0) -> Callable[..., Array]:
+def tao(H:Callable[..., Array],
+        binding:float=0.0,
+        gradient:Optional[Callable[..., Array]] = None) -> Callable[..., Array]:
     """
     Generate Tao integrator
 
@@ -63,15 +72,18 @@ def tao(H:Callable[..., Array], binding:float=0.0) -> Callable[..., Array]:
         Hamiltonian function H(q, p, dt, *args)
     binding: float, default=0.0
         binding factor
+    gradient: Optional[Callable[..., Array]], default=None
+        gradient function (defaults to jax.grad)
 
     Returns
     -------
     Callable[[Array, *Any], Array]
         integrator(qp, dt, *args)
 
     """
-    dHdq = jax.grad(H, argnums=0)
-    dHdp = jax.grad(H, argnums=1)
+    gradient = grad if gradient is None else gradient
+    dHdq = gradient(H, argnums=0)
+    dHdp = gradient(H, argnums=1)
     def fa(state:Array, dt:Array, *args:Array) -> Array:
         q, p, Q, P = state.reshape(4, -1)
         return jax.numpy.concatenate([q, p - dt*dHdq(q, P, *args), Q + dt*dHdp(q, P, *args), P])

sympint/yoshida.py

@@ -142,7 +142,7 @@ def sequence(ni:int,
     Output sequence mappings have (x, dt, *args) singnatures
 
     """
-    indices, weights = table(len(mappings), ni, nf, merge)
+    indices, weights, *_ = table(len(mappings), ni, nf, merge)
     parameters = [[] for _ in range(len(mappings))] if parameters is None else parameters
     parameters = [parameters[i] for i in indices]
     def wrapper(mapping, weight, parameter):