checkpoint a new pattern for modifying compartment values

src/dynode/pydantic_config/config_definition.py

@@ -1,7 +1,7 @@
 """Top level classes for DynODE configs."""
 
 from datetime import date
-from typing import Callable, List, Optional
+from typing import Callable, List, Optional, Union
 
 from jax import Array
 from jax import numpy as jnp
@@ -50,6 +50,13 @@ def shape(self) -> tuple[int, ...]:
         """Get shape of the compartment."""
         return tuple([len(d_i) for d_i in self.dimensions])
 
+    def __setitem__(
+        self, index: Union[int, slice, tuple], value: float
+    ) -> None:
+        """Experimental function that sets a value in the JAX array using functional update with slicing support."""
+        assert isinstance(self.values, Array), "values is not an array"
+        self.values = self.values.at[index].set(value)
+
 
 class ParamStore(BaseModel):
     """Miscellaneous parameters of an ODE model."""

src/dynode/pydantic_config/covid_seip_config.py

@@ -40,6 +40,22 @@ def __init__(self):
             ode_function=seip_ode,
         )
 
+    @property
+    def s(self) -> Compartment:
+        return self.compartments[0]
+
+    @property
+    def e(self) -> Compartment:
+        return self.compartments[1]
+
+    @property
+    def i(self) -> Compartment:
+        return self.compartments[2]
+
+    @property
+    def c(self) -> Compartment:
+        return self.compartments[3]
+
     def _get_param_store(self, strains: list[Strain]) -> ParamStore:
         return ParamStore(
             strains=strains,

src/dynode/pydantic_config/dimension.py

@@ -52,6 +52,7 @@ class FullStratifiedImmuneHistory(Dimension):
 
     def __init__(self, strains: list[Strain]) -> None:
         """Create a fully stratified immune history dimension."""
+        # TODO add a no-infection bin
         strain_names = [s.strain_name for s in strains]
         num_strains = len(strain_names)
         all_immune_histories = []
@@ -73,6 +74,7 @@ class LastStrainImmuneHistory(Dimension):
 
     def __init__(self, strains: list[Strain]) -> None:
         """Create an immune history dimension that only tracks last infected strain."""
+        # TODO add a no-infection bin
         strain_names = [s.strain_name for s in strains]
         bins: list[Bin] = [
             CategoricalBin(name=state) for state in strain_names