Undo some solution typing changes

Author: pnkraemer

docs/examples_advanced/neural_ode.py

@@ -170,7 +170,7 @@ def loss(
         marginal_likelihood = strategy.log_marginal_likelihood(
             data[:, None],
             standard_deviation=jnp.ones_like(grid) * stdev,
-            posterior=sol.full_solution,
+            posterior=sol.solution_full,
         )
         return -1 * marginal_likelihood, {"sol": sol}
 

docs/examples_basic/conditioning_on_zero_residual.py

@@ -80,7 +80,7 @@ def vector_field(y, t):  # noqa: ARG001
 dt0 = ivpsolve.dt0(lambda y: vector_field(y, t=t0), (u0,))
 solve = ivpsolve.solve_adaptive_save_at(solver=solver, errorest=errorest)
 sol = solve(init, save_at=ts, dt0=dt0, atol=1e-1, rtol=1e-1)
-markov_seq_posterior = sol.full_solution
+markov_seq_posterior = sol.solution_full
 
 
 # +

probdiffeq/ivpsolve.py

@@ -38,12 +38,9 @@ class Solution(Protocol, Generic[S]):
     num_steps: int
     """The number of steps taken by the solver."""
 
-    full_solution: Any
+    solution_full: Any
     """A full description of the solution (beyond 'u', e.g. for dense outputs)."""
 
-    hyperparams: Any
-    """A description of (calibrated) hyperparameters."""
-
 
 # Revisit this dependent typing one Python >=3.12 is enforced
 # Concretely, Something like Solver[T, S: Solution[T]](Protocol):...

probdiffeq/probdiffeq.py

@@ -513,10 +513,10 @@ class ProbabilisticSolution(Generic[C, T]):
     u: TaylorCoeffTarget[C, T]
     """The current ODE solution estimate."""
 
-    full_solution: T | MarkovSequence[T]
+    solution_full: T | MarkovSequence[T]
     """The current posterior estimate."""
 
-    # Todo: merge 'output_scale' and 'auxiliary'?
+    # Todo: merge 'output_scale' and 'auxiliary' and "fun_evals"?
     output_scale: Any
     """The current output scale."""
 
@@ -634,7 +634,7 @@ def offgrid_marginals(self, t, *, solution):
         def _extract_previous(pytree):
             return tree.tree_map(lambda s: s[index - 1, ...], pytree)
 
-        posterior_t0 = _extract_previous(solution.full_solution)
+        posterior_t0 = _extract_previous(solution.solution_full)
         t0 = _extract_previous(solution.t)
 
         # Extract the RHS
@@ -674,8 +674,8 @@ def interpolate(
 
         # Interpolate
         tmp = self.strategy.interpolate(
-            posterior_t0=interp_from.full_solution,
-            posterior_t1=interp_to.full_solution,
+            posterior_t0=interp_from.solution_full,
+            posterior_t1=interp_to.solution_full,
             transition_t0_t=transition_t0_t,
             transition_t_t1=transition_t_t1,
         )
@@ -684,7 +684,7 @@ def interpolate(
         step_from = ProbabilisticSolution(
             t=interp_to.t,
             # New:
-            full_solution=step_and_interpolate_from.step_from,
+            solution_full=step_and_interpolate_from.step_from,
             # Old:
             u=interp_to.u,
             output_scale=interp_to.output_scale,
@@ -696,7 +696,7 @@ def interpolate(
         interpolated = ProbabilisticSolution(
             t=t,
             # New:
-            full_solution=interpolated,
+            solution_full=interpolated,
             u=estimate,
             # Taken from the rhs point
             output_scale=interp_to.output_scale,
@@ -708,7 +708,7 @@ def interpolate(
         interp_from = ProbabilisticSolution(
             t=t,
             # New:
-            full_solution=step_and_interpolate_from.interp_from,
+            solution_full=step_and_interpolate_from.interp_from,
             # Old:
             u=interp_from.u,
             output_scale=interp_from.output_scale,
@@ -725,13 +725,13 @@ def interpolate_at_t1(
     ):
         """Interpolate the solution near a checkpoint."""
         del t
-        tmp = self.strategy.interpolate_at_t1(posterior_t1=interp_to.full_solution)
+        tmp = self.strategy.interpolate_at_t1(posterior_t1=interp_to.solution_full)
         (estimate, interpolated), step_and_interpolate_from = tmp
 
         prev = ProbabilisticSolution(
             t=interp_to.t,
             # New
-            full_solution=step_and_interpolate_from.interp_from,
+            solution_full=step_and_interpolate_from.interp_from,
             # Old
             u=interp_from.u,  # incorrect?
             output_scale=interp_from.output_scale,  # incorrect?
@@ -742,7 +742,7 @@ def interpolate_at_t1(
         sol = ProbabilisticSolution(
             t=interp_to.t,
             # New:
-            full_solution=interpolated,
+            solution_full=interpolated,
             u=estimate,
             # Old:
             output_scale=interp_to.output_scale,
@@ -753,7 +753,7 @@ def interpolate_at_t1(
         acc = ProbabilisticSolution(
             t=interp_to.t,
             # New:
-            full_solution=step_and_interpolate_from.step_from,
+            solution_full=step_and_interpolate_from.step_from,
             # Old
             u=interp_to.u,
             output_scale=interp_to.output_scale,
@@ -1339,7 +1339,7 @@ def init(self, t, u) -> ProbabilisticSolution:
         return ProbabilisticSolution(
             t=t,
             u=estimate,
-            full_solution=posterior,
+            solution_full=posterior,
             auxiliary=auxiliary,
             output_scale=output_scale_prior,
             num_steps=0,
@@ -1353,7 +1353,7 @@ def step(self, state, *, dt: float, damp: float):
 
         # Predict
         u, prediction = self.strategy.predict(
-            posterior=state.full_solution, transition=transition
+            posterior=state.solution_full, transition=transition
         )
 
         # Linearize
@@ -1381,7 +1381,7 @@ def step(self, state, *, dt: float, damp: float):
         return ProbabilisticSolution(
             t=state.t + dt,
             u=u,
-            full_solution=posterior,
+            solution_full=posterior,
             output_scale=state.output_scale,
             auxiliary=auxiliary,
             num_steps=state.num_steps + 1,
@@ -1398,8 +1398,8 @@ def userfriendly_output(
         ones = np.ones_like(output_scale)
         output_scale = output_scale[-1]
 
-        init = solution0.full_solution
-        posterior = solution.full_solution
+        init = solution0.solution_full
+        posterior = solution.solution_full
         estimate, posterior = self.strategy.finalize(
             posterior0=init, posterior=posterior, output_scale=output_scale
         )
@@ -1409,7 +1409,7 @@ def userfriendly_output(
         return ProbabilisticSolution(
             t=ts,
             u=estimate,
-            full_solution=posterior,
+            solution_full=posterior,
             output_scale=output_scale,
             num_steps=solution.num_steps,
             auxiliary=solution.auxiliary,
@@ -1456,7 +1456,7 @@ def init(self, t, u) -> ProbabilisticSolution:
         return ProbabilisticSolution(
             t=t,
             u=estimate,
-            full_solution=posterior,
+            solution_full=posterior,
             auxiliary=lin_state,
             output_scale=output_scale,
             num_steps=0,
@@ -1484,7 +1484,7 @@ def step(self, state: ProbabilisticSolution, *, dt: float, damp: float):
         # (Includes re-discretisation)
         transition = self.prior(dt, output_scale)
         u, prediction = self.strategy.predict(
-            state.full_solution, transition=transition
+            state.solution_full, transition=transition
         )
 
         # Relinearize
@@ -1502,7 +1502,7 @@ def step(self, state: ProbabilisticSolution, *, dt: float, damp: float):
         return ProbabilisticSolution(
             t=state.t + dt,
             u=u,
-            full_solution=posterior,
+            solution_full=posterior,
             num_steps=state.num_steps + 1,
             auxiliary=lin_state,
             output_scale=output_scale,
@@ -1517,8 +1517,8 @@ def userfriendly_output(
         ones = np.ones_like(solution.output_scale)
         output_scale = ones[-1, ...]
 
-        init = solution0.full_solution
-        posterior = solution.full_solution
+        init = solution0.solution_full
+        posterior = solution.solution_full
         estimate, posterior = self.strategy.finalize(
             posterior0=init, posterior=posterior, output_scale=output_scale
         )
@@ -1529,7 +1529,7 @@ def userfriendly_output(
         return ProbabilisticSolution(
             t=ts,
             u=estimate,
-            full_solution=posterior,
+            solution_full=posterior,
             output_scale=output_scale,
             num_steps=solution.num_steps,
             auxiliary=solution.auxiliary,
@@ -1567,7 +1567,7 @@ def init(self, t: Array, u: TaylorCoeffTarget) -> ProbabilisticSolution:
         return ProbabilisticSolution(
             t=t,
             u=u,
-            full_solution=posterior,
+            solution_full=posterior,
             num_steps=0,
             auxiliary=correction_state,
             output_scale=output_scale,
@@ -1581,7 +1581,7 @@ def step(self, state: ProbabilisticSolution, *, dt, damp):
 
         # Predict
         u, prediction = self.strategy.predict(
-            state.full_solution, transition=transition
+            state.solution_full, transition=transition
         )
 
         # Linearize
@@ -1598,7 +1598,7 @@ def step(self, state: ProbabilisticSolution, *, dt, damp):
         return ProbabilisticSolution(
             t=state.t + dt,
             u=u,
-            full_solution=posterior,
+            solution_full=posterior,
             output_scale=output_scale,
             auxiliary=auxiliary,
             num_steps=state.num_steps + 1,
@@ -1614,8 +1614,8 @@ def userfriendly_output(
         ones = np.ones_like(solution.output_scale)
         output_scale = np.ones_like(solution.output_scale[-1])
 
-        init = solution0.full_solution
-        posterior = solution.full_solution
+        init = solution0.solution_full
+        posterior = solution.solution_full
         u, posterior = self.strategy.finalize(
             posterior0=init, posterior=posterior, output_scale=output_scale
         )
@@ -1626,7 +1626,7 @@ def userfriendly_output(
         return ProbabilisticSolution(
             t=ts,
             u=u,
-            full_solution=posterior,
+            solution_full=posterior,
             output_scale=output_scale,
             num_steps=solution.num_steps,
             auxiliary=solution.auxiliary,

tests/test_probdiffeq/test_log_marginal_likelihood.py

@@ -31,7 +31,7 @@ def test_output_is_a_scalar_and_not_nan_and_not_inf(solution):
     data = tree.tree_map(lambda s: s + 0.005, sol.u.mean[0])
     std = tree.tree_map(lambda _s: np.ones_like(sol.t), sol.u.std[0])
     lml = strategy.log_marginal_likelihood(
-        data, standard_deviation=std, posterior=sol.full_solution
+        data, standard_deviation=std, posterior=sol.solution_full
     )
     assert lml.shape == ()
     assert not np.isnan(lml)
@@ -49,7 +49,7 @@ def test_that_function_raises_error_for_wrong_std_shape_too_many(solution):
 
     with testing.raises(ValueError, match="does not match"):
         _ = strategy.log_marginal_likelihood(
-            data, standard_deviation=std, posterior=sol.full_solution
+            data, standard_deviation=std, posterior=sol.solution_full
         )
 
 
@@ -63,7 +63,7 @@ def test_raises_error_for_terminal_values(solution):
     data = tree.tree_map(lambda s: s[-1] + 0.005, sol.u.mean[0])
     std = tree.tree_map(lambda _s: np.ones_like(sol.t[-1]), sol.u.std[0])
 
-    posterior_t1 = tree.tree_map(lambda s: s[-1], sol.full_solution)
+    posterior_t1 = tree.tree_map(lambda s: s[-1], sol.solution_full)
     with testing.raises(ValueError, match="expected"):
         _ = strategy.log_marginal_likelihood(
             data, standard_deviation=std, posterior=posterior_t1
@@ -91,7 +91,7 @@ def test_raises_error_for_filter(fact):
     std = tree.tree_map(np.ones_like, sol.u.std[0])
     with testing.raises(TypeError, match="ilter"):
         _ = strategy.log_marginal_likelihood(
-            data, standard_deviation=std, posterior=sol.full_solution
+            data, standard_deviation=std, posterior=sol.solution_full
         )
 
 
@@ -102,5 +102,5 @@ def test_raise_error_if_structures_dont_match(solution):
 
     with testing.raises(ValueError, match="tree structure"):
         _ = strategy.log_marginal_likelihood(
-            data, standard_deviation=std, posterior=sol.full_solution
+            data, standard_deviation=std, posterior=sol.solution_full
         )

tests/test_probdiffeq/test_log_marginal_likelihood_terminal_values.py

@@ -49,7 +49,7 @@ def test_output_is_scalar_and_not_inf_and_not_nan(solution):
     std = tree.tree_map(lambda _s: 1e-2 * np.ones(()), sol.u.std[0])
 
     mll = strategy.log_marginal_likelihood_terminal_values(
-        data, standard_deviation=std, posterior=sol.full_solution
+        data, standard_deviation=std, posterior=sol.solution_full
     )
 
     assert mll.shape == ()
@@ -64,5 +64,5 @@ def test_raise_error_if_structures_dont_match(solution):
 
     with testing.raises(ValueError, match="structure"):
         _ = strategy.log_marginal_likelihood_terminal_values(
-            data, standard_deviation=std, posterior=sol.full_solution
+            data, standard_deviation=std, posterior=sol.solution_full
         )

tests/test_probdiffeq/test_sample.py

@@ -30,7 +30,7 @@ def test_sample_shape(approximation_and_strategy, shape):
 
     key = random.prng_key(seed=15)
     samples = strategy.markov_sample(
-        key, approximation.full_solution, shape=shape, reverse=True
+        key, approximation.solution_full, shape=shape, reverse=True
     )
     for s, u in zip(samples, approximation.u.mean):
         s_shape = tree.tree_map(lambda x: x.shape, s)

tests/test_probdiffeq/test_strategy_smoother_fixedinterval_vs_fixedpoint.py

@@ -52,17 +52,16 @@ def test_fixedpoint_smoother_equivalent_same_grid(solver_setup, solution_smoothe
     assert testing.allclose(sol_fp.u.mean, sol_sm.u.mean)
     assert testing.allclose(sol_fp.u.std, sol_sm.u.std)
     assert testing.allclose(sol_fp.u.marginals, sol_sm.u.marginals)
-    assert testing.allclose(sol_fp.output_scale, sol_sm.output_scale)
     assert testing.allclose(sol_fp.num_steps, sol_sm.num_steps)
     assert testing.allclose(
-        sol_fp.full_solution.marginal, sol_sm.full_solution.marginal
+        sol_fp.solution_full.marginal, sol_sm.solution_full.marginal
     )
 
     # The backward conditionals use different parametrisations
     # but implement the same transitions
     cond_fp, cond_sm = (
-        sol_fp.full_solution.conditional,
-        sol_sm.full_solution.conditional,
+        sol_fp.solution_full.conditional,
+        sol_sm.solution_full.conditional,
     )
     cond_fp = func.vmap(ssm.conditional.preconditioner_apply)(cond_fp)
     cond_sm = func.vmap(ssm.conditional.preconditioner_apply)(cond_sm)