Merge pull request #75 from pnkraemer/odefilter-interface

odesolver -> odefilter

Author: pnkraemer

tests/test_odesolver.py renamed to tests/test_odefilter.py

@@ -17,7 +17,7 @@ class EulerState:
     reference_state: jnp.array
 
 
-class EulerAsODEFilter(tornado.odesolver.ODEFilter):
+class EulerAsODEFilter(tornado.odefilter.ODEFilter):
     def initialize(self, ivp):
         return EulerState(
             ivp=ivp, y=ivp.y0, t=ivp.t0, error_estimate=None, reference_state=ivp.y0
@@ -31,15 +31,15 @@ def attempt_step(self, state, dt):
         )
 
 
-def test_odesolver():
+def test_odefilter():
     ivp = tornado.ivp.vanderpol(t0=0.0, tmax=1.5)
     constant_steps = tornado.step.ConstantSteps(dt=0.1)
     solver_order = 2
     solver = EulerAsODEFilter(
         steprule=constant_steps,
         num_derivatives=solver_order,
     )
-    assert isinstance(solver, tornado.odesolver.ODEFilter)
+    assert isinstance(solver, tornado.odefilter.ODEFilter)
 
     gen_sol = solver.solution_generator(ivp)
     for idx, _ in enumerate(gen_sol):

tornado/__init__.py

@@ -11,7 +11,7 @@
     iwp,
     kalman,
     linops,
-    odesolver,
+    odefilter,
     rv,
     sqrt,
     step,

tornado/ek0.py

@@ -3,10 +3,10 @@
 import jax.numpy as jnp
 
 import tornado.iwp
-from tornado import init, ivp, iwp, odesolver, rv, sqrt, step
+from tornado import init, ivp, iwp, odefilter, rv, sqrt, step
 
 
-class ReferenceEK0(odesolver.ODEFilter):
+class ReferenceEK0(odefilter.ODEFilter):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.P0 = None
@@ -34,7 +34,7 @@ def initialize(self, ivp):
         y = rv.MultivariateNormal(
             mean=mean, cov_sqrtm=jnp.kron(jnp.eye(ivp.dimension), cov_sqrtm)
         )
-        return odesolver.ODEFilterState(
+        return odefilter.ODEFilterState(
             ivp=ivp,
             t=ivp.t0,
             y=y,
@@ -68,7 +68,7 @@ def attempt_step(self, state, dt, verbose=False):
 
         y_new = jnp.abs(self.E0 @ m_new)
 
-        return odesolver.ODEFilterState(
+        return odefilter.ODEFilterState(
             ivp=state.ivp,
             t=state.t + dt,
             error_estimate=error,
@@ -77,7 +77,7 @@ def attempt_step(self, state, dt, verbose=False):
         )
 
 
-class KroneckerEK0(odesolver.ODEFilter):
+class KroneckerEK0(odefilter.ODEFilter):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.P0 = None
@@ -109,7 +109,7 @@ def initialize(self, ivp):
 
         y = rv.MatrixNormal(mean=mean, cov_sqrtm_1=jnp.eye(d), cov_sqrtm_2=cov_sqrtm)
 
-        return odesolver.ODEFilterState(
+        return odefilter.ODEFilterState(
             ivp=ivp,
             t=ivp.t0,
             error_estimate=None,
@@ -162,7 +162,7 @@ def attempt_step(self, state, dt, verbose=False):
         d = z.shape[0]
         error_estimate = jnp.stack([jnp.sqrt(sigma_squared * HQH)] * d)
 
-        return odesolver.ODEFilterState(
+        return odefilter.ODEFilterState(
             ivp=state.ivp,
             t=t_new,
             error_estimate=error_estimate,

tornado/ek1.py

@@ -5,10 +5,10 @@
 import jax.numpy as jnp
 import jax.scipy.linalg
 
-from tornado import init, iwp, linops, odesolver, rv, sqrt
+from tornado import init, iwp, linops, odefilter, rv, sqrt
 
 
-class ReferenceEK1(odesolver.ODEFilter):
+class ReferenceEK1(odefilter.ODEFilter):
     """Naive, reference EK1 implementation. Use this to test against."""
 
     def __init__(self, *args, **kwargs):
@@ -34,7 +34,7 @@ def initialize(self, ivp):
         )
         mean = extended_dy0.reshape((-1,), order="F")
         y = rv.MultivariateNormal(mean, jnp.kron(jnp.eye(ivp.dimension), cov_sqrtm))
-        return odesolver.ODEFilterState(
+        return odefilter.ODEFilterState(
             ivp=ivp,
             t=ivp.t0,
             y=y,
@@ -65,7 +65,7 @@ def attempt_step(self, state, dt):
         reference_state = jnp.maximum(y1, y2)
 
         # Return new state
-        return odesolver.ODEFilterState(
+        return odefilter.ODEFilterState(
             ivp=state.ivp,
             t=t,
             y=new_rv,
@@ -127,7 +127,7 @@ def estimate_error(h, sq, z):
         return error_estimate, sigma
 
 
-class BatchedEK1(odesolver.ODEFilter):
+class BatchedEK1(odefilter.ODEFilter):
     """Common functionality for EK1 variations that act on batched multivariate normals."""
 
     def __init__(self, *args, **kwargs):
@@ -159,7 +159,7 @@ def initialize(self, ivp):
         d, n = self.iwp.wiener_process_dimension, self.iwp.num_derivatives + 1
         cov_sqrtm = jnp.stack([cov_sqrtm] * d)
         new_rv = rv.BatchedMultivariateNormal(extended_dy0, cov_sqrtm)
-        return odesolver.ODEFilterState(
+        return odefilter.ODEFilterState(
             ivp=ivp,
             t=ivp.t0,
             y=new_rv,
@@ -186,7 +186,7 @@ def attempt_step(self, state, dt):
         reference_state = jnp.maximum(y1, y2)
 
         new_rv = rv.BatchedMultivariateNormal(new_mean, cov_sqrtm)
-        return odesolver.ODEFilterState(
+        return odefilter.ODEFilterState(
             ivp=state.ivp,
             t=t,
             y=new_rv,
@@ -420,7 +420,7 @@ def correct_cov_sqrtm(p_1d_raw, Jx, sc_bd, kgain):
         return new_sc
 
 
-class EarlyTruncationEK1(odesolver.ODEFilter):
+class EarlyTruncationEK1(odefilter.ODEFilter):
     """Use full Jacobians for mean-updates, but truncate cleverly to enforce a block-diagonal posterior covariance.
 
     "Cleverly" means:
@@ -458,7 +458,7 @@ def initialize(self, ivp):
         d, n = self.iwp.wiener_process_dimension, self.iwp.num_derivatives + 1
         cov_sqrtm = jnp.stack([cov_sqrtm] * d)
         new_rv = rv.BatchedMultivariateNormal(extended_dy0, cov_sqrtm)
-        return odesolver.ODEFilterState(
+        return odefilter.ODEFilterState(
             ivp=ivp,
             t=ivp.t0,
             y=new_rv,
@@ -616,7 +616,7 @@ def attempt_step(self, state, dt):
 
         # Return new state
         new_rv = rv.BatchedMultivariateNormal(new_mean, cov_sqrtm)
-        return odesolver.ODEFilterState(
+        return odefilter.ODEFilterState(
             ivp=state.ivp,
             t=t,
             y=new_rv,

tornado/ivpsolve.py

@@ -4,10 +4,10 @@
 
 import jax.numpy as jnp
 
-from tornado import ek0, ek1, ivp, odesolver, rv, step
+from tornado import ek0, ek1, ivp, odefilter, rv, step
 
 # Will be extended in the dev process
-_SOLVER_REGISTRY: Dict[str, odesolver.ODEFilter] = {
+_SOLVER_REGISTRY: Dict[str, odefilter.ODEFilter] = {
     "ek1_reference": ek1.ReferenceEK1,
     "ek1_diagonal": ek1.DiagonalEK1,
     "ek1_truncation": ek1.TruncationEK1,