13_03_2025: added symplectic tao integrator for non-separatable hamiltonians

Author: i-a-morozov

sympint/__init__.py

@@ -13,7 +13,8 @@
     'coefficients',
     'table',
     'sequence',
-    'midpoint'
+    'midpoint',
+    'tao'
 ]
 
 from sympint.functional import nest
@@ -26,4 +27,5 @@
 from sympint.yoshida import table
 from sympint.yoshida import sequence
 
-from sympint.integrators import midpoint
+from sympint.integrators import midpoint
+from sympint.integrators import tao

sympint/integrators.py

@@ -11,6 +11,9 @@
 from jax import Array
 
 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]:
     """
@@ -47,4 +50,50 @@ def newton(state: Array) -> Array:
             delta, *_ = jax.numpy.linalg.lstsq(jacobian, -error)
             return state + delta
         return nest(ns, newton)(state)
-    return integrator
+    return integrator
+
+
+def tao(H:Callable[..., Array], binding:float=0.0) -> Callable[..., Array]:
+    """
+    Generate Tao integrator
+
+    Parameters
+    ----------
+    H: Callable[[Array, *Any], Array]
+        Hamiltonian function H(q, p, dt, *args)
+    binding: float, default=0.0
+        binding factor
+
+    Returns
+    -------
+    Callable[[Array, *Any], Array]
+        integrator(qp, dt, *args)
+
+    """
+    dHdq = jax.grad(H, argnums=0)
+    dHdp = jax.grad(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])
+    def fb(state:Array, dt:Array, *args:Array) -> Array:
+        q, p, Q, P = state.reshape(4, -1)
+        return jax.numpy.concatenate([q + dt*dHdp(Q, p, *args), p, Q, P - dt*dHdq(Q, p, *args)])
+    def fc(state:Array, dt:Array, *args:Array) -> Array:
+        q, p, Q, P = state.reshape(4, -1)
+        omega = 2*binding*dt
+        cos = jax.numpy.cos(omega)
+        sin = jax.numpy.sin(omega)
+        dq = q - Q
+        dp = p - P
+        return jax.numpy.concatenate([
+            0.5*(q + Q + cos*dq + sin*dp),
+            0.5*(p + P - sin*dq + cos*dp),
+            0.5*(q + q - cos*dq - sin*dp),
+            0.5*(p + P + sin*dq - cos*dp)
+        ])
+    step = fold(sequence(0, 0, [fa, fb, fc] if binding != 0.0 else [fa, fb], merge=True))
+    def integrator(state:Array, dt:Array, *args:Array) -> Array:
+        local = step(jax.numpy.concatenate([state, state]), dt, *args)
+        q, p, *_ = local.reshape(4, -1)
+        return jax.numpy.concatenate([q, p])
+    return integrator