Reintroduce output scale into ibm transitions

Author: pnkraemer

probdiffeq/ivpsolvers.py

@@ -22,7 +22,7 @@ def prior_wiener_integrated(
     """Construct an adaptive(/continuous-time), multiply-integrated Wiener process."""
     ssm = impl.choose(ssm_fact, tcoeffs_like=tcoeffs)
     if output_scale is None:
-        output_scale = ssm.prototypes.output_scale()
+        output_scale = np.ones_like(ssm.prototypes.output_scale())
     discretize = ssm.conditional.ibm_transitions(base_scale=output_scale)
     init = ssm.normal.from_tcoeffs(tcoeffs)
     return init, discretize, ssm
@@ -35,7 +35,10 @@ def prior_wiener_integrated_discrete(
     init, discretize, ssm = prior_wiener_integrated(
         tcoeffs_like, output_scale=output_scale, ssm_fact=ssm_fact
     )
-    transitions, (p, p_inv) = functools.vmap(discretize)(np.diff(ts))
+
+    scales = np.ones_like(ssm.prototypes.output_scale())
+    discretize_vmap = functools.vmap(discretize, in_axes=(0, None))
+    transitions, (p, p_inv) = discretize_vmap(np.diff(ts), scales)
 
     preconditioner_apply_vmap = functools.vmap(ssm.conditional.preconditioner_apply)
     conditionals = preconditioner_apply_vmap(transitions, p, p_inv)
@@ -798,7 +801,7 @@ def _calibration_running_mean(*, ssm) -> _Calibration:
     #  In this case, the _calibration_most_recent() stuff becomes void.
 
     def init():
-        prior = ssm.prototypes.output_scale()
+        prior = np.ones_like(ssm.prototypes.output_scale())
         return prior, prior, 0.0
 
     def update(state, /, observed):
@@ -820,7 +823,7 @@ def solver_dynamic(strategy, *, correction, prior, ssm):
 
     def step_dynamic(state, /, *, dt, calibration):
         # Estimate error and calibrate the output scale
-        ones = ssm.prototypes.output_scale()
+        ones = np.ones_like(ssm.prototypes.output_scale())
         transition = prior(dt, ones)
         hidden = strategy.extrapolate_mean(state.rv, transition=transition)
         t = state.t + dt
@@ -855,7 +858,7 @@ def step_dynamic(state, /, *, dt, calibration):
 
 def _calibration_most_recent(*, ssm) -> _Calibration:
     def init():
-        return ssm.prototypes.output_scale()
+        return np.ones_like(ssm.prototypes.output_scale())
 
     def update(_state, /, observed):
         return ssm.stats.mahalanobis_norm_relative(0.0, observed)
@@ -906,7 +909,7 @@ def step(state: _State, *, dt, calibration):
 
 def _calibration_none(*, ssm) -> _Calibration:
     def init():
-        return ssm.prototypes.output_scale()
+        return np.ones_like(ssm.prototypes.output_scale())
 
     def update(_state, /, observed):
         raise NotImplementedError

probdiffeq/taylor.py

@@ -42,8 +42,11 @@ def starter(vf, initial_values, /, num: int, t):
         init = (rv_t0, conditional_t0)
 
         # Discretised prior
-        discretise = ssm.conditional.ibm_transitions(output_scale=1.0)
-        ibm_transitions = functools.vmap(discretise)(np.diff(ts))
+        scale = ssm.prototypes.output_scale()
+        discretise = ssm.conditional.ibm_transitions(scale)
+        ibm_transitions = functools.vmap(discretise, in_axes=(0, None))(
+            np.diff(ts), scale
+        )
 
         # Generate an observation-model for the QOI
         # (1e-7 observation noise for nuggets and for reusing existing code)

tests/test_impl/test_logpdfs.py

@@ -35,6 +35,6 @@ def random_variable(fact):
     tcoeffs = [np.ones((3,))] * 5  # values irrelevant
     ssm = impl.choose(fact, tcoeffs_like=tcoeffs)
     output_scale = np.ones_like(ssm.prototypes.output_scale())
-    discretize = ssm.conditional.ibm_transitions(output_scale=output_scale)
-    rv = discretize(0.1)
+    discretize = ssm.conditional.ibm_transitions(output_scale)
+    rv = discretize(0.1, output_scale)
     return rv[0].noise, ssm