add direction for saveat + test (#427)

* add direction for saveat + test

* remove direction, make changes before

* fix `save_y` test

* expand tests

* fix inf initialization

* fix test

Author: lockwo

diffrax/_integrate.py

@@ -37,7 +37,7 @@
     NoProgressMeter,
 )
 from ._root_finder import use_stepsize_tol
-from ._saveat import SaveAt, SubSaveAt
+from ._saveat import save_y, SaveAt, SubSaveAt
 from ._solution import is_okay, is_successful, RESULTS, Solution
 from ._solver import (
     AbstractImplicitSolver,
@@ -881,6 +881,18 @@ def _check_subsaveat_ts(ts):
 
     saveat = eqx.tree_at(_get_subsaveat_ts, saveat, replace_fn=_check_subsaveat_ts)
 
+    def _subsaveat_direction_fn(x):
+        if _is_subsaveat(x):
+            if x.fn is not save_y:
+                direction_fn = lambda t, y, args: x.fn(direction * t, y, args)
+                return eqx.tree_at(lambda x: x.fn, x, direction_fn)
+            else:
+                return x
+        else:
+            return x
+
+    saveat = jtu.tree_map(_subsaveat_direction_fn, saveat, is_leaf=_is_subsaveat)
+
     # Initialise states
     tprev = t0
     error_order = solver.error_order(terms)
@@ -924,7 +936,7 @@ def _allocate_output(subsaveat: SubSaveAt) -> SaveState:
             out_size += 1
         saveat_ts_index = 0
         save_index = 0
-        ts = jnp.full(out_size, jnp.inf, dtype=time_dtype)
+        ts = jnp.full(out_size, direction * jnp.inf, dtype=time_dtype)
         struct = eqx.filter_eval_shape(subsaveat.fn, t0, y0, args)
         ys = jtu.tree_map(
             lambda y: jnp.full((out_size,) + y.shape, jnp.inf, dtype=y.dtype), struct
@@ -1013,7 +1025,6 @@ def _allocate_output(subsaveat: SubSaveAt) -> SaveState:
     #
 
     progress_meter.close(final_state.progress_meter_state)
-
     is_save_state = lambda x: isinstance(x, SaveState)
     ts = jtu.tree_map(
         lambda s: s.ts * direction, final_state.save_state, is_leaf=is_save_state

test/test_integrate.py

@@ -389,7 +389,18 @@ def g(t, y, args):
         or saveat.subs.ts is not None
         or saveat.subs.steps
     ):
-        assert tree_allclose(sol1.ts, -cast(Array, sol2.ts), equal_nan=True)
+        assert sol1.ts is not None
+        assert sol2.ts is not None
+        assert tree_allclose(
+            sol1.ts[~jnp.isinf(sol1.ts)],
+            -cast(Array, sol2.ts[~jnp.isinf(sol2.ts)]),
+            equal_nan=True,
+        )
+        assert tree_allclose(
+            sol1.ts[jnp.isinf(sol1.ts)],
+            cast(Array, sol2.ts[jnp.isinf(sol2.ts)]),
+            equal_nan=True,
+        )
         assert tree_allclose(sol1.ys, sol2.ys, equal_nan=True)
     if saveat.dense:
         t = jnp.linspace(0.3, 4, 20)
@@ -398,6 +409,36 @@ def g(t, y, args):
             assert tree_allclose(sol1.derivative(ti), -sol2.derivative(-ti))
 
 
+@pytest.mark.parametrize(
+    "saveat",
+    (
+        diffrax.SaveAt(t0=True, fn=lambda t, y, args: t),
+        diffrax.SaveAt(t1=True, fn=lambda t, y, args: t),
+        diffrax.SaveAt(dense=True, fn=lambda t, y, args: t),
+        diffrax.SaveAt(steps=True, fn=lambda t, y, args: t),
+        diffrax.SaveAt(ts=jnp.linspace(3.0, 1.0, 5), fn=lambda t, y, args: t),
+    ),
+)
+def test_reverse_time_saveat(saveat):
+    def f(t, y, args):
+        return -y
+
+    t0 = 4
+    t1 = 0.3
+    dt0 = -1 / 50
+    y0 = 1.0
+    sol1 = diffrax.diffeqsolve(
+        diffrax.ODETerm(f),
+        diffrax.Euler(),
+        t0,
+        t1,
+        dt0,
+        y0,
+        saveat=saveat,
+    )
+    assert tree_allclose(sol1.ys, sol1.ts)
+
+
 def test_semi_implicit_euler():
     term1 = diffrax.ODETerm(lambda t, y, args: -y)
     term2 = diffrax.ODETerm(lambda t, y, args: y)