diff --git a/CHANGELOG.md b/CHANGELOG.md index ab5ce07..000f8a9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,14 @@ versioning while the project remains pre-1.0. parameter/state reuse and compatible manual meshes. A structured analytic example demonstrates accuracy-driven refresh outside differentiation (#205). +- A public synthetic hierarchical nonlinear ODE fixture defines shared/local + rates and local initial conditions, independent/coupled and nonstiff/stiff + regimes, deterministic observations and a hierarchical objective. Independent + cell equations and SciPy trajectories qualify NumPy/JAX frozen replay, external + public-step scans, forward/reverse gradients, Hessians and a regularized + Hessian log-determinant gradient. Small correctness evidence is mandatory; + larger comparison measurements remain a separate follow-up (#204). + ### Fixed - Namespace-aware advection and diffusion builders honor a floating reference diff --git a/benchmarks/__init__.py b/benchmarks/__init__.py new file mode 100644 index 0000000..4aa2da8 --- /dev/null +++ b/benchmarks/__init__.py @@ -0,0 +1 @@ +"""Synthetic numerical evidence and opt-in measurement workloads.""" diff --git a/benchmarks/hierarchical_ode.py b/benchmarks/hierarchical_ode.py new file mode 100644 index 0000000..fd5d393 --- /dev/null +++ b/benchmarks/hierarchical_ode.py @@ -0,0 +1,445 @@ +"""Generic hierarchical nonlinear two-pool model with independent references. + +Each group carries one or more local parameter/initial-condition cells. Shared +positive gain/recovery rates combine with group-local log rate offsets. The +coupled variant uses the group mean of the first pool in its forward flux. +Each cell conserves its two-pool total. This synthetic workload is numerical +inference evidence, not a calibrated model of a particular scientific system. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import TYPE_CHECKING, cast + +import numpy as np +from scipy.integrate import solve_ivp + +from op_engine import Array, CoreSolver, ModelCore, array_namespace +from op_engine.core_solver import AdaptiveConfig, AdaptiveStepSchedule, RunConfig +from op_engine.model_core import ModelCoreOptions + +if TYPE_CHECKING: + from collections.abc import Callable + + +@dataclass(frozen=True, slots=True) +class HierarchicalProblem: + """Static layout and regime for a reproducible hierarchical ODE workload.""" + + groups: int = 2 + local_dimension: int = 1 + coupled: bool = False + stiff: bool = False + + def __post_init__(self) -> None: + """Validate positive fixture dimensions. + + Raises: + ValueError: If a dimension is not a positive integer. + """ + for name in ("groups", "local_dimension"): + value = getattr(self, name) + if not isinstance(value, int) or isinstance(value, bool) or value < 1: + msg = f"{name} must be a positive integer" + raise ValueError(msg) + + @property + def cells(self) -> int: + """Number of local rate/initial-condition pairs.""" + return self.groups * self.local_dimension + + @property + def parameter_count(self) -> int: + """Two shared log rates plus local rate offsets and initial values.""" + return 2 + 2 * self.cells + + @property + def state_shape(self) -> tuple[int, int, int]: + """Pool/group/local-cell axis order.""" + return (2, self.groups, self.local_dimension) + + @property + def model_id(self) -> str: + """Structural identity excluding fitted parameter values.""" + return ( + f"hierarchical-two-pool:v1:groups={self.groups}:" + f"local={self.local_dimension}:coupled={self.coupled}" + ) + + +def nominal_parameters(problem: HierarchicalProblem) -> np.ndarray: + """Return deterministic concrete shared/local reference parameters. + + Returns: + Parameter vector in shared gain/recovery, local rates, local IC order. + """ + return np.concatenate(( + [math.log(0.4), math.log(50.0 if problem.stiff else 0.15)], + np.linspace(-0.2, 0.25, problem.cells), + np.linspace(-0.3, 0.4, problem.cells), + )) + + +def initial_state(problem: HierarchicalProblem, parameters: Array) -> Array: + """Construct differentiable group-local initial values with unit totals. + + Returns: + Native state tensor with one positive two-pool pair per local cell. + """ + xp = array_namespace(parameters) + first = 0.55 + 0.1 * xp.tanh( + xp.reshape( + parameters[2 + problem.cells :], (problem.groups, problem.local_dimension) + ) + ) + return cast("Array", xp.stack((first, 1 - first))) + + +def rhs( + problem: HierarchicalProblem, parameters: Array +) -> Callable[[object, Array], Array]: + """Bind native nonlinear equations without host parameter conversion. + + Returns: + Array-API RHS with shared/global and group-local dynamic inputs. + """ + xp = array_namespace(parameters) + gain, recovery = xp.exp(parameters[:2]) + local = xp.exp( + xp.reshape( + parameters[2 : 2 + problem.cells], (problem.groups, problem.local_dimension) + ) + ) + + def evaluate(_time: object, state: Array) -> Array: + first, second = state + driver = first + (0.15 * xp.mean(first, axis=0) if problem.coupled else 0.0) + forward = gain * local * first * driver / (1 + first) + net = forward - recovery * second + return cast("Array", xp.stack((-net, net))) + + return evaluate + + +def reference_initial( + problem: HierarchicalProblem, parameters: np.ndarray +) -> np.ndarray: + """Evaluate initial cells independently using scalar standard-library math. + + Returns: + Concrete reference state. + """ + result = np.empty(problem.state_shape) + for group in range(problem.groups): + for cell in range(problem.local_dimension): + index = group * problem.local_dimension + cell + first = 0.55 + 0.1 * math.tanh(float(parameters[2 + problem.cells + index])) + result[:, group, cell] = [first, 1 - first] + return result + + +def reference_rhs( + problem: HierarchicalProblem, parameters: np.ndarray, state: np.ndarray +) -> np.ndarray: + """Independent explicit cell equations, with no generated/native RHS calls. + + Returns: + Concrete two-pool derivatives. + """ + result = np.empty_like(state) + gain, recovery = (math.exp(float(value)) for value in parameters[:2]) + for group in range(problem.groups): + for cell in range(problem.local_dimension): + first, second = (float(value) for value in state[:, group, cell]) + mean_first = 0.0 + if problem.coupled: + mean_first = ( + sum(float(state[0, g, cell]) for g in range(problem.groups)) + / problem.groups + ) + local = math.exp( + float(parameters[2 + group * problem.local_dimension + cell]) + ) + forward = gain * local * first * (first + 0.15 * mean_first) / (1 + first) + result[:, group, cell] = [ + recovery * second - forward, + forward - recovery * second, + ] + return result + + +def reference_trajectory( + problem: HierarchicalProblem, parameters: np.ndarray, times: np.ndarray +) -> np.ndarray: + """Compute a tight-tolerance independent SciPy trajectory. + + Returns: + Reference output states including the initial point. + + Raises: + RuntimeError: If reference integration fails. + """ + initial = reference_initial(problem, parameters) + solution = solve_ivp( + lambda _time, flat: reference_rhs( + problem, parameters, flat.reshape(problem.state_shape) + ).ravel(), + (float(times[0]), float(times[-1])), + initial.ravel(), + t_eval=times, + method="DOP853", + rtol=2e-12, + atol=2e-14, + ) + if not solution.success: + raise RuntimeError(solution.message) + return solution.y.T.reshape((len(times), *problem.state_shape)) + + +def record_schedule( + problem: HierarchicalProblem, parameters: np.ndarray, times: np.ndarray +) -> tuple[AdaptiveStepSchedule, RunConfig]: + """Discover a nominal accepted mesh at an eager boundary. + + Returns: + Accepted schedule and its structural/reference-labeled configuration. + + Raises: + RuntimeError: If discovery fails to produce a mesh. + """ + core = ModelCore( + 2, + problem.groups, + times, + options=ModelCoreOptions( + other_axes=(problem.local_dimension,), axis_names=("pool", "group", "local") + ), + ) + core.set_initial_state(initial_state(problem, parameters)) + config = RunConfig( + method="dopri5", + adaptive=True, + adaptive_cfg=AdaptiveConfig(rtol=2e-8, atol=2e-10, dt_init=0.02), + replay_model_id=problem.model_id, + replay_reference_regime="nominal_parameters()" + + (":stiff" if problem.stiff else ":nonstiff"), + ) + solver = CoreSolver(core) + solver.run(rhs(problem, parameters), config=config) + schedule = solver.last_adaptive_schedule + if schedule is None: + msg = "Adaptive discovery did not produce a schedule" + raise RuntimeError(msg) + return schedule, config + + +def replay_trajectory( + problem: HierarchicalProblem, + parameters: Array, + schedule: AdaptiveStepSchedule, + config: RunConfig, + *, + external: bool = False, +) -> Array: + """Differentiate a frozen mesh through built-in or public external steps. + + Returns: + Native output trajectory, including the initial state. + """ + import jax # noqa: PLC0415 + + xp = array_namespace(parameters) + initial = initial_state(problem, parameters) + core = ModelCore( + 2, + problem.groups, + np.asarray(schedule.output_times), + options=ModelCoreOptions( + other_axes=(problem.local_dimension,), + axis_names=("pool", "group", "local"), + ), + ) + solver = CoreSolver(core) + function = rhs(problem, parameters) + if not external: + from dataclasses import replace # noqa: PLC0415 + + core.set_initial_state(initial) + solver.replay_adaptive_schedule( + function, + schedule, + config=replace(config, replay_loop="scan", replay_checkpoint=True), + ) + return cast("Array", core.state_array) + starts, sizes, selections = _flatten_mesh(schedule) + first, fsal = solver.fixed_explicit_step( + function, method="dopri5", t=starts[0], dt=sizes[0], y=initial + ) + + def advance( + carry: tuple[Array, Array | None], step: tuple[object, object] + ) -> tuple[tuple[Array, Array | None], Array]: + state, cached = carry + time, dt = step + result, stage = solver.fixed_explicit_step( + function, method="dopri5", t=time, dt=dt, y=state, first_stage=cached + ) + return (result, stage), result + + _, tail = jax.lax.scan( + jax.checkpoint(advance), + (first, fsal), + (xp.asarray(starts[1:]), xp.asarray(sizes[1:])), + ) + complete = xp.concat((initial[None], first[None], tail), axis=0) + return cast("Array", complete[xp.asarray(selections)]) + + +def numpy_replay_trajectory( + problem: HierarchicalProblem, + parameters: np.ndarray, + schedule: AdaptiveStepSchedule, + config: RunConfig, +) -> np.ndarray: + """Replay the identical mesh eagerly with NumPy for value parity. + + Returns: + Output trajectory in NumPy. + """ + core = ModelCore( + 2, + problem.groups, + np.asarray(schedule.output_times), + options=ModelCoreOptions( + other_axes=(problem.local_dimension,), + axis_names=("pool", "group", "local"), + ), + ) + core.set_initial_state(initial_state(problem, parameters)) + CoreSolver(core).replay_adaptive_schedule( + rhs(problem, parameters), schedule, config=config + ) + return np.asarray(core.state_array) + + +def _flatten_mesh( + schedule: AdaptiveStepSchedule, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + starts: list[float] = [] + sizes: list[float] = [] + selections = [0] + for start, interval in zip( + schedule.output_times[:-1], schedule.step_sizes, strict=True + ): + time = start + for size in interval: + starts.append(time) + sizes.append(size) + time += size + selections.append(len(sizes)) + return np.asarray(starts), np.asarray(sizes), np.asarray(selections) + + +def refined_schedule(schedule: AdaptiveStepSchedule) -> AdaptiveStepSchedule: + """Halve every accepted step while retaining the verified structural record. + + Returns: + Valid refined mesh with the same output times and provenance. + """ + return AdaptiveStepSchedule( + schedule.output_times, + tuple( + tuple(value for size in interval for value in (size / 2, size / 2)) + for interval in schedule.step_sizes + ), + provenance=schedule.provenance, + ) + + +def observations(problem: HierarchicalProblem, times: np.ndarray) -> np.ndarray: + """Deterministic reference observations with a nonzero residual pattern. + + Returns: + First-pool observations at noninitial output times. + """ + trajectory = reference_trajectory(problem, nominal_parameters(problem), times) + signal = trajectory[1:, 0] + return signal + 0.025 * 0.2 * np.sin( + np.arange(signal.size).reshape(signal.shape) + 0.5 + ) + + +def objective( + problem: HierarchicalProblem, parameters: Array, trajectory: Array, data: Array +) -> Array: + """Combine group likelihoods and shared/local hierarchical prior regularization. + + Returns: + Scalar negative log objective up to parameter-independent constants. + """ + xp = array_namespace(parameters) + residual = (trajectory[1:, 0] - data) / 0.025 + global_center = xp.asarray(nominal_parameters(problem)[:2], dtype=parameters.dtype) + shared = xp.sum((parameters[:2] - global_center) ** 2) + local_rate = parameters[2 : 2 + problem.cells] + local_initial = parameters[2 + problem.cells :] + # Local log-rate offsets shrink toward a shared, fitted log-gain mean. + prior = ( + shared + + xp.sum(((local_rate - 0.1 * parameters[0]) / 0.5) ** 2) + + xp.sum((local_initial / 0.8) ** 2) + ) + return cast("Array", 0.5 * (xp.sum(residual**2) + prior)) + + +def reference_objective( + problem: HierarchicalProblem, + parameters: np.ndarray, + times: np.ndarray, + data: np.ndarray, +) -> float: + """Evaluate the same public statistical problem with independent scalar priors. + + Returns: + Concrete reference negative log objective. + """ + trajectory = reference_trajectory(problem, parameters, times) + likelihood = sum( + ((float(value) - float(observed)) / 0.025) ** 2 + for value, observed in zip(trajectory[1:, 0].ravel(), data.ravel(), strict=True) + ) + center = nominal_parameters(problem) + shared = sum((float(parameters[i]) - float(center[i])) ** 2 for i in range(2)) + local = sum( + ((float(value) - 0.1 * float(parameters[0])) / 0.5) ** 2 + for value in parameters[2 : 2 + problem.cells] + ) + initial = sum( + (float(value) / 0.8) ** 2 for value in parameters[2 + problem.cells :] + ) + return 0.5 * (likelihood + shared + local + initial) + + +def loss_function( + problem: HierarchicalProblem, + schedule: AdaptiveStepSchedule, + config: RunConfig, + data: Array, + *, + external: bool = False, +) -> Callable[[Array], Array]: + """Bind the fixed-mesh objective while retaining all dynamic parameters. + + Returns: + Scalar objective callable for JAX transforms. + """ + + def evaluate(parameters: Array) -> Array: + trajectory = replay_trajectory( + problem, parameters, schedule, config, external=external + ) + return objective(problem, parameters, trajectory, data) + + return evaluate diff --git a/docs/guides/hierarchical-ode.md b/docs/guides/hierarchical-ode.md new file mode 100644 index 0000000..e953ce0 --- /dev/null +++ b/docs/guides/hierarchical-ode.md @@ -0,0 +1,88 @@ +# Hierarchical nonlinear ODE evidence + +`benchmarks/hierarchical_ode.py` defines a synthetic generic two-pool model for +structured numerical inference. This is a correctness workload with explicit +reference equations, not a calibrated model or a new inference framework. +The fixture separates group count and local parameter dimension, with state +shape `(2 pools, groups, local cells)`. + +## Problem definition and parameter layout + +For group `g` and local cell `j`, write first pool `a[g,j]` and second pool +`b[g,j]`. Each pair has unit total. The forward and return flows are: + +```text +q[g,j] = exp(theta_gain) * exp(theta_local[g,j]) + * a[g,j] * (a[g,j] + c * mean_group(a[:,j])) / (1 + a[g,j]) +r[g,j] = exp(theta_recovery) * b[g,j] +a' = r - q +b' = q - r +``` + +`c=0` gives independent groups; `c=0.15` gives the coupled variant. Parameters +are ordered as two shared log rates, all group-local log-rate offsets, then all +local initial-condition coefficients. Thus there are +`2 + 2 * groups * local_dimension` fitted values. Initial values are +`a0=0.55+0.1*tanh(theta_initial)` and `b0=1-a0`; positivity and unit totals are +maintained at initialization without stopping differentiation. + +Nominal shared gain is 0.4. Recovery is 0.15 in the nonstiff regime and 50 in +the stiff regime, creating a fast relaxation mode. Local rate offsets and IC +coefficients vary deterministically across cells. Stiffness is a parameter +regime rather than a change to the structural model identity. + +## Independent trajectories and objective + +The reference RHS uses scalar standard-library math and explicit group/cell +loops. `reference_trajectory` passes that independent RHS to SciPy DOP853 with +`rtol=2e-12`, `atol=2e-14`. Native equations are compared directly with the +reference before trajectory comparisons. NumPy replay, JAX built-in scan replay, +and an external JAX scan agree on one frozen nominal mesh. + +Observations are first-pool values at noninitial output times, plus the +reproducible residual `0.025*0.2*sin(cell_index+0.5)`. The Gaussian likelihood +uses scale 0.025. Shared log rates have unit-scale priors about their nominal +centers; local log-rate offsets have scale 0.5 about `0.1*theta_gain`, and local +IC coefficients have scale 0.8 about zero. The sum of group likelihood terms +and shared/local priors defines the scalar objective up to constants. + +The nonzero residual is deliberate: derivatives include curvature of the +trajectory, rather than only the squared sensitivity at a perfect fit. +The evidence checks all shared/local/IC parameters, with independent-reference +finite differences for gradients and centered differences of native gradients +for Hessians. Forward and reverse mode agree. Step halving checks gradient +stability as well as trajectory agreement. + +The Laplace-style diagnostic is `0.5*logdet(H + 25*I)`, where `H` is the +objective Hessian. The fixed diagonal regularization is part of this numerical +diagnostic, not a claim that an unregularized posterior is convex. Tests verify +positive eigenvalues before interpreting the log determinant, require a +nonzero gradient, and compare that gradient with centered log-determinant +differences. Accuracy or positive definiteness outside these fixtures requires +fresh evidence. + +## Public step boundary and conditional derivatives + +Nominal schedules are discovered eagerly with Dormand–Prince at `rtol=2e-8`, +`atol=2e-10`. Built-in replay uses `replay_loop="scan"` and checkpointing. +The external loop uses only `CoreSolver.fixed_explicit_step`, reuses the FSAL +stage, and checkpoints its scan body. Output selections preserve the original +time grid. Both invoke the same accepted-step formulas; neither differentiates +adaptive acceptance decisions. Follow the [provenance/refresh policy](replay-provenance.md) +when parameters move. The halved comparison mesh is a manually refined mesh, +not a fresh set of controller acceptance decisions. + +## Reproduce the small correctness evidence + +```shell +uv sync --dev +uv run pytest tests/op_engine/test_hierarchical_ode.py -q +``` + +The eight cases cover independent/coupled and nonstiff/stiff regimes, direct +RHS agreement, local totals, accurate reference trajectories, NumPy/JAX parity, +built-in/external loops, forward/reverse gradients, Hessians, log-determinant +gradients, and mesh refinement. The NumPy/JAX required qualification jobs reject skipped fixture cases, and +installed core wheels execute the same evidence. Large dimension/mesh timing sweeps and fair +Diffrax comparisons are a separate opt-in measurement layer; no speedup or +memory-scaling claim follows from these correctness cases alone. diff --git a/justfile b/justfile index dd0059b..a0d6993 100644 --- a/justfile +++ b/justfile @@ -103,6 +103,8 @@ build-test-core: cp pyproject.toml "${CLEANROOM}/pyproject.toml" cp -R tests "${CLEANROOM}/tests" cp -R examples "${CLEANROOM}/examples" + + cp -R benchmarks "${CLEANROOM}/benchmarks" cd "${CLEANROOM}" "${CLEANROOM}/venv/bin/pytest" --import-mode=importlib tests --quiet --exitfirst diff --git a/mkdocs.yml b/mkdocs.yml index d9bc7c8..020fd48 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -59,6 +59,8 @@ nav: - Backend Qualification: guides/backend-qualification.md - Replay Provenance: guides/replay-provenance.md + + - Hierarchical ODE Evidence: guides/hierarchical-ode.md - Backend Strategies: guides/backends.md - API Reference: - Adaptive Tau: api-reference/adaptive_tau.md diff --git a/pyproject.toml b/pyproject.toml index b5a644b..998a095 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,6 +67,7 @@ dev = [ ] [tool.pytest.ini_options] +pythonpath = ["."] testpaths = [ "docs/", "examples/", diff --git a/scripts/qualify_backend.py b/scripts/qualify_backend.py index 166552a..bf8bd12 100644 --- a/scripts/qualify_backend.py +++ b/scripts/qualify_backend.py @@ -43,6 +43,7 @@ def main() -> int: "pytest", "tests/op_engine/test_backend_qualification.py", "tests/op_engine/test_operator_precision.py", + "tests/op_engine/test_hierarchical_ode.py", *( ["tests/op_engine/test_array_namespace.py"] if args.backend == "torch" diff --git a/tests/op_engine/test_hierarchical_ode.py b/tests/op_engine/test_hierarchical_ode.py new file mode 100644 index 0000000..0e58642 --- /dev/null +++ b/tests/op_engine/test_hierarchical_ode.py @@ -0,0 +1,167 @@ +"""Independent trajectory and higher-order AD evidence on hierarchical states.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, cast + +import numpy as np +import pytest + +from benchmarks.hierarchical_ode import ( + HierarchicalProblem, + initial_state, + loss_function, + nominal_parameters, + numpy_replay_trajectory, + observations, + record_schedule, + reference_initial, + reference_objective, + reference_rhs, + reference_trajectory, + refined_schedule, + replay_trajectory, + rhs, +) + +if TYPE_CHECKING: + from collections.abc import Iterator + + from op_engine import Array + +TIMES = np.asarray([0.0, 0.2, 0.5, 1.0]) +REGIMES = [(False, False), (False, True), (True, False), (True, True)] + + +@pytest.fixture +def jax_x64() -> Iterator[None]: + """Scope double precision to independent derivative evidence.""" + jax = pytest.importorskip("jax") + enable = getattr(jax, "enable_x64", None) + if enable is None: + enable = jax.experimental.enable_x64 + with enable(): + yield + + +@pytest.mark.parametrize(("stiff", "coupled"), REGIMES) +def test_numpy_hierarchical_equations_and_trajectories( + *, stiff: bool, coupled: bool +) -> None: + """Cell-loop equations and an accurate independent solve agree for all regimes.""" + problem = HierarchicalProblem( + groups=3, local_dimension=2, stiff=stiff, coupled=coupled + ) + parameters = nominal_parameters(problem) + initial = np.asarray(initial_state(problem, parameters)) + np.testing.assert_allclose( + initial, reference_initial(problem, parameters), rtol=1e-14 + ) + compiled = np.asarray(rhs(problem, parameters)(0.0, initial)) + expected = reference_rhs(problem, parameters, initial) + np.testing.assert_allclose(compiled, expected, rtol=1e-14) + np.testing.assert_allclose(compiled.sum(axis=0), np.zeros((3, 2)), atol=1e-14) + schedule, config = record_schedule(problem, parameters, TIMES) + trajectory = numpy_replay_trajectory(problem, parameters, schedule, config) + independent = reference_trajectory(problem, parameters, TIMES) + np.testing.assert_allclose(trajectory, independent, rtol=2e-7, atol=2e-9) + np.testing.assert_allclose( + trajectory.sum(axis=1), np.ones((len(TIMES), 3, 2)), atol=2e-13 + ) + refined = numpy_replay_trajectory( + problem, parameters, refined_schedule(schedule), config + ) + assert ( + np.max(abs(refined - independent)) + <= np.max(abs(trajectory - independent)) + 2e-13 + ) + + +@pytest.mark.usefixtures("jax_x64") +@pytest.mark.parametrize(("stiff", "coupled"), REGIMES) +def test_jax_hierarchical_forward_reverse_hessian_and_logdet( # noqa: PLR0914 + *, stiff: bool, coupled: bool +) -> None: + """Shared/local/IC derivatives use a nonzero-residual hierarchical objective.""" + jax = pytest.importorskip("jax") + jnp = pytest.importorskip("jax.numpy") + problem = HierarchicalProblem(stiff=stiff, coupled=coupled) + nominal = nominal_parameters(problem) + parameters = jnp.asarray(nominal) + data_np = observations(problem, TIMES) + data = jnp.asarray(data_np) + schedule, config = record_schedule(problem, nominal, TIMES) + trajectory = jax.jit( + lambda values: replay_trajectory(problem, values, schedule, config) + )(parameters) + external = jax.jit( + lambda values: replay_trajectory( + problem, values, schedule, config, external=True + ) + )(parameters) + reference = reference_trajectory(problem, nominal, TIMES) + np.testing.assert_allclose(trajectory, reference, rtol=2e-7, atol=2e-9) + np.testing.assert_allclose(external, trajectory, rtol=1e-12, atol=1e-13) + np.testing.assert_allclose( + trajectory, + numpy_replay_trajectory(problem, nominal, schedule, config), + rtol=1e-12, + atol=1e-13, + ) + assert np.max(abs(reference[1:, 0] - data_np)) > 0.001 + loss = loss_function(problem, schedule, config, data) + external_loss = loss_function(problem, schedule, config, data, external=True) + value, reverse = jax.jit(jax.value_and_grad(loss))(parameters) + forward = jax.jit(jax.jacfwd(loss))(parameters) + np.testing.assert_allclose(reverse, forward, rtol=1e-10, atol=1e-10) + np.testing.assert_allclose( + value, + reference_objective(problem, nominal, TIMES, data_np), + rtol=1e-6, + atol=1e-7, + ) + perturbations = 1e-4 * np.eye(problem.parameter_count) + reference_gradient = np.asarray([ + ( + reference_objective(problem, nominal + shift, TIMES, data_np) + - reference_objective(problem, nominal - shift, TIMES, data_np) + ) + / 2e-4 + for shift in perturbations + ]) + np.testing.assert_allclose(reverse, reference_gradient, rtol=2e-5, atol=2e-6) + hessian = jax.jit(jax.hessian(loss))(parameters) + gradient = jax.jit(jax.grad(loss)) + hessian_fd = np.column_stack([ + ( + np.asarray(gradient(parameters + shift)) + - np.asarray(gradient(parameters - shift)) + ) + / 2e-4 + for shift in perturbations + ]) + np.testing.assert_allclose(hessian, hessian_fd, rtol=1e-5, atol=2e-6) + np.testing.assert_allclose(hessian, np.asarray(hessian).T, rtol=1e-12, atol=1e-12) + regularization = 25.0 + regularized = hessian + regularization * jnp.eye(problem.parameter_count) + assert np.min(np.linalg.eigvalsh(np.asarray(regularized))) > 0 + + def logdet(values: Array) -> Array: + matrix = jax.hessian(loss)(values) + regularization * jnp.eye( + problem.parameter_count + ) + return cast("Array", 0.5 * jnp.linalg.slogdet(matrix)[1]) + + laplace_gradient = jax.jit(jax.grad(logdet))(parameters) + evaluate = jax.jit(logdet) + logdet_fd = np.asarray([ + (evaluate(parameters + shift) - evaluate(parameters - shift)) / 2e-4 + for shift in perturbations + ]) + assert np.linalg.norm(np.asarray(laplace_gradient)) > 1e-5 + np.testing.assert_allclose(laplace_gradient, logdet_fd, rtol=2e-5, atol=2e-6) + external_gradient = jax.jit(jax.grad(external_loss))(parameters) + np.testing.assert_allclose(external_gradient, reverse, rtol=1e-10, atol=1e-10) + refined = loss_function(problem, refined_schedule(schedule), config, data) + refined_gradient = jax.jit(jax.grad(refined))(parameters) + np.testing.assert_allclose(refined_gradient, reverse, rtol=2e-5, atol=2e-6)