Coupled SDE System Implementation

Hi all,

thanks for the great library. I'm having an issue implementing a coupled system of SDEs. I'm getting an ValueError: `terms` must be a PyTree of `AbstractTerms` (such as `ODETerm`) error. The system is:

\begin{aligned}
		\frac{\mathrm{d} S(t)}{\mathrm{d} t} &= -\beta(t)S(t)\frac{I(t)}{N} \mathrm{d} t, \\
		\frac{\mathrm{d} I(t)}{\mathrm{d}t} &= (\beta(t)S(t)\frac{I(t)}{N} - \gamma(t) I(t)) \mathrm{d} t,\\
		\frac{\mathrm{d} R(t)}{\mathrm{d}t} &= \gamma(t) I(t)\, \mathrm{d} t,\\
		\frac{\mathrm{d} \log\beta(t)}{\mathrm{d}t} &= w_3\mathrm{d} B_w(t),\\
        \frac{\mathrm{d} \log\gamma(t)}{\mathrm{d}t} &= u_3 \mathrm{d}B_u(t)
\end{aligned}

The code is

import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt

import diffrax

def sde_drift(t, y, args):
    N, _ = args
    beta = jnp.exp(y[3])
    gamma = jnp.exp(y[4])
    dS = -(beta * y[0] * y[1]) / N
    dI = (beta * y[0] * y[1]) / N - y[1] * gamma
    dR = y[1] * gamma
    # only diffusion, no drift
    dbeta = 0.0  # jnp.array([0.0])
    dgamma = 0.0  # jnp.array([0.0])
    dy = jnp.array([dS, dI, dR, dbeta, dgamma])

    return dy

def sde_diffusion(t, y, args):
    _, sigma_1 = args
    y1, y2, y3, y4, y5 = y
    diagonal = jnp.array([0.0, 0.0, 0.0, sigma_1 * y4, sigma_1 * y5])
    return diagonal 


def sde():

    t0 = 0
    t1 = 100
    dt0 = 0.1
    y0 = jnp.array([3990.0, 10.0, 0.01, jnp.log(0.25), jnp.log(0.05)])
    args = (4000.0, 0.2)

    bm = diffrax.VirtualBrownianTree(t0, t1, tol=1e-2, shape=(5,), key=jr.PRNGKey(42))
    terms = diffrax.MultiTerm(diffrax.ODETerm(sde_drift), diffrax.ControlTerm(sde_diffusion, bm))
    solver = diffrax.SEA()
    saveat = diffrax.SaveAt(dense=True)

    print(type(terms))

    sol = diffrax.diffeqsolve(terms, solver, t0, t1, dt0=dt0, y0=y0, args=args, saveat=saveat)
    print(sol)

Printing the type of terms yields 'diffrax._term.MultiTerm, so I'm not entirely sure where to look. What can you suggest to look at?

Thanks in advance.

Reproduce the supplied example and begin at diffrax.diffeqsolve, MultiTerm, and the AbstractTerms validation that emits the ValueError. Inspect how the SEA solver expects the coupled drift, diffusion, and VirtualBrownianTree terms to be shaped; done means the example no longer raises the validation error and its behavior is verified.