Initial condition returns an IVPSolution object now

Author: pnkraemer

probdiffeq/ivpsolve.py

@@ -148,7 +148,7 @@ def solve_adaptive_save_at(
 
     # I think the user expects the initial condition to be part of the state
     # (as well as marginals), so we compute those things here
-    posterior_t0, *_ = initial_condition
+    posterior_t0 = initial_condition.posterior
     posterior_save_at, output_scale = solution_save_at
     _tmp = _userfriendly_output(
         posterior=posterior_save_at, posterior_t0=posterior_t0, ssm=ssm
@@ -238,7 +238,7 @@ def solve_adaptive_save_every_step(
     t = np.concatenate((np.atleast_1d(t0), t))
 
     # I think the user expects marginals, so we compute them here
-    posterior_t0, *_ = initial_condition
+    posterior_t0 = initial_condition.posterior
     posterior, output_scale = solution_every_step
     _tmp = _userfriendly_output(posterior=posterior, posterior_t0=posterior_t0, ssm=ssm)
     marginals, posterior = _tmp
@@ -296,7 +296,8 @@ def body_fn(s, dt):
     _t, (posterior, output_scale) = solver.extract(result_state)
 
     # I think the user expects marginals, so we compute them here
-    posterior_t0, *_ = initial_condition
+    # posterior_t0, *_ = initial_condition
+    posterior_t0 = initial_condition.posterior
     _tmp = _userfriendly_output(posterior=posterior, posterior_t0=posterior_t0, ssm=ssm)
     marginals, posterior = _tmp
 

probdiffeq/ivpsolvers.py

@@ -1,6 +1,6 @@
 """Probabilistic IVP solvers."""
 
-from probdiffeq import stats
+from probdiffeq import ivpsolve, stats
 from probdiffeq.backend import (
     containers,
     control_flow,
@@ -694,13 +694,25 @@ def is_suitable_for_save_at(self):
     def is_suitable_for_save_every_step(self):
         return self.strategy.is_suitable_for_save_every_step
 
-    def init(self, t, initial_condition) -> _State:
-        posterior, output_scale = initial_condition
+    def initial_condition(self) -> ivpsolve.IVPSolution:
+        """Construct an initial condition."""
+        posterior = self.strategy.initial_condition(prior=self.prior)
+        return ivpsolve.IVPSolution(
+            t=None,
+            u=None,
+            u_std=None,
+            output_scale=self.prior.output_scale,
+            marginals=None,
+            posterior=posterior,
+            num_steps=None,
+            ssm=None,
+        )
 
-        rv, extra = self.strategy.init(posterior)
+    def init(self, t, initial_condition: ivpsolve.IVPSolution) -> _State:
+        rv, extra = self.strategy.init(initial_condition.posterior)
         rv, corr = self.correction.init(rv)
 
-        calib_state = self.calibration.init(output_scale)
+        calib_state = self.calibration.init(initial_condition.output_scale)
         return _State(t=t, hidden=rv, aux_extra=extra, output_scale=calib_state)
 
     def step(self, state: _State, *, vector_field, dt):
@@ -766,11 +778,6 @@ def _state(t_, s, scale):
         acc = _state(t, step_from_, interp_to.output_scale)
         return _InterpRes(step_from=acc, interpolated=sol, interp_from=prev)
 
-    def initial_condition(self):
-        """Construct an initial condition."""
-        posterior = self.strategy.initial_condition(prior=self.prior)
-        return posterior, self.prior.output_scale
-
 
 def solver_mle(strategy, *, correction, prior, ssm):
     """Create a solver that calibrates the output scale via maximum-likelihood.