SDE tests failing with complex inputs

[Randl]

There are some tests for SDEs failing when adding complex inputs

 =========================== short test summary info ============================
FAILED test/test_integrate.py::test_sde_strong_order[complex128-ItoMilstein-any-0.5] - assert -0.2 < (Array(-0.00617573+0.01165244j, dtype=complex128) - 0.5)
FAILED test/test_integrate.py::test_sde_strong_order[complex128-SPaRK-any-0.5] - assert (Array(0.8286123-0.19839093j, dtype=complex128) - 0.5) < 0.2
FAILED test/test_integrate.py::test_sde_strong_order[complex128-GeneralShARK-any-0.5] - assert (Array(0.79761363-0.21382705j, dtype=complex128) - 0.5) < 0.2
FAILED test/test_integrate.py::test_sde_strong_order[complex128-StratonovichMilstein-com-1] - assert (Array(1.21753621+0.37906851j, dtype=complex128) - 1) < 0.2
FAILED test/test_integrate.py::test_sde_strong_order[complex128-SPaRK-com-1] - assert (Array(1.21429825+0.13029023j, dtype=complex128) - 1) < 0.2

Note that not all instances of this test are failing, so the problem is likely in the specific solvers.
I'd love to look into it, but I'm not an expert on stochastic SDEs, and some guidance would be appreciated. Specifically, are there any "usual suspects" regarding what could go wrong? Is there a toy example worth starting from?
Also, can the order be complex or is it itself a problem? If it is, should we take the abs or the real part for comparison?

[patrick-kidger]

The order should definitely be real. I suspect the answer isn't "just take the absolute part", though.

FWIW the "strong order" of an SDE method is basically just the exact same thing as the "order" of an ODE method. (The "strong" is mreaning the order is measuring with a fixed choice of Brownian input, rather than any kind of probabilistic thing going on.)

[andyElking]

Hi! I think you just need to add a jnp.abs on the line

diffrax/test/helpers.py

Line 103 in 80d1f34

square_diff = jnp.square(y1 - y2)

If that doesn't work, I can take a closer look.

[Randl]

That fixed part of the tests, but one is still failing:

=========================== short test summary info ============================
FAILED test/test_integrate.py::test_sde_strong_order[complex128-ItoMilstein-any-0.5] - assert -0.2 < (Array(0.02747656, dtype=float64) - 0.5)

[andyElking]

Great to see the others passed!!
As far as I know, Milstein is the only SDE solver in Diffrax which uses derivatives, so that could be the cause. But otherwise I have no clue.

[Randl]

Hmm, I fixed Milstein, but it looks like there are additional failures for SRK, specifically with additive noise. Most likely some conj() missing somewhere