IID Brownian motion?

[jajcayn]

Hey there, First off, thanks a lot for this library! I wondered whether it is possible to solve a 2-dimensional SDE with IID noise? As a simple example, consider the Ornstein-Uhlenback process (I need it to drive other ODE, so I want to use sol_ou.evaluate(t) in the later vector field).

t0, t1 = 0.0, 20.0
DT = 0.001
dim = 2

def drift_ou(t, y, args):
    return (args["mu"] - y) / args["tau"]

def diffustion_ou(t, y, args):
    return jnp.ones((dim))*args["sigma"]

brownian_motion = dfx.VirtualBrownianTree(
    t0,
    t1,
    tol=1e-3,
    shape=jax.ShapeDtypeStruct((dim,), np.float64),
    key=jr.PRNGKey(1),
)

terms = dfx.MultiTerm(dfx.ODETerm(drift_ou), dfx.ControlTerm(diffustion_ou, brownian_motion))
args = {"mu": 1.0, "tau": 0.1, "sigma": 0.3}
sol_ou = dfx.diffeqsolve(
    terms,
    dfx.Heun(),
    t0,
    t1,
    dt0=DT,
    y0=jnp.ones((dim))*args["mu"],
    saveat=dfx.SaveAt(dense=True),
    max_steps=int(t1*(1./DT)),
    args=args
)

In this code, I am trying to create a 2-dimensional Ornstein-Uhlenback, and VirtualBrownianTree gives me the dW in 2 dimensions; however, the Brownian motion is the same for both dimensions. Is there some way to make the "samples" independent?

Cheers!