Merge pull request #95 from ins-amu/jules/implicit-schemes-9047034567365458749

Add implicit Theta scheme and Jacobi solver

Author: maedoc

vbjax/__init__.py

@@ -21,6 +21,7 @@ def _use_many_cores():
 
 # import stuff
 from .loops import make_sde, make_ode, make_dde, make_sdde, heun_step, make_continuation
+from .implicit import jacobi, make_implicit_sde
 from .shtlc import make_shtdiff
 from .neural_mass import (
     JRState, JRTheta, jr_dfun, jr_default_theta,

vbjax/implicit.py

@@ -0,0 +1,143 @@
+import jax
+import jax.numpy as np
+
+def jacobi(A, b, w=2.0/3.0, tol=1e-9, max_iters=100):
+    """
+    Jacobi iterative linear solver for Ax = b.
+
+    Parameters
+    ----------
+    A : array
+        Matrix A.
+    b : array
+        Vector b.
+    w : float
+        Relaxation parameter.
+    tol : float
+        Tolerance.
+    max_iters : int
+        Maximum number of iterations.
+
+    Returns
+    -------
+    x : array
+        Solution vector.
+    n_iter : int
+        Number of iterations performed.
+    """
+
+    # Precompute diagonal inverse and LU
+    diag_A = np.diag(A)
+    inv_diag_A = 1.0 / diag_A
+    w_m_invD = w * np.diag(inv_diag_A)
+    LU = A - np.diag(diag_A)
+
+    x0 = b
+
+    def cond(state):
+        x, dx, n_iter = state
+        return (n_iter < max_iters) & (np.linalg.norm(dx) > tol)
+
+    def body(state):
+        x, dx, n_iter = state
+        xn = w_m_invD @ (b - LU @ x) + (1.0 - w) * x
+        dx = x - xn
+        return xn, dx, n_iter + 1
+
+    dx = np.ones_like(x0)
+    # Using jax.lax.while_loop for the iterative solver
+    x, dx, n_iter = jax.lax.while_loop(cond, body, (x0, dx, 0))
+
+    return x, n_iter
+
+
+def _theta_dy(f, j, h, th, y0, y1, f_y0, pars, tol=1e-4):
+    J = np.eye(y1.size) - h * th * j(y1, pars)
+    F = y0 + h * (th * f(y1, pars) + f_y0) - y1
+
+    dy, _ = jacobi(J, F, w=2.0/3.0, tol=tol, max_iters=10)
+    return dy
+
+def _compute_noise(gfun, x, p, sqrt_dt, z_t):
+    g = gfun(x, p)
+    return g * sqrt_dt * z_t
+
+
+def make_implicit_sde(dt, dfun, jfun, gfun, th=0.5, tol=1e-4, max_iters=10):
+    """
+    Construct an implicit SDE integrator using the Theta method.
+
+    Parameters
+    ----------
+    dt : float
+        Time step.
+    dfun : callable
+        Drift function dfun(x, p).
+    jfun : callable
+        Jacobian of drift function jfun(x, p).
+    gfun : callable or float
+        Diffusion function gfun(x, p) or constant sigma.
+    th : float
+        Theta parameter (0.5 for Crank-Nicolson, 1.0 for backward Euler).
+    tol : float
+        Tolerance for the Newton-Jacobi iteration.
+    max_iters : int
+        Maximum iterations for the Newton loop.
+
+    Returns
+    -------
+    step : callable
+        Step function step(x, z_t, p).
+    loop : callable
+        Loop function loop(x0, zs, p).
+    """
+
+    if not hasattr(gfun, '__call__'):
+        sig = gfun
+        gfun = lambda *_: sig
+
+    sqrt_dt = np.sqrt(dt)
+
+    def step(x, z_t, p):
+        # Euler guess
+        f_val = dfun(x, p)
+        f_y0 = (1 - th) * f_val
+        y1_euler = x + dt * f_val # Initial guess using explicit Euler
+
+        # Refine if implicit
+        def refine(y1):
+             # First Newton step
+             # We use y1 (Euler guess) to compute Jacobian and Residual
+             dy = _theta_dy(dfun, jfun, dt, th, x, y1, f_y0, p, tol=tol)
+             y1 = y1 + dy
+
+             def refinement_cond(state):
+                 y1, dy, n_iter = state
+                 return (n_iter < max_iters) & (np.linalg.norm(dy) > tol)
+
+             def refinement_body(state):
+                 y1, _, n_iter = state
+                 dy = _theta_dy(dfun, jfun, dt, th, x, y1, f_y0, p, tol=tol)
+                 y1 = y1 + dy
+                 return y1, dy, n_iter + 1
+
+             init_state = (y1, dy, 1)
+             final_state = jax.lax.while_loop(refinement_cond, refinement_body, init_state)
+             return final_state[0]
+
+        y1 = jax.lax.cond(th > 0.0, refine, lambda x: x, y1_euler)
+
+        noise = _compute_noise(gfun, x, p, sqrt_dt, z_t)
+        y1 = y1 + noise
+
+        return y1
+
+    @jax.jit
+    def loop(x0, zs, p):
+        def op(x, z):
+            x = step(x, z, p)
+            return x, x
+        _, xs = jax.lax.scan(op, x0, zs)
+        return xs
+
+    return step, loop

vbjax/tests/test_implicit.py

@@ -0,0 +1,135 @@
+
+import jax
+import jax.numpy as np
+import vbjax
+import pytest
+
+def test_jacobi():
+    A = np.array([[4.0, 1.0], [1.0, 3.0]])
+    b = np.array([1.0, 2.0])
+    x, n_iter = vbjax.jacobi(A, b, tol=1e-6)
+    expected = np.linalg.solve(A, b)
+    assert np.allclose(x, expected, atol=1e-5)
+
+def test_make_implicit_sde_linear():
+    # dy = -k * y * dt
+    k = 1.0
+    def f(y, k):
+        return -k * y
+
+    def j_fn(y, k):
+        return -k * np.eye(y.size)
+
+    y0 = np.array([1.0])
+    dt = 0.1
+    tf = 1.0
+
+    # 0.5 = Crank-Nicolson -> trapezoidal rule
+    step, loop = vbjax.make_implicit_sde(dt, f, j_fn, 0.0, th=0.5)
+
+    # Create noise (zeros since deterministic test)
+    n_steps = int(tf / dt)
+    zs = np.zeros((n_steps, 1))
+
+    ys = loop(y0, zs, k)
+
+    ts = np.arange(1, n_steps + 1) * dt
+    exact = y0 * np.exp(-ts)
+
+    factor = (1 - k*dt/2) / (1 + k*dt/2)
+    expected_numeric = y0 * (factor ** np.arange(1, n_steps + 1))
+
+    error = np.abs(ys.flatten() - expected_numeric)
+    assert np.max(error) < 1e-5
+
+    # Also check against exact just to be sure it's close
+    error_exact = np.abs(ys.flatten() - exact)
+    assert np.max(error_exact) < 1e-2
+
+
+def test_make_implicit_sde_autodiff():
+    # dy = -y^3
+    def f(y, p):
+        return -y**3
+
+    # Auto-diff Jacobian
+    j_fn = jax.jacfwd(f)
+
+    y0 = np.array([1.0])
+    dt = 0.1
+    tf = 1.0
+
+    step, loop = vbjax.make_implicit_sde(dt, f, j_fn, 0.0, th=1.0) # Backward Euler
+
+    n_steps = int(tf / dt)
+    zs = np.zeros((n_steps, 1))
+
+    ys = loop(y0, zs, None)
+
+    # Backward Euler: y_{n+1} = y_n - h * y_{n+1}^3
+    # Check consistency
+    for i in range(n_steps):
+        prev = y0 if i == 0 else ys[i-1]
+        curr = ys[i]
+        res = curr + dt * curr**3 - prev
+        assert np.abs(res) < 1e-4
+
+@pytest.mark.benchmark(group="stiff_solver")
+def test_benchmark_heun_stiff(benchmark):
+    k = 1000.0
+    dt = 0.0001 # Stability limit requires small dt
+    tf = 1.0
+    n_steps = int(tf / dt)
+
+    def f(y, k):
+        return -k * y
+
+    def g(y, k):
+        return 0.1
+
+    y0 = np.ones(100)
+    zs = jax.random.normal(jax.random.PRNGKey(0), (n_steps, 100))
+
+    _, loop = vbjax.make_sde(dt, f, g)
+
+    # Warmup / Compile
+    loop(y0, zs, k).block_until_ready()
+
+    def run():
+        return loop(y0, zs, k).block_until_ready()
+
+    benchmark(run)
+
+@pytest.mark.benchmark(group="stiff_solver")
+def test_benchmark_implicit_stiff(benchmark):
+    k = 1000.0
+    dt = 0.01 # Implicit can take larger steps
+    tf = 1.0
+    n_steps = int(tf / dt)
+
+    def f(y, k):
+        return -k * y
+
+    def j_fn(y, k):
+        return -k * np.eye(y.size)
+
+    def g(y, k):
+        return 0.1
+
+    y0 = np.ones(100)
+    zs = jax.random.normal(jax.random.PRNGKey(0), (n_steps, 100))
+
+    step, loop = vbjax.make_implicit_sde(dt, f, j_fn, g, th=0.5)
+
+    # Warmup / Compile
+    loop(y0, zs, k).block_until_ready()
+
+    def run():
+        return loop(y0, zs, k).block_until_ready()
+
+    benchmark(run)
+
+if __name__ == "__main__":
+    test_jacobi()
+    test_make_implicit_sde_linear()
+    test_make_implicit_sde_autodiff()