Rename backend.containers into backend.structs to make code more compact

Author: pnkraemer

probdiffeq/backend/containers.py probdiffeq/backend/structs.pyprobdiffeq/backend/containers.py renamed to probdiffeq/backend/structs.py

@@ -1,13 +1,13 @@
 """LatentConds."""
 
-from probdiffeq.backend import abc, containers, func, linalg, np, tree
+from probdiffeq.backend import abc, func, linalg, np, structs, tree
 from probdiffeq.backend.typing import Array
 from probdiffeq.impl import _normal, _stats
 from probdiffeq.util import cholesky_util
 
 
 @tree.register_dataclass
-@containers.dataclass
+@structs.dataclass
 class LatentCond:
     """Conditional distributions in latent space."""
 

probdiffeq/impl/_conditional.py

@@ -1,4 +1,4 @@
-from probdiffeq.backend import abc, containers, func, linalg, np, random, tree
+from probdiffeq.backend import abc, func, linalg, np, random, structs, tree
 from probdiffeq.backend.typing import Any, Callable
 from probdiffeq.impl import _conditional, _normal
 from probdiffeq.util import cholesky_util
@@ -14,7 +14,7 @@ def linearize(self, fun, rv, state: None, *, damp: float):
         raise NotImplementedError
 
 
-@containers.dataclass
+@structs.dataclass
 class DenseTs0(Linearization):
     ode_order: int
     ode_shape: tuple
@@ -42,7 +42,7 @@ def a1(m):
         return cond, None
 
 
-@containers.dataclass
+@structs.dataclass
 class DenseTs1(Linearization):
     ode_order: int
     ode_shape: tuple
@@ -76,7 +76,7 @@ def constraint(m):
         return cond, None
 
 
-@containers.dataclass
+@structs.dataclass
 class DenseSlr0(Linearization):
     cubature_rule: Any
     ode_shape: tuple
@@ -149,7 +149,7 @@ def slr0(self, fn, x):
         return _normal.Normal(fx_mean, cov_sqrtm.T)
 
 
-@containers.dataclass
+@structs.dataclass
 class DenseSlr1(Linearization):
     cubature_rule: Any
     ode_shape: tuple
@@ -224,7 +224,7 @@ def slr1(self, fn, x):
         return linop_cond, rv_cond
 
 
-@containers.dataclass
+@structs.dataclass
 class IsotropicTs0(Linearization):
     ode_order: int
     unravel: Callable
@@ -253,7 +253,7 @@ def a1(m):
         return cond, None
 
 
-@containers.dataclass
+@structs.dataclass
 class IsotropicTs1(Linearization):
     ode_order: int
     unravel: Callable
@@ -305,7 +305,7 @@ def select_0(s):
         return cond, key
 
 
-@containers.dataclass
+@structs.dataclass
 class BlockDiagTs0(Linearization):
     ode_order: int
     unravel: Callable
@@ -336,7 +336,7 @@ def a1(s):
         return cond, None
 
 
-@containers.dataclass
+@structs.dataclass
 class BlockDiagTs1(Linearization):
     ode_order: int
     unravel: Callable

probdiffeq/impl/_linearize.py

@@ -1,11 +1,11 @@
-from probdiffeq.backend import abc, containers, linalg, np, tree
+from probdiffeq.backend import abc, linalg, np, structs, tree
 from probdiffeq.backend.typing import Generic, Sequence, TypeVar
 
 T = TypeVar("T")
 
 
 @tree.register_dataclass
-@containers.dataclass
+@structs.dataclass
 class Normal(Generic[T]):
     mean: T
     cholesky: T

probdiffeq/impl/_normal.py

@@ -1,11 +1,11 @@
 """State-space model implementations."""
 
-from probdiffeq.backend import containers, func, tree
+from probdiffeq.backend import func, structs, tree
 from probdiffeq.backend.typing import Callable
 from probdiffeq.impl import _conditional, _linearize, _normal, _prototypes, _stats
 
 
-@containers.dataclass
+@structs.dataclass
 class FactImpl:
     """Implementation of factorized state-space models."""
 

probdiffeq/impl/impl.py

@@ -7,7 +7,7 @@
 See the tutorials for example use cases.
 """
 
-from probdiffeq.backend import containers, flow, func, linalg, np, tree, warnings
+from probdiffeq.backend import flow, func, linalg, np, structs, tree, warnings
 from probdiffeq.backend.typing import Any, Array, Callable, Generic, Protocol, TypeVar
 
 T_contra = TypeVar("T_contra", contravariant=True)
@@ -225,7 +225,7 @@ def advance(sol_and_state: tuple, t_next) -> tuple[tuple, Any]:
             """
 
             @tree.register_dataclass
-            @containers.dataclass
+            @structs.dataclass
             class AdvanceState:
                 do_continue: bool
                 solution: Any
@@ -332,7 +332,7 @@ def dt0_adaptive(vf, initial_values, /, t0, *, error_contraction_rate, rtol, ato
 
 
 @tree.register_dataclass
-@containers.dataclass
+@structs.dataclass
 class TimeStepState(Generic[T]):
     """A state variable type for adaptive time-stepping."""
 
@@ -359,7 +359,7 @@ class TimeStepState(Generic[T]):
 
 
 @tree.register_dataclass
-@containers.dataclass
+@structs.dataclass
 class _RejectionLoopState:
     """State for a single rejection loop.
 

probdiffeq/ivpsolve.py

@@ -3,7 +3,7 @@
 See the tutorials for example use cases.
 """
 
-from probdiffeq.backend import containers, flow, func, linalg, np, random, special, tree
+from probdiffeq.backend import flow, func, linalg, np, random, special, structs, tree
 from probdiffeq.backend.typing import (
     Any,
     Array,
@@ -39,7 +39,7 @@ def __call__(self, *ys: T, t: ArrayLike) -> T: ...
 
 
 @tree.register_dataclass
-@containers.dataclass
+@structs.dataclass
 class CubaturePositiveWeights:
     """A datastructure for cubature rules that have positive weights.
 
@@ -220,7 +220,7 @@ def constraint_ode_slr1(*, ssm, cubature_fun=cubature_third_order_spherical):
 
 
 @tree.register_dataclass
-@containers.dataclass
+@structs.dataclass
 class TaylorCoeffTarget(Generic[C, T]):
     """A probabilistic description of Taylor coefficients.
 
@@ -240,7 +240,7 @@ class TaylorCoeffTarget(Generic[C, T]):
 
 
 @tree.register_dataclass
-@containers.dataclass
+@structs.dataclass
 class MarkovSequence(Generic[T]):
     """A datastructure for Markov sequences as batches of joint distributions.
 
@@ -492,7 +492,7 @@ def log_marginal_likelihood(self, u, /, *, standard_deviation, posterior: T):
 
 
 @tree.register_dataclass
-@containers.dataclass
+@structs.dataclass
 class ProbabilisticSolution(Generic[C, T]):
     """A datastructure for probabilistic solutions of differential equations."""
 
@@ -807,7 +807,7 @@ def prior_wiener_integrated_discrete(
 
 
 @tree.register_dataclass
-@containers.dataclass
+@structs.dataclass
 class _InterpRes(Generic[T]):
     """A datastructure to store interpolation results.
 

probdiffeq/probdiffeq.py

@@ -3,7 +3,7 @@
 Mostly **discrete** filtering and smoothing.
 """
 
-from probdiffeq.backend import containers, flow, tree
+from probdiffeq.backend import flow, structs, tree
 from probdiffeq.backend.typing import Any
 
 
@@ -40,7 +40,7 @@ def fixedpointsmoother_precon(*, ssm):
     """Construct a discrete, preconditioned fixedpoint-smoother."""
 
     @tree.register_dataclass
-    @containers.dataclass
+    @structs.dataclass
     class _FPState:
         rv: Any
         conditional: Any
@@ -78,7 +78,7 @@ def kalmanfilter_with_marginal_likelihood(*, ssm):
     """Construct a Kalman-filter-implementation of computing the marginal likelihood."""
 
     @tree.register_dataclass
-    @containers.dataclass
+    @structs.dataclass
     class _KFState:
         rv: Any
         num_data_points: float

probdiffeq/util/filter_util.py

@@ -1,7 +1,7 @@
 """Compare solve_fixed_grid to solve_adaptive_save_every_step."""
 
 from probdiffeq import ivpsolve, probdiffeq, taylor
-from probdiffeq.backend import containers, func, np, ode, testing, tree
+from probdiffeq.backend import func, np, ode, structs, testing, tree
 from probdiffeq.backend.typing import Array
 from probdiffeq.util import test_util
 
@@ -10,7 +10,7 @@
 def test_fixed_grid_result_matches_adaptive_grid_result_when_reusing_grid(fact):
     vf, u0, (t0, t1) = ode.ivp_lotka_volterra()
 
-    class Taylor(containers.NamedTuple):
+    class Taylor(structs.NamedTuple):
         state: Array
         velocity: Array
         acceleration: Array

tests/test_ivpsolve/test_fixed_grid_matches_adaptive_grid.py

@@ -1,11 +1,11 @@
 """Tests for interaction with the solution API."""
 
 from probdiffeq import ivpsolve, probdiffeq, taylor
-from probdiffeq.backend import containers, func, np, ode, testing
+from probdiffeq.backend import func, np, ode, structs, testing
 from probdiffeq.backend.typing import Array
 
 
-class Taylor(containers.NamedTuple):
+class Taylor(structs.NamedTuple):
     """A non-standard Taylor-coefficient data structure."""
 
     state: Array