Vector of Wiener processes

[julianiacoponi]

I want to try and use diffrax to simulate some correlated ("directional") noise in a 3D oscillator i.e. solving the following SDEs

I couldn’t parse from the docs what kind of solver would treat the noise term as a matrix and the dW term as a vector of Wiener processes, at the same time.

Does diffrax have such a solver / can I build one from what diffrax offers?

[hsimonfroy]

Hi!

You could try something like:

drift = lambda t, y, args: F  # F is a jnp array of shape (6)
diffusion = lambda t, y, args: G # G is a jnp array of shape (6,4)
brownian_motion = VirtualBrownianTree(t0, t1, tol=dt/2, shape=(4), key=seed)
terms = MultiTerm(ODETerm(drift), ControlTerm(diffusion, brownian_motion))

It wasn't clear to me, but for a drift function $f$ and a diffusion function $g$, what diffrax actually implements is: $dX = f \odot dt + g \cdot dW$, where $\odot$ is piecewise product and $\cdot$ is a tensor dot product (sum over all the dimensions of dW).

[patrick-kidger]

Yup! diffrax.ControlTerm is the way to implement a matrix-valued diffusion with vector-valued noise.

If this is unclear then I'd definitely be happy to consider a PR improving the documentation. :)