Add implicit Theta scheme and Jacobi solver to vbjax

- Added `vbjax/implicit.py` implementing `make_implicit_sde` (Theta method) and `jacobi` (linear solver).
- Exposed new functions in `vbjax/__init__.py`.
- Added tests in `vbjax/tests/test_implicit.py` covering solver accuracy and integrator convergence.
- Added `benchmark_implicit.py` to demonstrate speedup on stiff systems compared to explicit Heun integration.

Author: google-labs-jules[bot]

benchmark_implicit.py

@@ -0,0 +1,63 @@
+
+import jax
+import jax.numpy as np
+import vbjax
+import time
+
+def benchmark():
+    # Stiff system example: Robertson problem or simple stiff linear
+    # Let's use a stiff linear system: dy/dt = -1000 y
+    k = 1000.0
+    dt = 0.001 # 1/k.  Explicit stability limit is 2/k = 0.002.
+    # If we use dt = 0.01, explicit should explode, implicit should work.
+
+    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) # 100 dimensions
+    tf = 1.0
+
+    # 1. Heun (Explicit)
+    # We use a small dt for stability or just measure speed
+    dt_explicit = 0.0001
+    n_explicit = int(tf / dt_explicit)
+    zs_explicit = jax.random.normal(jax.random.PRNGKey(0), (n_explicit, 100))
+
+    _, loop_heun = vbjax.make_sde(dt_explicit, f, g)
+
+    # Compile
+    loop_heun(y0, zs_explicit, k).block_until_ready()
+
+    t0 = time.time()
+    for _ in range(10):
+        loop_heun(y0, zs_explicit, k).block_until_ready()
+    t_heun = (time.time() - t0) / 10
+    print(f"Heun (dt={dt_explicit}, steps={n_explicit}): {t_heun*1000:.2f} ms")
+
+    # 2. Implicit (Theta=0.5)
+    # Use larger dt
+    dt_implicit = 0.01
+    n_implicit = int(tf / dt_implicit)
+    zs_implicit = jax.random.normal(jax.random.PRNGKey(0), (n_implicit, 100))
+
+    step_imp, loop_imp = vbjax.make_implicit_sde(dt_implicit, f, j_fn, g, th=0.5)
+
+    # Compile
+    loop_imp(y0, zs_implicit, k).block_until_ready()
+
+    t0 = time.time()
+    for _ in range(10):
+        loop_imp(y0, zs_implicit, k).block_until_ready()
+    t_imp = (time.time() - t0) / 10
+    print(f"Implicit (dt={dt_implicit}, steps={n_implicit}): {t_imp*1000:.2f} ms")
+
+    print(f"Speedup factor (Total Time): {t_heun / t_imp:.2f}x")
+
+if __name__ == "__main__":
+    benchmark()

vbjax/__init__.py

@@ -21,7 +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 theta, jacobi
+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

@@ -52,94 +52,92 @@ def body(state):
 
 
 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
+    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 theta(f, j, y0, h, tf, *pars, th=0.5, sigma=0.0, tol=1e-4, key=None):
+
+def make_implicit_sde(dt, dfun, jfun, gfun, th=0.5, tol=1e-4, max_iters=10):
     """
-    Implicit Theta integration scheme.
+    Construct an implicit SDE integrator using the Theta method.
 
     Parameters
     ----------
-    f : callable
-        Drift function f(y, *pars).
-    j : callable
-        Jacobian function j(y, *pars).
-    y0 : array
-        Initial state.
-    h : float
+    dt : float
         Time step.
-    tf : float
-        Final time.
-    *pars : list
-        Parameters passed to f and j.
+    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 trapezoidal/Crank-Nicolson, 1.0 for backward Euler).
-    sigma : float or array
-        Noise standard deviation.
+        Theta parameter (0.5 for Crank-Nicolson, 1.0 for backward Euler).
     tol : float
         Tolerance for the Newton-Jacobi iteration.
-    key : jax.random.PRNGKey, optional
-        Random key for noise generation. If None, a default key is used.
+    max_iters : int
+        Maximum iterations for the Newton loop.
 
     Returns
     -------
-    ts : array
-        Time points.
-    ys : array
-        Trajectory.
+    step : callable
+        Step function step(x, z_t, p).
+    loop : callable
+        Loop function loop(x0, zs, p).
     """
-    if key is None:
-        key = jax.random.PRNGKey(42)
 
-    num_steps = int(tf / h)
-    add_noise = np.any(sigma > 0.0)
+    if not hasattr(gfun, '__call__'):
+        sig = gfun
+        gfun = lambda *_: sig
 
-    def scan_body(carry, _):
-        y0, rng_key = carry
+    sqrt_dt = np.sqrt(dt)
 
-        f_y0 = (1 - th) * f(y0, *pars)
-        y1 = y0 + h * f_y0
+    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):
-             dy = _theta_dy(f, j, h, th, y0, y0, f_y0, pars, tol=tol)
-             y1 = y0 + dy
+             # 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 < 10) & (np.linalg.norm(dy) > tol)
+                 return (n_iter < max_iters) & (np.linalg.norm(dy) > tol)
 
              def refinement_body(state):
                  y1, _, n_iter = state
-                 dy = _theta_dy(f, j, h, th, y0, y1, f_y0, pars, tol=tol)
+                 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)
+        y1 = jax.lax.cond(th > 0.0, refine, lambda x: x, y1_euler)
 
-        rng_key, step_key = jax.random.split(rng_key)
-        noise = jax.lax.cond(
-            add_noise,
-            lambda k: np.sqrt(sigma) * jax.random.normal(k, y0.shape),
-            lambda k: np.zeros_like(y0),
-            step_key
-        )
+        noise = _compute_noise(gfun, x, p, sqrt_dt, z_t)
         y1 = y1 + noise
 
-        return (y1, rng_key), y1
-
-    _, ys = jax.lax.scan(scan_body, (y0, key), None, length=num_steps)
+        return y1
 
-    # Prepend y0
-    ys = np.concatenate([y0[None, ...], ys], axis=0)
-    ts = np.arange(num_steps + 1) * h
+    @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 ts, ys
+    return step, loop

vbjax/tests/test_implicit.py

@@ -5,74 +5,80 @@
 import pytest
 
 def test_jacobi():
-    # Solve Ax = b
-    # A = [[4, 1], [1, 3]]
-    # b = [1, 2]
-    # Exact solution:
-    # 4x + y = 1
-    # x + 3y = 2
-    # x = 2 - 3y
-    # 4(2-3y) + y = 1 => 8 - 12y + y = 1 => -11y = -7 => y = 7/11
-    # x = 2 - 21/11 = 1/11
-
     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)
-    print(f"Jacobi solved in {n_iter} iterations. x={x}")
-
-
-def test_theta_linear():
-    # Test implicit scheme on a simple linear ODE: y' = -y
-    # Analytical solution: y(t) = y0 * exp(-t)
 
+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):
-        # Jacobian is scalar -k
-        # If y is vector, Jacobian is diagonal -k * I
         return -k * np.eye(y.size)
 
     y0 = np.array([1.0])
-    h = 0.1
+    dt = 0.1
     tf = 1.0
 
-    # Run with Theta=0.5 (Crank-Nicolson)
-    ts, ys = vbjax.theta(f, j_fn, y0, h, tf, k, th=0.5)
+    # 0.5 = Crank-Nicolson -> trapezoidal rule
+    step, loop = vbjax.make_implicit_sde(dt, f, j_fn, 0.0, th=0.5)
 
-    exact = y0 * np.exp(-ts)
-    error = np.abs(ys.flatten() - exact)
+    # Create noise (zeros since deterministic test)
+    n_steps = int(tf / dt)
+    zs = np.zeros((n_steps, 1))
 
-    print(f"Max error (Theta=0.5): {np.max(error)}")
-    assert np.max(error) < 1e-2
+    ys = loop(y0, zs, k)
 
+    ts = np.arange(1, n_steps + 1) * dt
+    exact = y0 * np.exp(-ts) # Actually Trapezoidal rule approximation will differ slightly from exact exp
+    # Trapezoidal: y_{n+1} = y_n + h/2 (f_n + f_{n+1}) = y_n - k*h/2 (y_n + y_{n+1})
+    # y_{n+1} (1 + kh/2) = y_n (1 - kh/2)
+    # y_{n+1} = y_n * (1 - kh/2) / (1 + kh/2)
 
-def test_theta_auto_jacobian():
-    # Same as above but using jax.jacfwd
-    k = 1.0
-    def f(y, k):
-        return -k * y
+    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
 
-    j_fn = jax.jacfwd(f, argnums=0)
+
+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])
-    h = 0.1
+    dt = 0.1
     tf = 1.0
 
-    ts, ys = vbjax.theta(f, j_fn, y0, h, tf, k, th=0.5)
+    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))
 
-    exact = y0 * np.exp(-ts)
-    error = np.abs(ys.flatten() - exact)
+    ys = loop(y0, zs, None)
 
-    print(f"Max error (Theta=0.5, AutoDiff): {np.max(error)}")
-    assert np.max(error) < 1e-2
+    # Backward Euler: y_{n+1} = y_n - h * y_{n+1}^3
+    # Solve y_{n+1} + h * y_{n+1}^3 - y_n = 0
+    # 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
 
 if __name__ == "__main__":
     test_jacobi()
-    test_theta_linear()
-    test_theta_auto_jacobian()
+    test_make_implicit_sde_linear()
+    test_make_implicit_sde_autodiff()