diff --git a/benchmarks/jump_step_timing.py b/benchmarks/jump_step_timing.py new file mode 100644 index 00000000..933b59fc --- /dev/null +++ b/benchmarks/jump_step_timing.py @@ -0,0 +1,116 @@ +from warnings import simplefilter + + +simplefilter(action="ignore", category=FutureWarning) + +import timeit +from functools import partial + +import diffrax +import equinox as eqx +import jax +import jax.numpy as jnp +import jax.random as jr +from old_pid_controller import OldPIDController + + +t0 = 0 +t1 = 5 +dt0 = 0.5 +y0 = 1.0 +drift = diffrax.ODETerm(lambda t, y, args: -0.2 * y) + + +def diffusion_vf(t, y, args): + return jnp.ones((), dtype=y.dtype) + + +def get_terms(key): + bm = diffrax.VirtualBrownianTree(t0, t1, 2**-5, (), key) + diffusion = diffrax.ControlTerm(diffusion_vf, bm) + return diffrax.MultiTerm(drift, diffusion) + + +solver = diffrax.Heun() +step_ts = jnp.linspace(t0, t1, 129, endpoint=True) +pid_controller = diffrax.PIDController( + rtol=0, atol=1e-3, dtmin=2**-9, dtmax=1.0, pcoeff=0.3, icoeff=0.7 +) +new_controller = diffrax.ClipStepSizeController( + pid_controller, + step_ts=step_ts, + store_rejected_steps=None, +) +old_controller = OldPIDController( + rtol=0, atol=1e-3, dtmin=2**-9, dtmax=1.0, pcoeff=0.3, icoeff=0.7, step_ts=step_ts +) + + +@eqx.filter_jit +@partial(jax.vmap, in_axes=(0, None)) +def solve(key, controller): + term = get_terms(key) + return diffrax.diffeqsolve( + term, + solver, + t0, + t1, + dt0, + y0, + stepsize_controller=controller, + saveat=diffrax.SaveAt(ts=step_ts), + ) + + +num_samples = 100 +keys = jr.split(jr.PRNGKey(0), num_samples) + + +def do_timing(controller): + @jax.jit + @eqx.debug.assert_max_traces(max_traces=1) + def time_controller_fun(): + sols = solve(keys, controller) + assert sols.ys is not None + assert sols.ys.shape == (num_samples, len(step_ts)) + return sols.ys + + def time_controller(): + jax.block_until_ready(time_controller_fun()) + + return min(timeit.repeat(time_controller, number=3, repeat=20)) + + +time_new = do_timing(new_controller) + +time_old = do_timing(old_controller) + +print(f"New controller: {time_new:.5} s, Old controller: {time_old:.5} s") + +# How expensive is revisiting rejected steps? +revisiting_controller_short = diffrax.ClipStepSizeController( + pid_controller, + step_ts=step_ts, + store_rejected_steps=10, +) + +revisiting_controller_long = diffrax.ClipStepSizeController( + pid_controller, + step_ts=step_ts, + store_rejected_steps=4096, +) + +time_revisiting_short = do_timing(revisiting_controller_short) +time_revisiting_long = do_timing(revisiting_controller_long) + +print( + f"Revisiting controller\n" + f"with buffer len 10: {time_revisiting_short:.5} s\n" + f"with buffer len 4096: {time_revisiting_long:.5} s" +) + +# ======= RESULTS ======= +# New controller: 0.23506 s, Old controller: 0.30735 s +# Revisiting controller +# with buffer len 10: 0.23636 s +# with buffer len 4096: 0.23965 s diff --git a/benchmarks/old_pid_controller.py b/benchmarks/old_pid_controller.py new file mode 100644 index 00000000..f6d78098 --- /dev/null +++ b/benchmarks/old_pid_controller.py @@ -0,0 +1,414 @@ +from collections.abc import Callable +from typing import cast, Optional, TypeVar + +import equinox as eqx +import equinox.internal as eqxi +import jax +import jax.lax as lax +import jax.numpy as jnp +import jax.tree_util as jtu +import lineax.internal as lxi +import optimistix as optx +from diffrax import AbstractTerm, ODETerm, RESULTS +from diffrax._custom_types import ( + Args, + BoolScalarLike, + IntScalarLike, + RealScalarLike, + VF, + Y, +) +from diffrax._misc import static_select, upcast_or_raise +from diffrax._step_size_controller import AbstractAdaptiveStepSizeController +from equinox.internal import ω +from jaxtyping import Array, PyTree, Real +from lineax.internal import complex_to_real_dtype + + +ω = cast(Callable, ω) + + +def _select_initial_step( + terms: PyTree[AbstractTerm], + t0: RealScalarLike, + y0: Y, + args: Args, + func: Callable[ + [PyTree[AbstractTerm], RealScalarLike, Y, Args], + VF, + ], + error_order: RealScalarLike, + rtol: RealScalarLike, + atol: RealScalarLike, + norm: Callable[[PyTree], RealScalarLike], +) -> RealScalarLike: + # TODO: someone needs to figure out an initial step size algorithm for SDEs. + if not isinstance(terms, ODETerm): + return 0.01 + + def fn(carry): + t, y, _h0, _d1, _f, _ = carry + f = func(terms, t, y, args) + return t, y, _h0, _d1, _f, f + + def intermediate(carry): + _, _, _, _, _, f0 = carry + d0 = norm((y0**ω / scale**ω).ω) + d1 = norm((f0**ω / scale**ω).ω) + _cond = (d0 < 1e-5) | (d1 < 1e-5) + _d1 = jnp.where(_cond, 1, d1) + h0 = jnp.where(_cond, 1e-6, 0.01 * (d0 / _d1)) + t1 = t0 + h0 + y1 = (y0**ω + h0 * f0**ω).ω + return t1, y1, h0, d1, f0, f0 + + scale = (atol + ω(y0).call(jnp.abs) * rtol).ω + dummy_h = t0 + dummy_d = eqxi.eval_empty(norm, y0) + dummy_f = eqxi.eval_empty(lambda: func(terms, t0, y0, args)) + _, _, h0, d1, f0, f1 = eqxi.scan_trick( + fn, [intermediate], (t0, y0, dummy_h, dummy_d, dummy_f, dummy_f) + ) + d2 = norm(((f1**ω - f0**ω) / scale**ω).ω) / h0 + max_d = jnp.maximum(d1, d2) + h1 = jnp.where( + max_d <= 1e-15, + jnp.maximum(1e-6, h0 * 1e-3), + (0.01 / max_d) ** (1 / error_order), + ) + return jnp.minimum(100 * h0, h1) + + +_ControllerState = TypeVar("_ControllerState") +_Dt0 = TypeVar("_Dt0", None, RealScalarLike, Optional[RealScalarLike]) + +_PidState = tuple[ + BoolScalarLike, BoolScalarLike, RealScalarLike, RealScalarLike, RealScalarLike +] + + +def _none_or_array(x): + if x is None: + return None + else: + return jnp.asarray(x) + + +class OldPIDController( + AbstractAdaptiveStepSizeController[_PidState, Optional[RealScalarLike]] +): + r"""See the doc of diffrax.PIDController for more information.""" + + rtol: RealScalarLike + atol: RealScalarLike + pcoeff: RealScalarLike = 0 + icoeff: RealScalarLike = 1 + dcoeff: RealScalarLike = 0 + dtmin: Optional[RealScalarLike] = None + dtmax: Optional[RealScalarLike] = None + force_dtmin: bool = True + step_ts: Optional[Real[Array, " steps"]] = eqx.field( + default=None, converter=_none_or_array + ) + jump_ts: Optional[Real[Array, " jumps"]] = eqx.field( + default=None, converter=_none_or_array + ) + factormin: RealScalarLike = 0.2 + factormax: RealScalarLike = 10.0 + norm: Callable[[PyTree], RealScalarLike] = optx.rms_norm + safety: RealScalarLike = 0.9 + error_order: Optional[RealScalarLike] = None + + def __check_init__(self): + if self.jump_ts is not None and not jnp.issubdtype( + self.jump_ts.dtype, jnp.inexact + ): + raise ValueError( + f"jump_ts must be floating point, not {self.jump_ts.dtype}" + ) + + def wrap(self, direction: IntScalarLike): + step_ts = None if self.step_ts is None else self.step_ts * direction + jump_ts = None if self.jump_ts is None else self.jump_ts * direction + return eqx.tree_at( + lambda s: (s.step_ts, s.jump_ts), + self, + (step_ts, jump_ts), + is_leaf=lambda x: x is None, + ) + + def init( + self, + terms: PyTree[AbstractTerm], + t0: RealScalarLike, + t1: RealScalarLike, + y0: Y, + dt0: Optional[RealScalarLike], + args: Args, + func: Callable[[PyTree[AbstractTerm], RealScalarLike, Y, Args], VF], + error_order: Optional[RealScalarLike], + ) -> tuple[RealScalarLike, _PidState]: + del t1 + if dt0 is None: + error_order = self._get_error_order(error_order) + dt0 = _select_initial_step( + terms, + t0, + y0, + args, + func, + error_order, + self.rtol, + self.atol, + self.norm, + ) + + dt0 = lax.stop_gradient(dt0) + if self.dtmax is not None: + dt0 = jnp.minimum(dt0, self.dtmax) + if self.dtmin is None: + at_dtmin = jnp.array(False) + else: + at_dtmin = dt0 <= self.dtmin + dt0 = jnp.maximum(dt0, self.dtmin) + + t1 = self._clip_step_ts(t0, t0 + dt0) + t1, jump_next_step = self._clip_jump_ts(t0, t1) + + y_leaves = jtu.tree_leaves(y0) + if len(y_leaves) == 0: + y_dtype = lxi.default_floating_dtype() + else: + y_dtype = jnp.result_type(*y_leaves) + return t1, ( + jump_next_step, + at_dtmin, + dt0, + jnp.array(1.0, dtype=complex_to_real_dtype(y_dtype)), + jnp.array(1.0, dtype=complex_to_real_dtype(y_dtype)), + ) + + def adapt_step_size( + self, + t0: RealScalarLike, + t1: RealScalarLike, + y0: Y, + y1_candidate: Y, + args: Args, + y_error: Optional[Y], + error_order: RealScalarLike, + controller_state: _PidState, + ) -> tuple[ + BoolScalarLike, + RealScalarLike, + RealScalarLike, + BoolScalarLike, + _PidState, + RESULTS, + ]: + del args + if y_error is None and y0 is not None: + # y0 is not None check is included to handle the edge case that the state + # is just a trivial `None` PyTree. In this case `y_error` has the same + # PyTree structure and thus overlaps with our special usage of `None` to + # indicate a lack of error estimate. + raise RuntimeError( + "Cannot use adaptive step sizes with a solver that does not provide " + "error estimates." + ) + ( + made_jump, + at_dtmin, + prev_dt, + prev_inv_scaled_error, + prev_prev_inv_scaled_error, + ) = controller_state + error_order = self._get_error_order(error_order) + prev_dt = jnp.where(made_jump, prev_dt, t1 - t0) + + # + # Figure out how things went on the last step: error, and whether to + # accept/reject it. + # + + def _scale(_y0, _y1_candidate, _y_error): + # In case the solver steps into a region for which the vector field isn't + # defined. + _nan = jnp.isnan(_y1_candidate).any() + _y1_candidate = jnp.where(_nan, _y0, _y1_candidate) + _y = jnp.maximum(jnp.abs(_y0), jnp.abs(_y1_candidate)) + with jax.numpy_dtype_promotion("standard"): + return _y_error / (self.atol + _y * self.rtol) + + scaled_error = self.norm(jtu.tree_map(_scale, y0, y1_candidate, y_error)) + keep_step = scaled_error < 1 + if self.dtmin is not None: + keep_step = keep_step | at_dtmin + # Make sure it's not a Python scalar and thus getting a ZeroDivisionError. + inv_scaled_error = 1 / jnp.asarray(scaled_error) + inv_scaled_error = lax.stop_gradient( + inv_scaled_error + ) # See note in init above. + # Note: if you ever remove this lax.stop_gradient, then you'll need to do a lot + # of work to get safe gradients through these operations. + # When `inv_scaled_error` has a (non-symbolic) zero cotangent, and `y_error` + # is either zero or inf, then we get a `0 * inf = nan` on the backward pass. + + # + # Adjust next step size + # + + _zero_coeff = lambda c: isinstance(c, (int, float)) and c == 0 + coeff1 = (self.icoeff + self.pcoeff + self.dcoeff) / error_order + coeff2 = -cast(RealScalarLike, self.pcoeff + 2 * self.dcoeff) / error_order + coeff3 = self.dcoeff / error_order + factor1 = 1 if _zero_coeff(coeff1) else inv_scaled_error**coeff1 + factor2 = 1 if _zero_coeff(coeff2) else prev_inv_scaled_error**coeff2 + factor3 = 1 if _zero_coeff(coeff3) else prev_prev_inv_scaled_error**coeff3 + factormin = jnp.where(keep_step, 1, self.factormin) + factor = jnp.clip( + self.safety * factor1 * factor2 * factor3, + min=factormin, + max=self.factormax, + ) + # Once again, see above. In case we have gradients on {i,p,d}coeff. + # (Probably quite common for them to have zero tangents if passed across + # a grad API boundary as part of a larger model.) + factor = lax.stop_gradient(factor) + factor = eqxi.nondifferentiable(factor) + dt = prev_dt * factor.astype(jnp.result_type(prev_dt)) + + # E.g. we failed an implicit step, so y_error=inf, so inv_scaled_error=0, + # so factor=factormin, and we shrunk our step. + # If we're using a PI or PID controller we shouldn't then force shrinking on + # the next or next two steps as well! + pred = (inv_scaled_error == 0) | jnp.isinf(inv_scaled_error) + inv_scaled_error = jnp.where(pred, 1, inv_scaled_error) + + # + # Clip next step size based on dtmin/dtmax + # + + result = RESULTS.successful + if self.dtmax is not None: + dt = jnp.minimum(dt, self.dtmax) + if self.dtmin is None: + at_dtmin = jnp.array(False) + else: + if not self.force_dtmin: + result = RESULTS.where(dt < self.dtmin, RESULTS.dt_min_reached, result) + at_dtmin = dt <= self.dtmin + dt = jnp.maximum(dt, self.dtmin) + + # + # Clip next step size based on step_ts/jump_ts + # + + if jnp.issubdtype(jnp.result_type(t1), jnp.inexact): + # Two nextafters. If made_jump then t1 = prevbefore(jump location) + # so now _t1 = nextafter(jump location) + # This is important because we don't know whether or not the jump is as a + # result of a left- or right-discontinuity, so we have to skip the jump + # location altogether. + _t1 = static_select(made_jump, eqxi.nextafter(eqxi.nextafter(t1)), t1) + else: + _t1 = t1 + next_t0 = jnp.where(keep_step, _t1, t0) + next_t1 = self._clip_step_ts(next_t0, next_t0 + dt) + next_t1, next_made_jump = self._clip_jump_ts(next_t0, next_t1) + + inv_scaled_error = jnp.where(keep_step, inv_scaled_error, prev_inv_scaled_error) + prev_inv_scaled_error = jnp.where( + keep_step, prev_inv_scaled_error, prev_prev_inv_scaled_error + ) + controller_state = ( + next_made_jump, + at_dtmin, + dt, + inv_scaled_error, + prev_inv_scaled_error, + ) + return keep_step, next_t0, next_t1, made_jump, controller_state, result + + def _get_error_order(self, error_order: Optional[RealScalarLike]) -> RealScalarLike: + # Attribute takes priority, if the user knows the correct error order better + # than our guess. + error_order = error_order if self.error_order is None else self.error_order + if error_order is None: + raise ValueError( + "The order of convergence for the solver has not been specified; pass " + "`PIDController(..., error_order=...)` manually instead. If solving " + "an ODE then this should be equal to the (global) order plus one. If " + "solving an SDE then should be equal to the (global) order plus 0.5." + ) + return error_order + + def _clip_step_ts(self, t0: RealScalarLike, t1: RealScalarLike) -> RealScalarLike: + if self.step_ts is None: + return t1 + + step_ts0 = upcast_or_raise( + self.step_ts, + t0, + "`PIDController.step_ts`", + "time (the result type of `t0`, `t1`, `dt0`, `SaveAt(ts=...)` etc.)", + ) + step_ts1 = upcast_or_raise( + self.step_ts, + t1, + "`PIDController.step_ts`", + "time (the result type of `t0`, `t1`, `dt0`, `SaveAt(ts=...)` etc.)", + ) + # TODO: it should be possible to switch this O(nlogn) for just O(n) by keeping + # track of where we were last, and using that as a hint for the next search. + t0_index = jnp.searchsorted(step_ts0, t0, side="right") + t1_index = jnp.searchsorted(step_ts1, t1, side="right") + # This minimum may or may not actually be necessary. The left branch is taken + # iff t0_index < t1_index <= len(self.step_ts), so all valid t0_index s must + # already satisfy the minimum. + # However, that branch is actually executed unconditionally and then where'd, + # so we clamp it just to be sure we're not hitting undefined behaviour. + t1 = jnp.where( + t0_index < t1_index, + step_ts1[jnp.minimum(t0_index, len(self.step_ts) - 1)], + t1, + ) + return t1 + + def _clip_jump_ts( + self, t0: RealScalarLike, t1: RealScalarLike + ) -> tuple[RealScalarLike, BoolScalarLike]: + if self.jump_ts is None: + return t1, False + assert jnp.issubdtype(self.jump_ts.dtype, jnp.inexact) + if not jnp.issubdtype(jnp.result_type(t0), jnp.inexact): + raise ValueError( + "`t0`, `t1`, `dt0` must be floating point when specifying `jump_ts`. " + f"Got {jnp.result_type(t0)}." + ) + if not jnp.issubdtype(jnp.result_type(t1), jnp.inexact): + raise ValueError( + "`t0`, `t1`, `dt0` must be floating point when specifying `jump_ts`. " + f"Got {jnp.result_type(t1)}." + ) + jump_ts0 = upcast_or_raise( + self.jump_ts, + t0, + "`PIDController.jump_ts`", + "time (the result type of `t0`, `t1`, `dt0`, `SaveAt(ts=...)` etc.)", + ) + jump_ts1 = upcast_or_raise( + self.jump_ts, + t1, + "`PIDController.jump_ts`", + "time (the result type of `t0`, `t1`, `dt0`, `SaveAt(ts=...)` etc.)", + ) + t0_index = jnp.searchsorted(jump_ts0, t0, side="right") + t1_index = jnp.searchsorted(jump_ts1, t1, side="right") + next_made_jump = t0_index < t1_index + t1 = jnp.where( + next_made_jump, + eqxi.prevbefore(jump_ts1[jnp.minimum(t0_index, len(self.jump_ts) - 1)]), + t1, + ) + return t1, next_made_jump diff --git a/diffrax/__init__.py b/diffrax/__init__.py index 42073a10..d35a7fac 100644 --- a/diffrax/__init__.py +++ b/diffrax/__init__.py @@ -121,6 +121,7 @@ from ._step_size_controller import ( AbstractAdaptiveStepSizeController as AbstractAdaptiveStepSizeController, AbstractStepSizeController as AbstractStepSizeController, + ClipStepSizeController as ClipStepSizeController, ConstantStepSize as ConstantStepSize, PIDController as PIDController, StepTo as StepTo, diff --git a/diffrax/_autocitation.py b/diffrax/_autocitation.py index 547177ce..c2cdcade 100644 --- a/diffrax/_autocitation.py +++ b/diffrax/_autocitation.py @@ -36,7 +36,7 @@ SRA1, Tsit5, ) -from ._step_size_controller import PIDController +from ._step_size_controller import ClipStepSizeController, PIDController def citation(*args, **kwargs): @@ -134,7 +134,7 @@ def citation(*args, **kwargs): _thesis_cite = r""" -phdthesis{kidger2021on, +@phdthesis{kidger2021on, title={{O}n {N}eural {D}ifferential {E}quations}, author={Patrick Kidger}, year={2021}, @@ -352,10 +352,10 @@ def _virtual_brownian_tree(terms): return ( r""" % You are simulating Brownian motion using a virtual Brownian tree, which was introduced -% in: +% in the following two papers: """ + vbt_ref - + "\n\n" + + "\n" + single_seed_ref ) @@ -570,6 +570,17 @@ def _auto_dt0(dt0): """ +@citation_rules.append +def _clip_controller(terms, stepsize_controller): + if type(stepsize_controller) is ClipStepSizeController: + if stepsize_controller.store_rejected_steps is not None and is_sde(terms): + return r""" +% You are adaptively solving an SDE whilst revisiting rejected time points. This is a +% subtle point required for the correctness of adaptive noncommutative SDE solves, as +% found in: +""" + _parse_reference(ClipStepSizeController) + + @citation_rules.append def _pid_controller(stepsize_controller, terms=None): if type(stepsize_controller) is PIDController: diff --git a/diffrax/_misc.py b/diffrax/_misc.py index 7c6fa53b..7d52fbde 100644 --- a/diffrax/_misc.py +++ b/diffrax/_misc.py @@ -1,5 +1,5 @@ from collections.abc import Callable -from typing import Any, cast, Optional +from typing import Any, cast, Optional, Union import jax import jax.core @@ -160,7 +160,10 @@ def static_select(pred: BoolScalarLike, a: ArrayLike, b: ArrayLike) -> ArrayLike def upcast_or_raise( - x: ArrayLike, array_for_dtype: ArrayLike, x_name: str, dtype_name: str + x: ArrayLike, + array_for_dtype: Union[ArrayLike, jnp.dtype], + x_name: str, + dtype_name: str, ): """If `JAX_NUMPY_DTYPE_PROMOTION=strict`, then this will raise an error if `jnp.result_type(x, array_for_dtype)` is not the same as `array_for_dtype.dtype`. diff --git a/diffrax/_solution.py b/diffrax/_solution.py index f1b8d21b..e99f2c15 100644 --- a/diffrax/_solution.py +++ b/diffrax/_solution.py @@ -21,6 +21,14 @@ class RESULTS(optx.RESULTS): # pyright: ignore event_occurred = ( "Terminating differential equation solve because an event occurred." ) + max_steps_rejected = ( + "Maximum number of rejected steps was reached. Consider increasing " + "`diffrax.ClipStepSizeController(store_rejected_steps==...)`." + ) + internal_error = ( + "An internal error occurred in Diffrax. This is a bug! Please open a GitHub " + "issue with a minimum working example. (<50 lines of code is ideal)" + ) # Backward compatibility diff --git a/diffrax/_step_size_controller/__init__.py b/diffrax/_step_size_controller/__init__.py index 18d19c00..74e9371f 100644 --- a/diffrax/_step_size_controller/__init__.py +++ b/diffrax/_step_size_controller/__init__.py @@ -1,6 +1,9 @@ -from .adaptive import ( +from .base import ( AbstractAdaptiveStepSizeController as AbstractAdaptiveStepSizeController, - PIDController as PIDController, + AbstractStepSizeController as AbstractStepSizeController, ) -from .base import AbstractStepSizeController as AbstractStepSizeController +from .clip import ClipStepSizeController as ClipStepSizeController from .constant import ConstantStepSize as ConstantStepSize, StepTo as StepTo +from .pid import ( + PIDController as PIDController, +) diff --git a/diffrax/_step_size_controller/base.py b/diffrax/_step_size_controller/base.py index 625bd6fb..9e6059ca 100644 --- a/diffrax/_step_size_controller/base.py +++ b/diffrax/_step_size_controller/base.py @@ -3,6 +3,7 @@ from typing import Generic, Optional, TypeVar import equinox as eqx +from equinox import AbstractVar from jaxtyping import PyTree from .._custom_types import Args, BoolScalarLike, IntScalarLike, RealScalarLike, VF, Y @@ -11,7 +12,7 @@ _ControllerState = TypeVar("_ControllerState") -_Dt0 = TypeVar("_Dt0", None, RealScalarLike, Optional[RealScalarLike]) +_Dt0 = TypeVar("_Dt0", bound=Optional[RealScalarLike]) class AbstractStepSizeController(eqx.Module, Generic[_ControllerState, _Dt0]): @@ -127,3 +128,33 @@ def adapt_step_size( happened successfully, or if it failed for some reason. (e.g. hitting a minimum allowed step size in the solver.) """ + + +class AbstractAdaptiveStepSizeController( + AbstractStepSizeController[_ControllerState, _Dt0] +): + """Indicates an adaptive step size controller. + + Accepts tolerances `rtol` and `atol`. When used in conjunction with an implicit + solver ([`diffrax.AbstractImplicitSolver`][]), then these tolerances will + automatically be used as the tolerances for the nonlinear solver passed to the + implicit solver, if they are not specified manually. + """ + + rtol: AbstractVar[RealScalarLike] + atol: AbstractVar[RealScalarLike] + norm: AbstractVar[Callable[[PyTree], RealScalarLike]] + + def __check_init__(self): + if self.rtol is None or self.atol is None: + raise ValueError( + "The default values for `rtol` and `atol` were removed in Diffrax " + "version 0.1.0. (As the choice of tolerance is nearly always " + "something that you, as an end user, should make an explicit choice " + "about.)\n" + "If you want to match the previous defaults then specify " + "`rtol=1e-3`, `atol=1e-6`. For example:\n" + "```\n" + "diffrax.PIDController(rtol=1e-3, atol=1e-6)\n" + "```\n" + ) diff --git a/diffrax/_step_size_controller/clip.py b/diffrax/_step_size_controller/clip.py new file mode 100644 index 00000000..0167e789 --- /dev/null +++ b/diffrax/_step_size_controller/clip.py @@ -0,0 +1,394 @@ +from collections.abc import Callable +from typing import cast, Generic, Optional, TypeVar + +import equinox as eqx +import equinox.internal as eqxi +import jax +import jax.numpy as jnp +from jaxtyping import Array, PyTree, Real + +from .._custom_types import ( + Args, + BoolScalarLike, + FloatScalarLike, + IntScalarLike, + RealScalarLike, + VF, + Y, +) +from .._misc import upcast_or_raise +from .._solution import is_okay, RESULTS +from .._term import AbstractTerm +from .base import AbstractStepSizeController + + +_ControllerState = TypeVar("_ControllerState") +_Dt0 = TypeVar("_Dt0", bound=Optional[RealScalarLike]) + + +class _ClipState(eqx.Module, Generic[_ControllerState]): + step_info: Optional[tuple[IntScalarLike, Array]] + jump_info: Optional[tuple[IntScalarLike, Array]] + reject_info: Optional[tuple[IntScalarLike, Array]] + inner_state: _ControllerState + + +def _none_or_sorted_array(x): + if x is None: + return None + else: + return jnp.sort(jnp.asarray(x)) + + +def _assert_floating(t: FloatScalarLike, name: str, dtype): + t_dtype = jnp.result_type(t) + if not jnp.issubdtype(t_dtype, jnp.floating): + raise ValueError(f"{name} must be floating-point, got {t_dtype}") + if t_dtype != dtype: + raise ValueError( + f"All timelike inputs must have the same dtype got both {dtype} and " + f"{t_dtype}." + ) + + +def _get_t(i: IntScalarLike, ts: Array) -> RealScalarLike: + # As `ts[i]`, but `ts[len(ts))]` returns `inf`. + # `i` must be in `{0, 1, ..., len(ts)}`. + if len(ts) == 0: + return jnp.inf + else: + i_min_len = jnp.minimum(i, len(ts) - 1) + return jnp.where(i == len(ts), jnp.inf, ts[i_min_len]) + + +def _clip_t( + t: FloatScalarLike, + i: IntScalarLike, + ts: Array, + prevbefore: bool, +) -> FloatScalarLike: + assert jnp.issubdtype(jnp.result_type(t), jnp.floating) + assert jnp.result_type(t) == jnp.result_type(ts) + _t = _get_t(i, ts) + if prevbefore: + _t = eqxi.prevbefore(_t) + return jnp.minimum(_t, t) + + +def _bump_next_t0(next_t0, ts): + # Our previous step may have been to prevbefore a jump. + # In this case we want to bump our next step to occur nextafter the jump. + # We don't test against just `jump_ts[jump_index]`. The index in the state + # is intended only as a hint to improve the efficiency of + # `_find_idx_with_hint`; it's not load-bearing. This is for safety, in case some + # other stepsize control is going on. (TODO: do we want to keep it like this, or + # do we want to switch to just the single check?) + nextafter_next_t0 = eqxi.nextafter(next_t0) + made_jump1 = jnp.any(nextafter_next_t0 == ts) + # For safety we also test `next_t0 == ts`, just in case some other stepsize control + # is going on. (I don't think this should actually be necessary.) + made_jump2 = jnp.any(next_t0 == ts) + # There are two nextafters. This is important because we don't know whether + # or not the jump is a left- or a right-discontinuity, so we skip the jump + # time altogether. + next_t0 = jnp.where(made_jump1, eqxi.nextafter(nextafter_next_t0), next_t0) + next_t0 = cast(Array, next_t0) + next_t0 = jnp.where(made_jump2, nextafter_next_t0, next_t0) + next_t0 = cast(Array, next_t0) + return next_t0, made_jump1 | made_jump2 + + +def _find_idx_with_hint(t: RealScalarLike, ts: Optional[Array], hint: IntScalarLike): + # Find index of first element of `ts` strictly greater than `t`. + # Uses a linear search starting from `hint`. The value `hint` is assumed to be in + # `{0, 1, ..., len(ts)}` + if ts is None: + return 0 + + def cond_up(_i): + return (_i < len(ts)) & (ts[_i] <= t) + + def cond_down(_i): + return (_i > 0) & (ts[_i - 1] > t) + + i = hint + i = jax.lax.while_loop(cond_up, lambda _i: _i + 1, i) + i = jax.lax.while_loop(cond_down, lambda _i: _i - 1, i) + return i + + +class ClipStepSizeController( + AbstractStepSizeController[_ClipState[_ControllerState], _Dt0] +): + """Wraps an existing step controller with three pieces of functionality: + + - Have the solver step exactly to certain times ('`step_ts`'). + - Have the solver step to just before and just after certain time ('`jump_ts`'). + - Have the solver record the times of rejected steps, and step exactly to those + times in future steps. + + In all cases this essentially corresponds to clipping steps so that any that are + 'too large' will instead by clipped from one of the three above cases. + + Stepping exactly to certain times can be useful if you want to ensure that your + solution is highly accurate at that exact time point -- by default Diffrax will + adaptively step wherever it likes, and then interpolate to produce the output values + in `SaveAt(ts=...)`. + + Specifying jump times is needed for computational efficiency when solving + differential equations for which the vector field has known jumps (e.g. due to a + discontinuous forcing term). Otherwise an adaptive solver must reject many steps as + it slows down to try and locate a jump. When using this, the solver will step to the + floating point number immediately before the jump, and then resume solving from the + floating point number immediately after it, with the jump itself not being + evaluated. + + Revisiting rejected steps is needed when adaptively solving SDEs with noncommutative + noise. Otherwise, a small bias may be introduced in the higher-order (Lévy area) + terms of the solution, as it is possible to reject a step *because* of the samples + drawn in these higher order terms. + + ??? Citation + + For more details on revisiting rejected steps when adaptively solving SDEs, see: + + ```bibtex + @misc{foster2024convergenceadaptiveapproximationsstochastic, + title={On the convergence of adaptive approximations for + stochastic differential equations}, + author={James Foster and Andraž Jelinčič}, + year={2024}, + eprint={2311.14201}, + archivePrefix={arXiv}, + primaryClass={math.NA}, + url={https://arxiv.org/abs/2311.14201}, + } + ``` + """ + + controller: AbstractStepSizeController[_ControllerState, _Dt0] + step_ts: Optional[Real[Array, " steps"]] + jump_ts: Optional[Real[Array, " jumps"]] + store_rejected_steps: Optional[int] = eqx.field(static=True) + callback_on_reject: Optional[Callable] = eqx.field(static=True) + + @eqxi.doc_remove_args("_callback_on_reject") + def __init__( + self, + controller, + step_ts=None, + jump_ts=None, + store_rejected_steps=None, + _callback_on_reject=None, + ): + """**Arguments**: + + - `controller`: The controller to wrap. + Can be any [`diffrax.AbstractAdaptiveStepSizeController`][]. + - `step_ts`: Denotes extra times that must be stepped to. + - `jump_ts`: Denotes extra times that must be stepped to, and at which the + vector field has a known discontinuity. (This is used to force FSAL solvers + to re-evaluate the vector field.) + `store_rejected_steps`: If this is set to a positive integer, then any + rejected steps will have their time stored, and that time will be stepped to + exactly in a later step. This is used when solving SDEs with noncommutative + noise, for which this ensures that the distribution coming from Lévy area + terms is correct. Setting this to e.g. `100` should be plenty, but if more + consecutive steps are rejected, then a runtime error will be raised. (Note + that this is not the total number of rejected steps in a solve, but just the + maximum number of *consecutive* rejected steps.) + """ + self.controller = controller + self.step_ts = _none_or_sorted_array(step_ts) + self.jump_ts = _none_or_sorted_array(jump_ts) + if (store_rejected_steps is not None) and (store_rejected_steps <= 0): + raise ValueError( + "`store_rejected_steps must either be `None`" + " or a non-negative integer." + ) + self.store_rejected_steps = store_rejected_steps + self.callback_on_reject = _callback_on_reject + + def __check_init__(self): + if self.jump_ts is not None and not jnp.issubdtype( + self.jump_ts.dtype, jnp.floating + ): + raise ValueError( + f"jump_ts must be floating point, not {self.jump_ts.dtype}" + ) + + def wrap(self, direction: IntScalarLike): + step_ts = None if self.step_ts is None else jnp.sort(self.step_ts * direction) + jump_ts = None if self.jump_ts is None else jnp.sort(self.jump_ts * direction) + controller = self.controller.wrap(direction) + return eqx.tree_at( + lambda s: (s.step_ts, s.jump_ts, s.controller), + self, + (step_ts, jump_ts, controller), + is_leaf=lambda x: x is None, + ) + + def init( + self, + terms: PyTree[AbstractTerm], + t0: RealScalarLike, + t1: RealScalarLike, + y0: Y, + dt0: _Dt0, + args: Args, + func: Callable[[PyTree[AbstractTerm], RealScalarLike, Y, Args], VF], + error_order: Optional[RealScalarLike], + ) -> tuple[RealScalarLike, _ClipState[_ControllerState]]: + t_dtype = jnp.result_type(t0) + _assert_floating(t0, "t0", t_dtype) + _assert_floating(t1, "t1", t_dtype) + if dt0 is not None: + _assert_floating(dt0, "dt0", t_dtype) + t1, inner_state = self.controller.init( + terms, t0, t1, y0, dt0, args, func, error_order + ) + _assert_floating(t1, "controller.init(...)", t_dtype) + + if self.step_ts is None: + step_info = None + else: + step_ts = upcast_or_raise( + self.step_ts, + t_dtype, + "`ClipStepSizeController.step_ts`", + "time (the result type of `t0`, `t1`, `dt0`, `SaveAt(ts=...)` etc.)", + ) + step_index = jnp.searchsorted(step_ts, t0, side="right") + t1 = _clip_t(t1, step_index, step_ts, False) + step_info = (step_index, step_ts) + + if self.jump_ts is None: + jump_info = None + else: + jump_ts = upcast_or_raise( + self.jump_ts, + t_dtype, + "`ClipStepSizeController.jump_ts`", + "time (the result type of `t0`, `t1`, `dt0`, `SaveAt(ts=...)` etc.)", + ) + jump_index = jnp.searchsorted(jump_ts, t0, side="right") + t1 = _clip_t(t1, jump_index, jump_ts, True) + jump_info = (jump_index, jump_ts) + + if self.store_rejected_steps is None: + reject_info = None + else: + reject_ts = jnp.zeros(self.store_rejected_steps, dtype=t_dtype) + reject_index = jnp.array(self.store_rejected_steps) + reject_info = (reject_index, reject_ts) + + state = _ClipState(step_info, jump_info, reject_info, inner_state) + return t1, state + + def adapt_step_size( + self, + t0: RealScalarLike, + t1: RealScalarLike, + y0: Y, + y1_candidate: Y, + args: Args, + y_error: Optional[Y], + error_order: RealScalarLike, + controller_state: _ClipState[_ControllerState], + ) -> tuple[ + BoolScalarLike, + RealScalarLike, + RealScalarLike, + BoolScalarLike, + _ClipState[_ControllerState], + RESULTS, + ]: + t_dtype = jnp.result_type(t0) + _assert_floating(t0, "t0", t_dtype) + _assert_floating(t1, "t1", t_dtype) + ( + keep_step, + next_t0, + next_t1, + made_jump, + inner_state, + result, + ) = self.controller.adapt_step_size( + t0, + t1, + y0, + y1_candidate, + args, + y_error, + error_order, + controller_state.inner_state, + ) + _assert_floating(next_t0, "next_t0", t_dtype) + _assert_floating(next_t1, "next_t1", t_dtype) + + # Logging utility for testing purposes + callback_on_reject = self.callback_on_reject + if callback_on_reject is not None: + + def callback(_keep_step, _t1): + callback_on_reject(_keep_step, _t1) + return _keep_step + + keep_step = jax.pure_callback(callback, keep_step, keep_step, t1) + + if controller_state.step_info is None: + step_info = None + else: + step_index, step_ts = controller_state.step_info + # We actaully bump `next_t0` past any `step_ts` whilst checking where to + # clip `next_t1`. This is in case we have a set up like the following: + # ```python + # ClipStepSizeController( + # ClipStepSizeController(..., step_ts=[x]), jump_ts=[x] + # ) + # ``` + # with a single value `x`. Otherwise in this case, the outer controller will + # propose a step over the interval [something, prevbefore(x)], then on the + # next step the inner controller will propose a step over [prevbefore(x), x] + # which definitely isn't desired! + _next_t0, _ = _bump_next_t0(next_t0, step_ts) + step_index = _find_idx_with_hint(_next_t0, step_ts, step_index) + next_t1 = _clip_t(next_t1, step_index, step_ts, False) + step_info = step_index, step_ts + if controller_state.jump_info is None: + jump_info = None + else: + jump_index, jump_ts = controller_state.jump_info + next_t0, made_jump2 = _bump_next_t0(next_t0, jump_ts) + made_jump = made_jump | made_jump2 + jump_index = _find_idx_with_hint(next_t0, jump_ts, jump_index) + next_t1 = _clip_t(next_t1, jump_index, jump_ts, True) + jump_info = jump_index, jump_ts + if controller_state.reject_info is None: + reject_info = None + else: + assert self.store_rejected_steps is not None + reject_index, reject_ts = controller_state.reject_info + # If the step ended at `t1==reject_ts[reject_index],` then we have + # successfully stepped to this time and we pop off this rejected time by + # incrementing `reject_index`. + # We do this increment even if the step is rejected, because we will + # re-add the rejected time to the buffer immediately. + rejected_t = _get_t(reject_index, reject_ts) + result = RESULTS.where( + (t1 > rejected_t) & is_okay(result), RESULTS.internal_error, result + ) + reject_index = reject_index + jnp.where(t1 == rejected_t, 1, 0) + # Now, if the step is rejected then we must store the rejected time in the + # buffer. + reject_index = reject_index - jnp.where(keep_step, 0, 1) + result = RESULTS.where( + (reject_index < 0) & is_okay(result), RESULTS.max_steps_rejected, result + ) + new_rejected_t = jnp.where(keep_step, reject_ts[reject_index], t1) + reject_ts = reject_ts.at[reject_index].set(new_rejected_t) + next_t1 = _clip_t(next_t1, reject_index, reject_ts, False) + reject_info = reject_index, reject_ts + + state = _ClipState(step_info, jump_info, reject_info, inner_state) + return keep_step, next_t0, next_t1, made_jump, state, result diff --git a/diffrax/_step_size_controller/adaptive.py b/diffrax/_step_size_controller/pid.py similarity index 74% rename from diffrax/_step_size_controller/adaptive.py rename to diffrax/_step_size_controller/pid.py index 9d181c95..710ae944 100644 --- a/diffrax/_step_size_controller/adaptive.py +++ b/diffrax/_step_size_controller/pid.py @@ -1,6 +1,6 @@ import typing from collections.abc import Callable -from typing import cast, Optional, TYPE_CHECKING, TypeVar +from typing import cast, Optional, TYPE_CHECKING import equinox as eqx import equinox.internal as eqxi @@ -10,15 +10,8 @@ import jax.tree_util as jtu import lineax.internal as lxi import optimistix as optx -from jaxtyping import Real - - -if TYPE_CHECKING: - from typing import ClassVar as AbstractVar -else: - from equinox import AbstractVar from equinox.internal import ω -from jaxtyping import Array, PyTree +from jaxtyping import PyTree from lineax.internal import complex_to_real_dtype from .._custom_types import ( @@ -29,10 +22,10 @@ VF, Y, ) -from .._misc import static_select, upcast_or_raise from .._solution import RESULTS from .._term import AbstractTerm, ODETerm -from .base import AbstractStepSizeController +from .base import AbstractAdaptiveStepSizeController +from .clip import ClipStepSizeController ω = cast(Callable, ω) @@ -89,50 +82,25 @@ def intermediate(carry): return jnp.minimum(100 * h0, h1) -_ControllerState = TypeVar("_ControllerState") -_Dt0 = TypeVar("_Dt0", None, RealScalarLike, Optional[RealScalarLike]) - - -class AbstractAdaptiveStepSizeController( - AbstractStepSizeController[_ControllerState, _Dt0] -): - """Indicates an adaptive step size controller. - - Accepts tolerances `rtol` and `atol`. When used in conjunction with an implicit - solver ([`diffrax.AbstractImplicitSolver`][]), then these tolerances will - automatically be used as the tolerances for the nonlinear solver passed to the - implicit solver, if they are not specified manually. - """ - - rtol: AbstractVar[RealScalarLike] - atol: AbstractVar[RealScalarLike] - norm: AbstractVar[Callable[[PyTree], RealScalarLike]] - - def __check_init__(self): - if self.rtol is None or self.atol is None: - raise ValueError( - "The default values for `rtol` and `atol` were removed in Diffrax " - "version 0.1.0. (As the choice of tolerance is nearly always " - "something that you, as an end user, should make an explicit choice " - "about.)\n" - "If you want to match the previous defaults then specify " - "`rtol=1e-3`, `atol=1e-6`. For example:\n" - "```\n" - "diffrax.PIDController(rtol=1e-3, atol=1e-6)\n" - "```\n" - ) +# _PidState = (prev_inv_scaled_error, prev_prev_inv_scaled_error) +_PidState = tuple[RealScalarLike, RealScalarLike] -_PidState = tuple[ - BoolScalarLike, BoolScalarLike, RealScalarLike, RealScalarLike, RealScalarLike -] +# We use a metaclass for backwards compatibility. When a user calls +# PIDController(... step_ts=s, jump_ts=j) this should return a +# ClipStepSizeController(PIDController(...), s, j). +class _MetaPID(type(eqx.Module)): + def __call__(cls, *args, **kwargs): + step_ts = kwargs.pop("step_ts", None) + jump_ts = kwargs.pop("jump_ts", None) + if step_ts is not None or jump_ts is not None: + return ClipStepSizeController(cls(*args, **kwargs), step_ts, jump_ts) + return super().__call__(*args, **kwargs) -def _none_or_array(x): - if x is None: - return None - else: - return jnp.asarray(x) +# Sneak the metaclass past pyright, as otherwise it disables the dataclass-ness of +# `eqx.Module`. +_set_metaclass = dict(metaclass=_MetaPID) if TYPE_CHECKING: @@ -157,7 +125,8 @@ def __repr__(self): # TODO: we don't currently offer a limiter, or a variant accept/reject scheme, as given # in Soderlind and Wang 2006. class PIDController( - AbstractAdaptiveStepSizeController[_PidState, Optional[RealScalarLike]] + AbstractAdaptiveStepSizeController[_PidState, Optional[RealScalarLike]], + **_set_metaclass, ): r"""Adapts the step size to produce a solution accurate to a given tolerance. The tolerance is calculated as `atol + rtol * y` for the evolving solution `y`. @@ -347,41 +316,20 @@ def dynamics(t, y, args): rtol: RealScalarLike atol: RealScalarLike + norm: Callable[[PyTree], RealScalarLike] = rms_norm pcoeff: RealScalarLike = 0 icoeff: RealScalarLike = 1 dcoeff: RealScalarLike = 0 dtmin: Optional[RealScalarLike] = None dtmax: Optional[RealScalarLike] = None force_dtmin: bool = True - step_ts: Optional[Real[Array, " steps"]] = eqx.field( - default=None, converter=_none_or_array - ) - jump_ts: Optional[Real[Array, " jumps"]] = eqx.field( - default=None, converter=_none_or_array - ) factormin: RealScalarLike = 0.2 factormax: RealScalarLike = 10.0 - norm: Callable[[PyTree], RealScalarLike] = rms_norm safety: RealScalarLike = 0.9 error_order: Optional[RealScalarLike] = None - def __check_init__(self): - if self.jump_ts is not None and not jnp.issubdtype( - self.jump_ts.dtype, jnp.inexact - ): - raise ValueError( - f"jump_ts must be floating point, not {self.jump_ts.dtype}" - ) - def wrap(self, direction: IntScalarLike): - step_ts = None if self.step_ts is None else self.step_ts * direction - jump_ts = None if self.jump_ts is None else self.jump_ts * direction - return eqx.tree_at( - lambda s: (s.step_ts, s.jump_ts), - self, - (step_ts, jump_ts), - is_leaf=lambda x: x is None, - ) + return self def init( self, @@ -444,26 +392,20 @@ def init( dt0 = lax.stop_gradient(dt0) if self.dtmax is not None: dt0 = jnp.minimum(dt0, self.dtmax) - if self.dtmin is None: - at_dtmin = jnp.array(False) - else: - at_dtmin = dt0 <= self.dtmin + if self.dtmin is not None: dt0 = jnp.maximum(dt0, self.dtmin) - t1 = self._clip_step_ts(t0, t0 + dt0) - t1, jump_next_step = self._clip_jump_ts(t0, t1) + t1 = t0 + dt0 y_leaves = jtu.tree_leaves(y0) if len(y_leaves) == 0: y_dtype = lxi.default_floating_dtype() else: y_dtype = jnp.result_type(*y_leaves) + real_dtype = complex_to_real_dtype(y_dtype) return t1, ( - jump_next_step, - at_dtmin, - dt0, - jnp.array(1.0, dtype=complex_to_real_dtype(y_dtype)), - jnp.array(1.0, dtype=complex_to_real_dtype(y_dtype)), + jnp.array(1.0, dtype=real_dtype), + jnp.array(1.0, dtype=real_dtype), ) def adapt_step_size( @@ -543,22 +485,11 @@ def adapt_step_size( "error estimates." ) ( - made_jump, - at_dtmin, - prev_dt, prev_inv_scaled_error, prev_prev_inv_scaled_error, ) = controller_state error_order = self._get_error_order(error_order) - # t1 - t0 is the step we actually took, so that's usually what we mean by the - # "previous dt". - # However if we made a jump then this t1 was clipped relatively to what it - # could have been, so for guessing the next step size it's probably better to - # use the size the step would have been, had there been no jump. - # There are cases in which something besides the step size controller modifies - # the step locations t0, t1; most notably the main integration routine clipping - # steps when we're right at the end of the interval. - prev_dt = jnp.where(made_jump, prev_dt, t1 - t0) + prev_dt = t1 - t0 # # Figure out how things went on the last step: error, and whether to @@ -576,8 +507,9 @@ def _scale(_y0, _y1_candidate, _y_error): scaled_error = self.norm(jtu.tree_map(_scale, y0, y1_candidate, y_error)) keep_step = scaled_error < 1 + # Automatically keep the step if we're at dtmin. if self.dtmin is not None: - keep_step = keep_step | at_dtmin + keep_step = keep_step | (prev_dt <= self.dtmin) # Make sure it's not a Python scalar and thus getting a ZeroDivisionError. inv_scaled_error = 1 / jnp.asarray(scaled_error) inv_scaled_error = lax.stop_gradient( @@ -600,10 +532,12 @@ def _scale(_y0, _y1_candidate, _y_error): factor2 = 1 if _zero_coeff(coeff2) else prev_inv_scaled_error**coeff2 factor3 = 1 if _zero_coeff(coeff3) else prev_prev_inv_scaled_error**coeff3 factormin = jnp.where(keep_step, 1, self.factormin) + # If the step is not kept, next step must be smaller, so factor must be <1. + factormax = jnp.where(keep_step, self.factormax, self.safety) factor = jnp.clip( self.safety * factor1 * factor2 * factor3, min=factormin, - max=self.factormax, + max=factormax, ) # Once again, see above. In case we have gradients on {i,p,d}coeff. # (Probably quite common for them to have zero tangents if passed across @@ -626,43 +560,22 @@ def _scale(_y0, _y1_candidate, _y_error): result = RESULTS.successful if self.dtmax is not None: dt = jnp.minimum(dt, self.dtmax) - if self.dtmin is None: - at_dtmin = jnp.array(False) - else: + if self.dtmin is not None: if not self.force_dtmin: result = RESULTS.where(dt < self.dtmin, RESULTS.dt_min_reached, result) - at_dtmin = dt <= self.dtmin dt = jnp.maximum(dt, self.dtmin) - # - # Clip next step size based on step_ts/jump_ts - # - - if jnp.issubdtype(jnp.result_type(t1), jnp.inexact): - # Two nextafters. If made_jump then t1 = prevbefore(jump location) - # so now _t1 = nextafter(jump location) - # This is important because we don't know whether or not the jump is as a - # result of a left- or right-discontinuity, so we have to skip the jump - # location altogether. - _t1 = static_select(made_jump, eqxi.nextafter(eqxi.nextafter(t1)), t1) - else: - _t1 = t1 - next_t0 = jnp.where(keep_step, _t1, t0) - next_t1 = self._clip_step_ts(next_t0, next_t0 + dt) - next_t1, next_made_jump = self._clip_jump_ts(next_t0, next_t1) + next_t0 = jnp.where(keep_step, t1, t0) + next_t1 = next_t0 + dt inv_scaled_error = jnp.where(keep_step, inv_scaled_error, prev_inv_scaled_error) prev_inv_scaled_error = jnp.where( keep_step, prev_inv_scaled_error, prev_prev_inv_scaled_error ) - controller_state = ( - next_made_jump, - at_dtmin, - dt, - inv_scaled_error, - prev_inv_scaled_error, - ) - return keep_step, next_t0, next_t1, made_jump, controller_state, result + controller_state = inv_scaled_error, prev_inv_scaled_error + # made_jump is handled by ClipStepSizeController, so we automatically set it to + # False + return keep_step, next_t0, next_t1, False, controller_state, result def _get_error_order(self, error_order: Optional[RealScalarLike]) -> RealScalarLike: # Attribute takes priority, if the user knows the correct error order better @@ -677,76 +590,6 @@ def _get_error_order(self, error_order: Optional[RealScalarLike]) -> RealScalarL ) return error_order - def _clip_step_ts(self, t0: RealScalarLike, t1: RealScalarLike) -> RealScalarLike: - if self.step_ts is None: - return t1 - - step_ts0 = upcast_or_raise( - self.step_ts, - t0, - "`PIDController.step_ts`", - "time (the result type of `t0`, `t1`, `dt0`, `SaveAt(ts=...)` etc.)", - ) - step_ts1 = upcast_or_raise( - self.step_ts, - t1, - "`PIDController.step_ts`", - "time (the result type of `t0`, `t1`, `dt0`, `SaveAt(ts=...)` etc.)", - ) - # TODO: it should be possible to switch this O(nlogn) for just O(n) by keeping - # track of where we were last, and using that as a hint for the next search. - t0_index = jnp.searchsorted(step_ts0, t0, side="right") - t1_index = jnp.searchsorted(step_ts1, t1, side="right") - # This minimum may or may not actually be necessary. The left branch is taken - # iff t0_index < t1_index <= len(self.step_ts), so all valid t0_index s must - # already satisfy the minimum. - # However, that branch is actually executed unconditionally and then where'd, - # so we clamp it just to be sure we're not hitting undefined behaviour. - t1 = jnp.where( - t0_index < t1_index, - step_ts1[jnp.minimum(t0_index, len(self.step_ts) - 1)], - t1, - ) - return t1 - - def _clip_jump_ts( - self, t0: RealScalarLike, t1: RealScalarLike - ) -> tuple[RealScalarLike, BoolScalarLike]: - if self.jump_ts is None: - return t1, False - assert jnp.issubdtype(self.jump_ts.dtype, jnp.inexact) - if not jnp.issubdtype(jnp.result_type(t0), jnp.inexact): - raise ValueError( - "`t0`, `t1`, `dt0` must be floating point when specifying `jump_ts`. " - f"Got {jnp.result_type(t0)}." - ) - if not jnp.issubdtype(jnp.result_type(t1), jnp.inexact): - raise ValueError( - "`t0`, `t1`, `dt0` must be floating point when specifying `jump_ts`. " - f"Got {jnp.result_type(t1)}." - ) - jump_ts0 = upcast_or_raise( - self.jump_ts, - t0, - "`PIDController.jump_ts`", - "time (the result type of `t0`, `t1`, `dt0`, `SaveAt(ts=...)` etc.)", - ) - jump_ts1 = upcast_or_raise( - self.jump_ts, - t1, - "`PIDController.jump_ts`", - "time (the result type of `t0`, `t1`, `dt0`, `SaveAt(ts=...)` etc.)", - ) - t0_index = jnp.searchsorted(jump_ts0, t0, side="right") - t1_index = jnp.searchsorted(jump_ts1, t1, side="right") - next_made_jump = t0_index < t1_index - t1 = jnp.where( - next_made_jump, - eqxi.prevbefore(jump_ts1[jnp.minimum(t0_index, len(self.jump_ts) - 1)]), - t1, - ) - return t1, next_made_jump - PIDController.__init__.__doc__ = """**Arguments:** @@ -761,10 +604,6 @@ def _clip_jump_ts( - `force_dtmin`: How to handle the step size hitting the minimum. If `True` then the step size is clipped to `dtmin`. If `False` then the differential equation solve halts with an error. -- `step_ts`: Denotes extra times that must be stepped to. -- `jump_ts`: Denotes extra times that must be stepped to, and at which the vector field - has a known discontinuity. (This is used to force FSAL solvers so re-evaluate the - vector field.) - `factormin`: Minimum amount a step size can be decreased relative to the previous step. - `factormax`: Maximum amount a step size can be increased relative to the previous diff --git a/docs/api/stepsize_controller.md b/docs/api/stepsize_controller.md index 6989c4c1..a59a3d64 100644 --- a/docs/api/stepsize_controller.md +++ b/docs/api/stepsize_controller.md @@ -2,10 +2,29 @@ The list of step size controllers is as follows. The most common cases are fixed step sizes with [`diffrax.ConstantStepSize`][] and adaptive step sizes with [`diffrax.PIDController`][]. -!!! warning - - To perform adaptive stepping with SDEs requires [commutative noise](../usage/how-to-choose-a-solver.md#stochastic-differential-equations). Note that this commutativity condition is not checked. +?? warning "Adaptive SDEs" + + When solving SDEs with an adaptive step controller, then three requirements must be met for the solution to converge to the correct result: + + 1. the Brownian motion must be generated with [`diffrax.VirtualBrownianTree`][]; + 2. the solver must satisfy certain technical conditions (in practice all SDE solvers except [`diffrax.Euler`][] satisfy these), + 3. the SDE must either have [commutative noise](../usage/how-to-choose-a-solver.md#stochastic-differential-equations), or `ClipStepSizeController(..., store_rejected_steps=...)` must be used. + + Conditions 1 and 2 are checked by Diffrax. Condition 3 is not (as there is no easy way to verify commutativity of the noise). + For more details about the convergence of adaptive solutions to SDEs, please refer to + + ```bibtex + @misc{foster2024convergenceadaptiveapproximationsstochastic, + title={On the convergence of adaptive approximations for stochastic differential equations}, + author={James Foster and Andraž Jelinčič}, + year={2024}, + eprint={2311.14201}, + archivePrefix={arXiv}, + primaryClass={math.NA}, + url={https://arxiv.org/abs/2311.14201}, + } + ``` ??? abstract "Abtract base classes" @@ -25,6 +44,7 @@ The list of step size controllers is as follows. The most common cases are fixed members: - rtol - atol + - norm --- @@ -41,3 +61,8 @@ The list of step size controllers is as follows. The most common cases are fixed selection: members: - __init__ + +::: diffrax.ClipStepSizeController + selection: + members: + - __init__ \ No newline at end of file diff --git a/test/test_adaptive_stepsize_controller.py b/test/test_adaptive_stepsize_controller.py index 4cc996c8..68508a2e 100644 --- a/test/test_adaptive_stepsize_controller.py +++ b/test/test_adaptive_stepsize_controller.py @@ -2,22 +2,30 @@ import diffrax import equinox as eqx +import equinox.internal as eqxi import jax import jax.numpy as jnp +import jax.random as jr import jax.tree_util as jtu +import pytest +from diffrax._step_size_controller.clip import _find_idx_with_hint from jaxtyping import Array from .helpers import tree_allclose -def test_step_ts(): +@pytest.mark.parametrize("backwards", [False, True]) +def test_step_ts(backwards): term = diffrax.ODETerm(lambda t, y, args: -0.2 * y) solver = diffrax.Dopri5() t0 = 0 t1 = 5 + if backwards: + t0, t1 = t1, t0 dt0 = None y0 = 1.0 - stepsize_controller = diffrax.PIDController(rtol=1e-4, atol=1e-6, step_ts=[3, 4]) + pid_controller = diffrax.PIDController(rtol=1e-4, atol=1e-6) + stepsize_controller = diffrax.ClipStepSizeController(pid_controller, step_ts=[3, 4]) saveat = diffrax.SaveAt(steps=True) sol = diffrax.diffeqsolve( term, @@ -33,7 +41,8 @@ def test_step_ts(): assert 4 in cast(Array, sol.ts) -def test_jump_ts(): +@pytest.mark.parametrize("backwards", [False, True]) +def test_jump_ts(backwards): # Tests no regression of https://github.com/patrick-kidger/diffrax/issues/58 def vector_field(t, y, args): @@ -45,12 +54,15 @@ def vector_field(t, y, args): solver = diffrax.Dopri5() t0 = 0 t1 = 15 + if backwards: + t0, t1 = t1, t0 dt0 = None y0 = 1.5, 0 saveat = diffrax.SaveAt(steps=True) def run(**kwargs): - stepsize_controller = diffrax.PIDController(rtol=1e-4, atol=1e-6, **kwargs) + pid_controller = diffrax.PIDController(rtol=1e-4, atol=1e-6) + stepsize_controller = diffrax.ClipStepSizeController(pid_controller, **kwargs) return diffrax.diffeqsolve( term, solver, @@ -75,13 +87,90 @@ def run(**kwargs): assert 8 in cast(Array, sol.ts) -def test_backprop(): +@pytest.mark.parametrize("backwards", [False, True]) +def test_revisit_steps(backwards): + t0 = 0.0 + t1 = 5.0 + dt0 = 0.5 + if backwards: + t0, t1 = t1, t0 + dt0 = -dt0 + y0 = 1.0 + drift = diffrax.ODETerm(lambda t, y, args: -0.2 * y) + + def diffusion_vf(t, y, args): + return jnp.ones((), dtype=y.dtype) + + bm = diffrax.VirtualBrownianTree(min(t0, t1), max(t0, t1), 2**-8, (), jr.key(0)) + diffusion = diffrax.ControlTerm(diffusion_vf, bm) + term = diffrax.MultiTerm(drift, diffusion) + solver = diffrax.Heun() + pid_controller = diffrax.PIDController( + rtol=0, atol=1e-3, dtmin=2**-7, pcoeff=0.5, icoeff=0.8 + ) + + rejected_ts_list = [] + + def callback_fun(keep_step, t1): + if not keep_step: + rejected_ts_list.append(t1.item()) + return None + + store_rejected_steps = 10 + stepsize_controller = diffrax.ClipStepSizeController( + pid_controller, + step_ts=[3, 4], + store_rejected_steps=store_rejected_steps, + _callback_on_reject=callback_fun, + ) + saveat = diffrax.SaveAt(steps=True, controller_state=True) + sol = diffrax.diffeqsolve( + term, + solver, + t0, + t1, + dt0, + y0, + stepsize_controller=stepsize_controller, + saveat=saveat, + ) + + assert sol.ts is not None + rejected_ts = jnp.array(rejected_ts_list) + if backwards: + rejected_ts = -rejected_ts + + # there should be many rejected steps, otherwise something went wrong + assert len(rejected_ts) > 10 + # check if all rejected ts are in the array sol.ts + ts = sol.ts[sol.ts != jnp.inf] + if backwards: + ts = ts[::-1] + for t in rejected_ts: + i = jnp.searchsorted(ts, t) + assert ts[i] == t + + assert 3 in cast(Array, sol.ts) + assert 4 in cast(Array, sol.ts) + + # Check that at the end of the run, the rejected stack is empty, + # i.e. rejected_index == store_rejected_steps + assert sol.controller_state is not None + reject_index, _ = sol.controller_state.reject_info + assert reject_index == store_rejected_steps + + +@pytest.mark.parametrize("use_clip", [True, False]) +def test_backprop(use_clip): + t0 = jnp.asarray(0, dtype=jnp.float64) + t1 = jnp.asarray(1, dtype=jnp.float64) + @eqx.filter_jit @eqx.filter_grad def run(ys, controller, state): y0, y1_candidate, y_error = ys _, tprev, tnext, _, state, _ = controller.adapt_step_size( - 0, 1, y0, y1_candidate, None, y_error, 5, state + t0, t1, y0, y1_candidate, None, y_error, 5, state ) with jax.numpy_dtype_promotion("standard"): return tprev + tnext + sum(jnp.sum(x) for x in jtu.tree_leaves(state)) @@ -90,12 +179,16 @@ def run(ys, controller, state): y1_candidate = jnp.array(2.0) term = diffrax.ODETerm(lambda t, y, args: -y) solver = diffrax.Tsit5() - stepsize_controller = diffrax.PIDController(rtol=1e-4, atol=1e-4) - _, state = stepsize_controller.init(term, 0, 1, y0, 0.1, None, solver.func, 5) + controller = diffrax.PIDController(rtol=1e-4, atol=1e-4) + if use_clip: + controller = diffrax.ClipStepSizeController( + controller, step_ts=[0.5], store_rejected_steps=20 + ) + _, state = controller.init(term, t0, t1, y0, 0.1, None, solver.func, 5) for y_error in (jnp.array(0.0), jnp.array(3.0), jnp.array(jnp.inf)): ys = (y0, y1_candidate, y_error) - grads = run(ys, stepsize_controller, state) + grads = run(ys, controller, state) assert not any(jnp.isnan(grad).any() for grad in grads) @@ -113,8 +206,12 @@ def run(t): t1 = 1 dt0 = None y0 = 1.0 - stepsize_controller = diffrax.PIDController( - rtol=1e-8, atol=1e-8, step_ts=t[None] + pid_controller = diffrax.PIDController( + rtol=1e-8, + atol=1e-8, + ) + stepsize_controller = diffrax.ClipStepSizeController( + pid_controller, step_ts=t[None] ) def forcing(s): @@ -139,3 +236,79 @@ def forcing(s): finite_diff = (r(0.5) - r(0.5 - eps)) / eps autodiff = jax.jit(jax.grad(run))(0.5) assert tree_allclose(finite_diff, autodiff) + + +def test_pid_meta(): + ts = jnp.array([3, 4], dtype=jnp.float64) + pid1 = diffrax.PIDController(rtol=1e-4, atol=1e-6) + pid2 = diffrax.PIDController(rtol=1e-4, atol=1e-6, step_ts=ts) # pyright: ignore + pid3 = diffrax.PIDController(rtol=1e-4, atol=1e-6, step_ts=ts, jump_ts=ts) # pyright: ignore + assert not isinstance(pid1, diffrax.ClipStepSizeController) + assert isinstance(pid1, diffrax.PIDController) + assert isinstance(pid2, diffrax.ClipStepSizeController) + assert isinstance(pid3, diffrax.ClipStepSizeController) + assert all(pid2.step_ts == ts) + assert all(pid3.step_ts == ts) + assert all(pid3.jump_ts == ts) + + +def test_nested_clip_wrappers(): + pid = diffrax.PIDController(rtol=0, atol=1.0) + wrap1 = diffrax.ClipStepSizeController(pid, jump_ts=[3.0, 13.0], step_ts=[23.0]) + wrap2 = diffrax.ClipStepSizeController(wrap1, step_ts=[2.0, 13.0], jump_ts=[23.0]) + func = lambda terms, t, y, args: -y + terms = diffrax.ODETerm(lambda t, y, args: -y) + _, state = wrap2.init(terms, -1.0, 0.0, 0.0, 4.0, None, func, 5) + + # test 1 + _, next_t0, next_t1, made_jump, state, _ = wrap2.adapt_step_size( + 0.0, 1.0, 0.0, 0.0, None, 0.0, 5, state + ) + assert next_t0 == 1 + assert next_t1 == 2 + assert not made_jump + _, next_t0, next_t1, made_jump, state, _ = wrap2.adapt_step_size( + next_t0, next_t1, 0.0, 0.0, None, 0.0, 5, state + ) + assert next_t0 == 2 + assert next_t1 == eqxi.prevbefore(jnp.asarray(3.0)) + assert not made_jump + + # test 2 + _, next_t0, next_t1, made_jump, state, _ = wrap2.adapt_step_size( + 10.0, 11.0, 0.0, 0.0, None, 0.0, 5, state + ) + assert next_t0 == 11 + assert next_t1 == eqxi.prevbefore(jnp.asarray(13.0)) + assert not made_jump + _, next_t0, next_t1, made_jump, state, _ = wrap2.adapt_step_size( + next_t0, next_t1, 0.0, 0.0, None, 0.0, 5, state + ) + assert next_t0 == eqxi.nextafter(jnp.asarray(13.0)) + assert next_t1 == eqxi.prevbefore(jnp.asarray(23.0)) + assert made_jump + + # test 3 + _, next_t0, next_t1, made_jump, state, _ = wrap2.adapt_step_size( + 20.0, 21.0, 0.0, 0.0, None, 0.0, 5, state + ) + assert next_t0 == 21 + assert next_t1 == eqxi.prevbefore(jnp.asarray(23.0)) + assert not made_jump + _, next_t0, next_t1, made_jump, state, _ = wrap2.adapt_step_size( + next_t0, next_t1, 0.0, 0.0, None, 0.0, 5, state + ) + assert next_t0 == eqxi.nextafter(jnp.asarray(23.0)) + assert next_t1 > next_t0 + assert made_jump + + +def test_find_idx_with_hint(): + ts = jnp.arange(5.0) + for hint in (0, 2, 3, 5): + idx = _find_idx_with_hint(2.5, ts, hint) + assert idx == 3 + idx = _find_idx_with_hint(2, ts, hint) + assert idx == 3 # not 2; we want the first value *strictly* greater. + idx = _find_idx_with_hint(1.9, ts, hint) + assert idx == 2