fix backwards-in-time solve for srk's

Author: sammccallum

diffrax/_solver/srk.py

@@ -360,8 +360,15 @@ def step(
         else:
             ignore_stage_g = jnp.array(self.tableau.ignore_stage_g)
 
-        # time increment
+        # The internal time step, h, is always positive. The direction is
+        # handled by _term.py. However, the drift control should be the signed
+        # value of h, `signed_h`.
+        # Futher, since _term.py blanket multiplies the output of the control
+        # by `direction`, we must preserve the symmetry of the SpaceTimeLevyArea,
+        # H, such that H(h) = H(-h).
         h = t1 - t0
+        signed_h = drift.contr(t0, t1)
+        direction = jnp.sign(signed_h)
 
         # First the drift related stuff
         a = self._embed_a_lower(self.tableau.a, dtype)
@@ -402,7 +409,7 @@ def make_zeros_aux(leaf):
         levy_areas = []
         if self.tableau.coeffs_hh is not None:  # space-time Lévy area
             assert isinstance(bm_inc, AbstractSpaceTimeLevyArea)
-            levy_areas.append(bm_inc.H)
+            levy_areas.append((direction * bm_inc.H**ω).ω)
             b_levy_list.append(jnp.asarray(self.tableau.coeffs_hh.b_sol, dtype=dtype))
 
             if self.tableau.coeffs_kk is not None:  # space-time-time Lévy area
@@ -443,7 +450,7 @@ def _comp_g(_t):
 
             if self.tableau.coeffs_hh is not None:  # space-time Lévy area
                 assert isinstance(bm_inc, AbstractSpaceTimeLevyArea)
-                levylist_kgs.append(diffusion.prod(g0, bm_inc.H))
+                levylist_kgs.append(diffusion.prod(g0, (direction * bm_inc.H**ω).ω))
                 a_levy.append(jnp.asarray(self.tableau.coeffs_hh.a, dtype=dtype))
 
             if self.tableau.coeffs_kk is not None:  # space-time-time Lévy area
@@ -545,7 +552,7 @@ def stage(
             z_j = (y0**ω + _drift_result**ω + _diffusion_result**ω).ω
 
             def compute_and_insert_kf_j(_h_kfs_in):
-                h_kf_j = drift.vf_prod(t0 + c_j * h, z_j, args, h)
+                h_kf_j = drift.vf_prod(t0 + c_j * h, z_j, args, signed_h)
                 return insert_jth_stage(_h_kfs_in, h_kf_j, j)
 
             if ignore_stage_f is None:
@@ -611,7 +618,7 @@ def compute_and_insert_kg_j(_w_kgs_in, _levylist_kgs_in):
             # g depends on t. This term is of the form $(g1 - g0) * (0.5*W_n - H_n)$.
             if self.tableau.coeffs_hh is not None:  # space-time Lévy area
                 assert isinstance(bm_inc, AbstractSpaceTimeLevyArea)
-                time_var_contr = (bm_inc.W**ω - 2.0 * bm_inc.H**ω).ω
+                time_var_contr = (bm_inc.W**ω - 2.0 * direction * bm_inc.H**ω).ω
                 time_var_term = diffusion.prod(g_delta, time_var_contr)
             else:
                 time_var_term = diffusion.prod(g_delta, bm_inc.W)

test/helpers.py

@@ -160,8 +160,8 @@ def _batch_sde_solve(
     else:
         struct = w_shape
     bm = diffrax.VirtualBrownianTree(
-        t0=t0,
-        t1=t1,
+        t0=jnp.minimum(t0, t1),
+        t1=jnp.maximum(t0, t1),
         shape=struct,
         tol=bm_tol,
         key=key,

test/test_sde1.py

@@ -40,19 +40,19 @@ def _solvers_and_orders():
 # converges to its own limit (i.e. using itself as reference), and then in a
 # different test check whether that limit is the same as the Euler/Heun limit.
 @pytest.mark.parametrize("solver_ctr,noise,theoretical_order", _solvers_and_orders())
-@pytest.mark.parametrize(
-    "dtype",
-    (jnp.float64,),
-)
+@pytest.mark.parametrize("dtype", (jnp.float64,))
+@pytest.mark.parametrize("direction", ("forwards", "backwards"))
 def test_sde_strong_order_new(
-    solver_ctr, noise: Literal["any", "com", "add"], theoretical_order, dtype
+    solver_ctr, noise: Literal["any", "com", "add"], theoretical_order, dtype, direction
 ):
     bmkey = jr.key(5678)
     sde_key = jr.key(11)
     num_samples = 100
     bmkeys = jr.split(bmkey, num=num_samples)
     t0 = 0.3
     t1 = 5.3
+    if direction == "backwards":
+        t0, t1 = t1, t0
 
     if noise == "add":
         sde = get_time_sde(t0, t1, dtype, sde_key, noise_dim=7)
@@ -98,14 +98,20 @@ def get_dt_and_controller(level):
 # This is to avoid recomputing the correct solutions for every solver.
 solutions = {
     "Ito": {
-        "any": None,
-        "com": None,
-        "add": None,
+        ("forwards", "any"): None,
+        ("forwards", "com"): None,
+        ("forwards", "add"): None,
+        ("backwards", "any"): None,
+        ("backwards", "com"): None,
+        ("backwards", "add"): None,
     },
     "Stratonovich": {
-        "any": None,
-        "com": None,
-        "add": None,
+        ("forwards", "any"): None,
+        ("forwards", "com"): None,
+        ("forwards", "add"): None,
+        ("backwards", "any"): None,
+        ("backwards", "com"): None,
+        ("backwards", "add"): None,
     },
 }
 
@@ -115,15 +121,18 @@ def get_dt_and_controller(level):
 # and Heun if the solver is Stratonovich.
 @pytest.mark.parametrize("solver_ctr,noise,theoretical_order", _solvers_and_orders())
 @pytest.mark.parametrize("dtype", (jnp.float64,))
+@pytest.mark.parametrize("direction", ("forwards", "backwards"))
 def test_sde_strong_limit(
-    solver_ctr, noise: Literal["any", "com", "add"], theoretical_order, dtype
+    solver_ctr, noise: Literal["any", "com", "add"], theoretical_order, dtype, direction
 ):
     bmkey = jr.key(5678)
     sde_key = jr.key(11)
     num_samples = 100
     bmkeys = jr.split(bmkey, num=num_samples)
     t0 = 0.3
     t1 = 5.3
+    if direction == "backwards":
+        t0, t1 = t1, t0
 
     if noise == "add":
         sde = get_time_sde(t0, t1, dtype, sde_key, noise_dim=3)
@@ -164,16 +173,16 @@ def test_sde_strong_limit(
     saveat = diffrax.SaveAt(ts=save_ts)
     levy_area = diffrax.SpaceTimeLevyArea  # must be common for all solvers
 
-    if solutions[sol_type][noise] is None:
+    if solutions[sol_type][(direction, noise)] is None:
         correct_sol, _ = simple_batch_sde_solve(
             bmkeys, sde, ref_solver, levy_area, None, contr_fine, 2**-10, saveat
         )
-        solutions[sol_type][noise] = correct_sol
+        solutions[sol_type][(direction, noise)] = correct_sol
     else:
-        correct_sol = solutions[sol_type][noise]
+        correct_sol = solutions[sol_type][(direction, noise)]
 
     sol, _ = simple_batch_sde_solve(
         bmkeys, sde, solver_ctr(), levy_area, None, contr_coarse, 2**-10, saveat
     )
     error = path_l2_dist(correct_sol, sol)
-    assert error < 0.05
+    assert error < 0.1