diff --git a/README.md b/README.md index 46b2f2b..f28e18b 100644 --- a/README.md +++ b/README.md @@ -18,6 +18,8 @@ Extras for [diffrax][diffrax-link]. +- `subsample`: shrink dense solution storage, preserving the original curve when + accepted steps fit and resampling with cubic Hermite interpolation otherwise. - `DiffEqSolver`: an object-oriented interface to `diffrax.diffeqsolve`. - `VectorizedDenseInterpolation`: a vectorized form of `diffrax.DenseInterpolation` that works on batched results from @@ -330,6 +332,51 @@ Array([[[1. , 0.368], ``` +### `subsample` + +Reduce the dense interpolation buffers after solving: + +```python +from diffraxtra import subsample + +# `soln` was computed with SaveAt(dense=True). +compact = subsample(soln, max_steps=4) +compact.evaluate(0.5) +``` + +A budget at or above the current capacity leaves the interpolation unchanged. If +all accepted steps fit in `max_steps`, the original knots and solver +interpolation are preserved exactly and unused padding is trimmed. Otherwise, +the original dense curve's values and derivatives are sampled at `max_steps + 1` +uniformly spaced times and fitted with cubic Hermite pieces. These match the +original curve at sampling knots, including both endpoints; between knots they +introduce approximation error. Increase `max_steps` and check convergence for +your application. Solver tolerances do not bound this additional error. Use this +for smooth ODEs; jumps and stochastic paths are not preserved by resampling. + +The returned object is still a `diffrax.Solution`. `eqx.tree_at` replaces only +its interpolation; saved `ts`, `ys`, solver state, and statistics retain their +original meaning. In particular, `stats["max_steps"]` describes the original +solve. The original solution is unchanged. Release it when you no longer need +its dense buffers. Saved step arrays (`SaveAt(steps=True)`) are not compacted. + +Both paths use a common packed coefficient buffer with capacity proportional to +`max_steps`, allowing the choice to happen under JIT, including mixed batches +where some solutions trim and others resample. Bytes per interval depend on the +original solver's coefficient layout. `max_steps` must be a static positive +integer under JIT. PyTree states, backward solves, and +`VectorizedDenseInterpolation` are supported. Apply `jax.vmap` to compact an +ordinary batch of dense solutions. + +See [the runnable orbit example](examples/subsample_orbits.py) for a harmonic +orbit and an eccentric Kepler orbit, with storage measurements and interpolation +convergence checks. It uses handwritten potential functions; matplotlib is the +only extra dependency: + +```bash +uv run --with matplotlib python examples/subsample_orbits.py --output /tmp/subsample-orbits.png +``` + ## Citation [![DOI][zenodo-badge]][zenodo-link] diff --git a/examples/__init__.py b/examples/__init__.py new file mode 100644 index 0000000..d5c7e8d --- /dev/null +++ b/examples/__init__.py @@ -0,0 +1 @@ +"""Runnable examples for diffraxtra.""" diff --git a/examples/figures/subsample-orbits.png b/examples/figures/subsample-orbits.png new file mode 100644 index 0000000..3fa96ec Binary files /dev/null and b/examples/figures/subsample-orbits.png differ diff --git a/examples/subsample_orbits.py b/examples/subsample_orbits.py new file mode 100644 index 0000000..2330ab9 --- /dev/null +++ b/examples/subsample_orbits.py @@ -0,0 +1,144 @@ +"""Compare dense storage and interpolation error for two analytic potentials. + +Run with matplotlib and this checkout of diffraxtra installed:: + + uv run --with matplotlib python examples/subsample_orbits.py + +Positions use kpc, times Myr, and velocities kpc/Myr. Potential values are in +kpc²/Myr². Plotted errors measure the additional error from subsampling, +relative to the original dense solution. +""" + +import argparse +import functools as ft +from collections.abc import Callable +from pathlib import Path +from typing import Any + +import diffrax as dfx +import equinox as eqx +import jax +import jax.numpy as jnp +import matplotlib.pyplot as plt +import numpy as np +from matplotlib.axes import Axes + +from diffraxtra import subsample + +Potential = Callable[[jax.Array], jax.Array] +BUDGETS = (8, 16, 32, 64) + + +def harmonic_potential(position: jax.Array, *, omega: float = 0.03) -> jax.Array: + """Return Φ = ω²r²/2, with ω in Myr⁻¹.""" + return 0.5 * omega**2 * jnp.sum(position**2) + + +def kepler_potential(position: jax.Array, *, mu: float = 0.45) -> jax.Array: + """Return Φ = -μ/r, with μ = GM in kpc³/Myr².""" + return -mu / jnp.linalg.norm(position) + + +def vector_field(_t: jax.Array, state: jax.Array, potential: Potential) -> jax.Array: + """Evolve position and velocity using acceleration = -∇Φ.""" + return jnp.concatenate((state[3:], -jax.grad(potential)(state[:3]))) + + +def solve_orbit(potential: Potential, *, speed_fraction: float = 1.0) -> dfx.Solution: + """Integrate from 8 kpc for one circular-orbit period of the potential.""" + position = jnp.array([8.0, 0.0, 0.0]) + circular_speed = jnp.sqrt(position[0] * jax.grad(potential)(position)[0]) + period = float(2 * jnp.pi * position[0] / circular_speed) + velocity = jnp.array([0.0, speed_fraction * circular_speed, 0.0]) + return dfx.diffeqsolve( + dfx.ODETerm(vector_field), + dfx.Dopri8(), + t0=0.0, + t1=period, + dt0=period / 1000, + y0=jnp.concatenate((position, velocity)), + args=potential, + saveat=dfx.SaveAt(dense=True), + stepsize_controller=dfx.PIDController(rtol=1e-10, atol=1e-12), + max_steps=4096, + ) + + +def storage_bytes(tree: Any) -> int: + """Count array storage in the dense interpolation only.""" + return sum(x.nbytes for x in jax.tree.leaves(tree) if eqx.is_array(x)) + + +def compare_solution( + sol: dfx.Solution, budget: int, times: jax.Array, reference: jax.Array +) -> tuple[jax.Array, jax.Array, int]: + """Check a compact solution and return its values, position errors, and bytes.""" + small = eqx.filter_jit(ft.partial(subsample, max_steps=budget))(sol) + values = jax.vmap(small.evaluate)(times) + errors = jnp.linalg.norm(values[:, :3] - reference[:, :3], axis=-1) + size = storage_bytes(small.interpolation) + np.testing.assert_array_less(size, storage_bytes(sol.interpolation)) + np.testing.assert_allclose( + values[jnp.array([0, -1])], reference[jnp.array([0, -1])], rtol=0, atol=1e-12 + ) + if budget >= int(sol.stats["num_accepted_steps"]): + roundoff = 32 * np.finfo(reference.dtype).eps + np.testing.assert_allclose( + values, + reference, + rtol=roundoff, + atol=roundoff * float(jnp.max(jnp.abs(reference))), + ) + return values, errors, size + + +def plot_orbit(axes: tuple[Axes, Axes], name: str, sol: dfx.Solution) -> None: + """Plot an orbit and its subsampling errors, checking convergence.""" + orbit_ax, error_ax = axes + times = jnp.linspace(sol.t0, sol.t1, 2001) + reference = jax.vmap(sol.evaluate)(times) + original_bytes = storage_bytes(sol.interpolation) + accepted = int(sol.stats["num_accepted_steps"]) + orbit_ax.plot(reference[:, 0], reference[:, 1], label="original", color="black") + print(f"{name}: {accepted} accepted steps, {original_bytes:,} dense bytes") # noqa: T201 + previous_error = np.inf + for budget in BUDGETS: + values, errors, size = compare_solution(sol, budget, times, reference) + error = float(jnp.max(errors)) + np.testing.assert_array_less(error, previous_error + 1e-12) + previous_error = error + label = f"{budget} intervals" + orbit_ax.plot(values[:, 0], values[:, 1], "--", label=label) + error_ax.semilogy(times, jnp.maximum(errors, 1e-15), label=label) + print( # noqa: T201 + f" {budget:2d} intervals: {size:7,d} bytes " + f"({original_bytes / size:.1f}x smaller), " + f"max position error {error:.3e} kpc" + ) + orbit_ax.set(title=name, xlabel="x [kpc]", ylabel="y [kpc]", aspect="equal") + error_ax.set( + title=f"{name}: error relative to dense reference", + xlabel="t [Myr]", + ylabel="position error [kpc]", + ) + for ax in axes: + ax.legend(fontsize="small") + + +def main(output: Path) -> None: + """Render the comparisons using a circular harmonic and eccentric Kepler orbit.""" + jax.config.update("jax_enable_x64", True) # noqa: FBT003 + fig, axes = plt.subplots(2, 2, figsize=(10, 8), constrained_layout=True) + plot_orbit(tuple(axes[0]), "Harmonic oscillator", solve_orbit(harmonic_potential)) + plot_orbit( + tuple(axes[1]), "Kepler", solve_orbit(kepler_potential, speed_fraction=0.8) + ) + output.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(output, dpi=160) + plt.close(fig) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, default=Path("subsample-orbits.png")) + main(parser.parse_args().output) diff --git a/src/diffraxtra/__init__.py b/src/diffraxtra/__init__.py index 9e62292..8e39066 100644 --- a/src/diffraxtra/__init__.py +++ b/src/diffraxtra/__init__.py @@ -5,6 +5,7 @@ "DiffEqSolver", "AbstractVectorizedDenseInterpolation", "VectorizedDenseInterpolation", + "subsample", ] from ._src import ( @@ -13,6 +14,7 @@ DiffEqSolver, VectorizedDenseInterpolation, ) +from ._src.subsample import subsample from ._version import ( # noqa: F401 version as __version__, version_tuple as __version_tuple__, diff --git a/src/diffraxtra/_src/subsample.py b/src/diffraxtra/_src/subsample.py new file mode 100644 index 0000000..e4fe773 --- /dev/null +++ b/src/diffraxtra/_src/subsample.py @@ -0,0 +1,210 @@ +"""Reduce the storage of a dense solution by cubic Hermite resampling.""" + +__all__ = ["subsample"] + +import functools as ft +import operator +from collections.abc import Callable +from typing import Any, cast +from typing_extensions import override + +import diffrax as dfx +import equinox as eqx +import jax +import jax.numpy as jnp +from jax.flatten_util import ravel_pytree + +from .interp import VectorizedDenseInterpolation + +# Keep typing compatible with JAX releases where ravel_pytree is untyped. +_ravel: Callable[[Any], tuple[jax.Array, Callable[[jax.Array], Any]]] = ravel_pytree + + +LocalFactory = Callable[[Any, Any, jax.Array], dfx.AbstractLocalInterpolation] + + +def _packed_factory( + cls: Callable[..., dfx.AbstractLocalInterpolation], infos: Any +) -> tuple[jax.Array, LocalFactory]: + """Describe a coefficient layout without retaining the dense buffers.""" + flat, unravel = _ravel(infos) + size = flat.size + + def factory(t0: Any, t1: Any, data: jax.Array) -> dfx.AbstractLocalInterpolation: + return cls(t0=t0, t1=t1, **unravel(data[:size])) + + return flat, factory + + +class _LocalInterpolation(dfx.AbstractLocalInterpolation): + """Evaluate packed coefficients in their original or Hermite format.""" + + t0: Any # type: ignore[misc] + t1: Any # type: ignore[misc] + data: jax.Array + exact: jax.Array + original_factory: LocalFactory = eqx.field(static=True) + hermite_factory: LocalFactory = eqx.field(static=True) + + @override + def evaluate(self, t0: Any, t1: Any = None, left: bool = True) -> Any: + def evaluate(factory: LocalFactory) -> Any: + return factory(self.t0, self.t1, self.data).evaluate(t0, t1, left=left) + + return jax.lax.cond( + self.exact, + lambda: evaluate(self.original_factory), + lambda: evaluate(self.hermite_factory), + ) + + +def _subsample_interpolation( + interp: dfx.DenseInterpolation, max_steps: int +) -> dfx.DenseInterpolation: + # The two interpolation layouts each require shape and reconstruction metadata. + # pylint: disable=too-many-locals + if interp.ts.ndim != 1: + msg = ( + "Use jax.vmap(subsample) or VectorizedDenseInterpolation " + "for batched solutions." + ) + raise ValueError(msg) + if max_steps >= interp.ts.shape[0] - 1: + return interp + + # Both paths have the same static layout, so mixed batches work under JIT. + exact = interp.ts_size <= max_steps + 1 + original_example = jax.tree.map( + lambda x: jnp.zeros(x.shape[1:], x.dtype), interp.infos + ) + y_example = jax.tree.map(jnp.zeros_like, interp.y0_if_trivial) + hermite_example = dict.fromkeys(("y0", "y1", "k0", "k1"), y_example) + flat_original, original_factory = _packed_factory( + interp.interpolation_cls, original_example + ) + flat_hermite, hermite_factory = _packed_factory( + dfx.ThirdOrderHermitePolynomialInterpolation, hermite_example + ) + width = max(flat_original.size, flat_hermite.size) + dtype = jnp.result_type(flat_original, flat_hermite) + + def pack(infos: Any) -> jax.Array: + data = jax.vmap(lambda info: _ravel(info)[0])(infos) + return jnp.pad(data.astype(dtype), ((0, 0), (0, width - data.shape[1]))) + + def resample() -> tuple[jax.Array, Any]: + # Diffrax stores increasing internal times even for backward integration. + ts = jnp.linspace(interp.ts[0], interp.ts[interp.ts_size - 1], max_steps + 1) + times = ts * interp.direction + ys = jax.vmap(interp.evaluate)(times) + dys = jax.vmap(interp.derivative)(times) + + def increments(dy: jax.Array) -> jax.Array: + dt = jnp.diff(ts).reshape((max_steps,) + (1,) * (dy.ndim - 1)) + return dy * dt * interp.direction + + infos = { + "y0": jax.tree.map(lambda y: y[:-1], ys), + "y1": jax.tree.map(lambda y: y[1:], ys), + "k0": jax.tree.map(lambda dy: increments(dy[:-1]), dys), + "k1": jax.tree.map(lambda dy: increments(dy[1:]), dys), + } + return ts, pack(infos) + + def trim() -> tuple[jax.Array, Any]: + original = jax.tree.map(lambda x: x[:max_steps], interp.infos) + return interp.ts[: max_steps + 1], pack(original) + + ts, data = jax.lax.cond(exact, trim, resample) + return dfx.DenseInterpolation( + ts=ts, + ts_size=jnp.where(exact, interp.ts_size, max_steps + 1), + infos={"data": data, "exact": jnp.full((max_steps,), exact)}, + interpolation_cls=ft.partial( + _LocalInterpolation, + original_factory=original_factory, + hermite_factory=hermite_factory, + ), + direction=interp.direction, + t0_if_trivial=interp.t0_if_trivial, + y0_if_trivial=interp.y0_if_trivial, + ) + + +def subsample(soln: dfx.Solution, /, *, max_steps: int) -> dfx.Solution: + """Return a dense solution with storage for ``max_steps`` intervals. + + If the accepted steps fit, only unused padding is removed and the original + interpolation is preserved exactly. Otherwise, values and time derivatives + of the original dense curve are sampled at + ``max_steps + 1`` uniformly spaced times over its actual integration interval. + These define cubic Hermite pieces; simply slicing the solver's original + coefficients would give incorrect interpolation across skipped steps. + + A budget at or above the original capacity leaves the interpolation unchanged. + When resampling, the new curve agrees at the sampling knots (both endpoints). + Between knots it is an approximation, with no inherited solver tolerance + guarantee. Increase ``max_steps`` and check convergence for your application. + Intended for smooth ODE solutions; jumps and stochastic paths are not preserved. + + Only the interpolation is replaced, using ``equinox.tree_at``. Saved ``ts``, + ``ys``, statistics, and solver state retain their original meaning and storage. + The input is unchanged. The returned interpolation stores no reference to the + original dense buffers. This does not change the integration's step budget. + + Supports PyTree states, backward integration, and + :class:`~diffraxtra.VectorizedDenseInterpolation`. For ordinary batched dense + solutions, apply with ``jax.vmap``. Under JIT, ``max_steps`` must be static. + + Args: + soln: A solution saved with ``diffrax.SaveAt(dense=True)``. + max_steps: Positive integer number of intervals in the new interpolation. + + Returns: + A new ``diffrax.Solution`` containing the smaller dense interpolation. + + Examples: + >>> import diffrax as dfx + >>> from diffraxtra import subsample + >>> sol = dfx.diffeqsolve( + ... dfx.ODETerm(lambda t, y, args: -y), dfx.Dopri5(), + ... t0=0., t1=1., dt0=0.01, y0=1., + ... saveat=dfx.SaveAt(dense=True), max_steps=256) + >>> compact = subsample(sol, max_steps=16) + >>> compact.interpolation.ts.shape + (17,) + >>> bool(abs(compact.evaluate(0.5) - sol.evaluate(0.5)) < 1e-6) + True + >>> trimmed = subsample(sol, max_steps=128) + >>> bool(trimmed.evaluate(0.5) == sol.evaluate(0.5)) + True + + """ + if isinstance(max_steps, bool): + msg = "max_steps must be a positive integer." + raise TypeError(msg) + try: + max_steps = operator.index(max_steps) + except TypeError as error: + msg = "max_steps must be a static positive integer." + raise TypeError(msg) from error + if max_steps < 1: + msg = "max_steps must be positive." + raise ValueError(msg) + + interp = cast("Any", soln.interpolation) + if isinstance(interp, VectorizedDenseInterpolation): + scalar = jax.vmap(lambda x: _subsample_interpolation(x, max_steps))( + interp.scalar_interpolation + ) + new_interp = eqx.tree_at(lambda x: x.scalar_interpolation, interp, scalar) + elif isinstance(interp, dfx.DenseInterpolation): + new_interp = _subsample_interpolation(interp, max_steps) + elif interp is None: + msg = "subsample requires a dense solution; use SaveAt(dense=True)." + raise ValueError(msg) + else: + msg = "subsample requires a diffrax.DenseInterpolation." + raise TypeError(msg) + result: dfx.Solution = eqx.tree_at(lambda x: x.interpolation, soln, new_interp) + return result diff --git a/tests/test_subsample.py b/tests/test_subsample.py new file mode 100644 index 0000000..39ba82e --- /dev/null +++ b/tests/test_subsample.py @@ -0,0 +1,284 @@ +"""Regression tests for reducing dense solution storage.""" + +import diffrax as dfx +import equinox as eqx +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from diffraxtra import VectorizedDenseInterpolation, subsample + + +def solve(t0=0.0, t1=3.0, y0=1.0, *, dense=True): + """Solve a smooth equation with deliberately oversized dense buffers.""" + return dfx.diffeqsolve( + dfx.ODETerm(lambda _t, y, _args: jax.tree.map(lambda x: -x, y)), + dfx.Dopri5(), + t0, + t1, + jnp.where(t1 >= t0, 0.01, -0.01), + y0, + saveat=dfx.SaveAt(t0=True, t1=True, dense=dense), + max_steps=512, + ) + + +@pytest.mark.parametrize("backward", [False, True]) +def test_dense_accuracy_and_storage(backward): + """Keep endpoints and metadata while actually reducing interpolation storage.""" + sol = solve(3.0, 0.0) if backward else solve() + small = subsample(sol, max_steps=32) + assert isinstance(small, dfx.Solution) + assert small is not sol + assert small.interpolation.ts.shape == (33,) + assert int(small.interpolation.ts_size) == 33 + assert sol.interpolation.ts.shape == (513,) + assert small.ts is sol.ts + assert small.ys is sol.ys + assert eqx.tree_equal(small.stats, sol.stats) + assert small.result == sol.result + assert small.interpolation.t0 == sol.interpolation.t0 + assert small.interpolation.t1 == sol.interpolation.t1 + times = jnp.linspace(sol.t0, sol.t1, 503) + np.testing.assert_allclose( + jax.vmap(small.evaluate)(times), jax.vmap(sol.evaluate)(times), rtol=3e-6 + ) + for t in [sol.t0, sol.t1, 1.5]: + np.testing.assert_allclose(small.evaluate(t), sol.evaluate(t), rtol=1e-12) + np.testing.assert_allclose(small.derivative(t), sol.derivative(t), rtol=1e-10) + np.testing.assert_allclose( + small.evaluate(0.2, 0.8), small.evaluate(0.8) - small.evaluate(0.2) + ) + assert jnp.isnan(small.evaluate(-1.0)) + assert all(x.shape[0] == 32 for x in jax.tree.leaves(small.interpolation.infos)) + + +def test_pytree_jit_and_grad(): + """Resampling supports structured states, compilation, and autodiff.""" + sol = solve(y0={"a": jnp.array([1.0, 2.0]), "b": (jnp.array(3.0),)}) + small = eqx.filter_jit(lambda s: subsample(s, max_steps=16))(sol) + for actual, expected in zip( + jax.tree.leaves(small.evaluate(0.7)), + jax.tree.leaves(sol.evaluate(0.7)), + strict=True, + ): + np.testing.assert_allclose(actual, expected, rtol=1e-5) + derivative = jax.grad(lambda t: small.evaluate(t)["b"][0])(0.7) + np.testing.assert_allclose(derivative, -3 * jnp.exp(-0.7), rtol=1e-4) + + +def test_vectorized_batch(): + """Flattened vectorized interpolation retains its batch and time semantics.""" + sol = jax.vmap(solve)(jnp.zeros((2,)), jnp.full((2,), 3.0), jnp.array([1.0, 2.0])) + sol = VectorizedDenseInterpolation.apply_to_solution(sol) + small = eqx.filter_jit(lambda s: subsample(s, max_steps=32))(sol) + assert isinstance(small.interpolation, VectorizedDenseInterpolation) + assert small.interpolation.ts.shape == (2, 33) + times = jnp.linspace(0, 3, 20).reshape(4, 5) + assert small.evaluate(times).shape == (2, 4, 5) + np.testing.assert_allclose(small.evaluate(times), sol.evaluate(times), rtol=3e-6) + + +@pytest.mark.parametrize("size", [0, -1, 1.5, True]) +def test_invalid_size(size): + """A fixed positive integer is required for output shapes.""" + with pytest.raises((TypeError, ValueError), match="max_steps"): + subsample(solve(), max_steps=size) + + +def test_requires_dense(): + """Missing dense output is diagnosed explicitly.""" + with pytest.raises(ValueError, match="dense"): + subsample(solve(dense=False), max_steps=10) + + +def test_trivial_interval(): + """A zero-duration solve retains its sole defined value.""" + sol = solve(1.0, 1.0) + small = subsample(sol, max_steps=4) + np.testing.assert_array_equal(small.evaluate(1.0), sol.evaluate(1.0)) + assert jnp.isnan(small.evaluate(0.0)) + + +def test_vmap_subsample(): + """Resampling can be vmapped over ordinary dense solutions.""" + sol = jax.vmap(lambda y: subsample(solve(y0=y), max_steps=8))(jnp.array([1.0, 2.0])) + assert sol.interpolation.ts.shape == (2, 9) + + +@pytest.mark.parametrize("solver", [dfx.Tsit5(), dfx.Dopri8(), dfx.Euler()]) +def test_adaptive_and_other_interpolants(solver): + """Resampling reads each solver's curve rather than reusing its coefficients.""" + sol = dfx.diffeqsolve( + dfx.ODETerm(lambda _t, y, _args: -y), + solver, + 0.0, + 2.0, + 0.01, + 1.0, + saveat=dfx.SaveAt(dense=True), + stepsize_controller=( + dfx.ConstantStepSize() + if isinstance(solver, dfx.Euler) + else dfx.PIDController(rtol=1e-5, atol=1e-7) + ), + ) + small = subsample(sol, max_steps=64) + times = jnp.linspace(0, 2, 201) + np.testing.assert_allclose( + jax.vmap(small.evaluate)(times), jax.vmap(sol.evaluate)(times), atol=3e-5 + ) + + +def test_convergence(): + """Cubic resampling error decreases as the sampling grid is refined.""" + sol = solve() + times = jnp.linspace(0, 3, 501) + reference = jax.vmap(sol.evaluate)(times) + errors = [ + jnp.max( + jnp.abs(jax.vmap(subsample(sol, max_steps=n).evaluate)(times) - reference) + ) + for n in [8, 16, 32] + ] + assert errors[1] < errors[0] / 10 + assert errors[2] < errors[1] / 10 + + +def test_event_uses_actual_interval(): + """An early termination samples only the interval actually integrated.""" + sol = dfx.diffeqsolve( + dfx.ODETerm(lambda _t, y, _args: -y), + dfx.Dopri5(), + 0.0, + 10.0, + 0.1, + 1.0, + saveat=dfx.SaveAt(dense=True), + event=dfx.Event(lambda _t, y, _args, **_kwargs: y < 0.5), + ) + small = subsample(sol, max_steps=16) + end = sol.interpolation.t1 + assert end < sol.t1 + assert small.interpolation.t1 == end + np.testing.assert_allclose(small.evaluate(end), sol.evaluate(end), rtol=1e-12) + + +def test_one_interval(): + """The smallest valid output still interpolates the whole domain.""" + sol = solve() + small = subsample(sol, max_steps=1) + assert small.interpolation.ts.shape == (2,) + for t in [0.0, 3.0]: + np.testing.assert_allclose(small.evaluate(t), sol.evaluate(t), rtol=1e-12) + + +def test_trim_is_exact(): + """A large enough budget preserves original knots, coefficients, and curve.""" + sol = solve() + small = eqx.filter_jit(lambda s: subsample(s, max_steps=320))(sol) + assert small.interpolation.ts.shape == (321,) + assert small.interpolation.ts_size == sol.interpolation.ts_size + np.testing.assert_array_equal(small.interpolation.ts, sol.interpolation.ts[:321]) + times = jnp.linspace(0, 3, 401) + # Coefficients are bitwise unchanged; compiler fusion may round evaluations + # differently on older JAX versions when the same polynomial is wrapped. + coefficients = np.concatenate( + [ + np.asarray(x[:320]).reshape(320, -1) + for x in jax.tree.leaves(sol.interpolation.infos) + ], + axis=1, + ) + np.testing.assert_array_equal(small.interpolation.infos["data"], coefficients) + np.testing.assert_allclose( + jax.vmap(small.evaluate)(times), + jax.vmap(sol.evaluate)(times), + rtol=32 * np.finfo(float).eps, + atol=0, + ) + + +def test_mixed_trim_and_resample_batch(): + """One compiled batch can trim one member and resample another.""" + + @jax.vmap + def solve_variable(end): + sol = solve(t1=end) + return subsample(sol, max_steps=32) + + small = eqx.filter_jit(solve_variable)(jnp.array([0.1, 3.0])) + assert int(small.interpolation.ts_size[0]) < 33 + assert int(small.interpolation.ts_size[1]) == 33 + values = jax.vmap(lambda s: s.evaluate(0.05))(small) + np.testing.assert_allclose(values, jnp.exp(-0.05), rtol=3e-6) + + +@pytest.mark.parametrize("budget", [32, 400]) +def test_storage_bytes(budget): + """A compact interpolation retains neither a reference nor a second payload.""" + sol = solve() + small = subsample(sol, max_steps=budget) + original = jax.tree.leaves(sol.interpolation) + compact = jax.tree.leaves(small.interpolation) + assert sum(x.nbytes for x in compact) < sum(x.nbytes for x in original) + assert not any(x is y for x in compact for y in original if y.ndim > 0) + + +def test_repeated_subsample(): + """A previously compacted curve can be compacted again.""" + small = subsample(subsample(solve(), max_steps=320), max_steps=32) + np.testing.assert_allclose(small.evaluate(0.7), jnp.exp(-0.7), rtol=3e-6) + + +def test_gradient_through_resampling(): + """Differentiate compaction with respect to the input curve's coefficients.""" + sol = solve() + # Padding is unused, but differentiating scale * NaN would contaminate grads. + infos = jax.tree.map( + lambda x: jnp.where(jnp.isfinite(x), x, 0), sol.interpolation.infos + ) + + def evaluate(scale): + scaled = eqx.tree_at( + lambda s: (s.interpolation.infos, s.interpolation.y0_if_trivial), + sol, + ( + jax.tree.map(lambda x: scale * x, infos), + scale * sol.interpolation.y0_if_trivial, + ), + ) + return subsample(scaled, max_steps=16).evaluate(0.7) + + np.testing.assert_allclose(jax.grad(evaluate)(1.0), jnp.exp(-0.7), rtol=1e-5) + + +def test_vectorized_mixed_modes(): + """The vectorized wrapper also supports mixed exact and approximate members.""" + sol = jax.vmap(lambda end: solve(t1=end))(jnp.array([0.1, 3.0])) + sol = VectorizedDenseInterpolation.apply_to_solution(sol) + small = eqx.filter_jit(lambda s: subsample(s, max_steps=32))(sol) + assert int(small.interpolation.ts_size[0]) < 33 + assert int(small.interpolation.ts_size[1]) == 33 + np.testing.assert_allclose(small.evaluate(0.05), sol.evaluate(0.05), rtol=3e-6) + + +@pytest.mark.filterwarnings("ignore:Complex dtype support:UserWarning") +@pytest.mark.parametrize("budget", [16, 320]) +def test_float32_and_complex(budget): + """Packing preserves state dtypes, including complex-valued states.""" + for dtype in [jnp.float32, jnp.complex64]: + sol = solve(y0=jnp.asarray(1.0, dtype=dtype)) + small = subsample(sol, max_steps=budget) + actual = small.evaluate(0.7) + assert actual.dtype == sol.evaluate(0.7).dtype + np.testing.assert_allclose(actual, sol.evaluate(0.7), rtol=1e-5) + + +def test_larger_budget_is_noop(): + """A request above the original capacity does not grow the buffers.""" + sol = solve() + small = subsample(sol, max_steps=1024) + assert small.interpolation.ts.shape == sol.interpolation.ts.shape + np.testing.assert_array_equal(small.evaluate(0.7), sol.evaluate(0.7))