migrate to pydantic v2

Author: JR-1991

catalax/model/base.py

@@ -6,7 +6,6 @@ class CatalaxBase(BaseModel):
     class Config:
         use_enum_values = True
         arbitrary_types_allowed = True
-        allow_mutation = True
         validate_assignment = True
 
-    __repr_fields__: List[str] = PrivateAttr(default=["__all__"])
+    _repr_fields: List[str] = PrivateAttr(default=["__all__"])

catalax/model/model.py

@@ -50,9 +50,6 @@ class Model(CatalaxBase):
 
     class Config:
         arbitrary_types_allowed = True
-        fields = {
-            "term": {"exclude": True},
-        }
 
     name: str
     odes: Dict[str, ODE] = Field(default_factory=DottedDict)
@@ -116,7 +113,7 @@ def add_ode(
             equation=equation, species=self.species[species], observable=observable
         )
 
-        self.odes[species].__model__ = self
+        self.odes[species]._model = self
 
     def add_species(self, species_string: str = "", **species_map):
         """Adds a single or multiple species to the model, which can later be used in ODEs.
@@ -164,7 +161,6 @@ def add_species(self, species_string: str = "", **species_map):
 
             # Make sure the symbol is valid
             check_symbol(symbol)
-
             self.species[symbol] = Species(name=name, symbol=Symbol(symbol))
 
     @staticmethod

catalax/model/ode.py

@@ -14,17 +14,13 @@
 class ODE(CatalaxBase):
     class Config:
         arbitrary_types_allowed = True
-        fields = {
-            "parameters": {"exclude": True},
-            "species": {"exclude": True},
-        }
 
     species: Species
     equation: Expr
     observable: bool = True
     parameters: Dict[Union[str, Expr], Parameter] = Field(default_factory=DottedDict)
 
-    __model__: Optional["Model"] = PrivateAttr(default=None)  # type: ignore
+    _model: Optional["Model"] = PrivateAttr(default=None)  # type: ignore
 
     @validator("equation", pre=True)
     def converts_ode_to_sympy(cls, value):
@@ -37,7 +33,7 @@ def __setattr__(self, name, value):
 
         super().__setattr__(name, value)
 
-        if name == "__model__":
+        if name == "_model":
             self.add_parameters_to_model()
 
     def add_parameters_to_model(self):
@@ -48,20 +44,20 @@ def add_parameters_to_model(self):
         done upon addition of the model due to no given knowledge of the species.
         """
 
-        if self.__model__ is None:
+        if self._model is None:
             return None
 
         for symbol in self.equation.free_symbols:
-            if str(symbol) in self.__model__.species or str(symbol) == "t":
+            if str(symbol) in self._model.species or str(symbol) == "t":
                 # Skip species and time symbol
                 continue
-            elif parameter_exists(str(symbol), self.__model__.parameters):
+            elif parameter_exists(str(symbol), self._model.parameters):
                 # Assign parameter if it is already present in the model
-                self.parameters[str(symbol)] = self.__model__.parameters[str(symbol)]
+                self.parameters[str(symbol)] = self._model.parameters[str(symbol)]
                 continue
 
             # Create a new one and add it to the model and ODE
             parameter = Parameter(name=str(symbol), symbol=symbol)  # type: ignore
 
             self.parameters[str(symbol)] = parameter
-            self.__model__.parameters[str(symbol)] = parameter
+            self._model.parameters[str(symbol)] = parameter

catalax/model/parameter.py

@@ -25,7 +25,7 @@ class Parameter(CatalaxBase):
     prior: Any = None  # TODO: Fix this typing
     _prior_str_: Optional[str] = None
 
-    @root_validator()
+    @root_validator(skip_on_failure=True)
     def _assign_prior_string(cls, values):
         if isinstance(values["prior"], tuple):
             prior, prior_str = values["prior"]
@@ -34,7 +34,7 @@ def _assign_prior_string(cls, values):
 
         return values
 
-    __repr_fields__: List[str] = PrivateAttr(
+    _repr_fields: List[str] = PrivateAttr(
         default={
             "name": "name",
             "symbol": "symbol",

catalax/model/utils.py

@@ -37,10 +37,10 @@ def __repr__(self):
 
     @staticmethod
     def _fields_to_print(cls_: CatalaxBase) -> Dict[str, str]:
-        if len(cls_.__repr_fields__) == 1 and cls_.__repr_fields__[0] == "__all__":
+        if len(cls_._repr_fields) == 1 and cls_._repr_fields[0] == "__all__":
             return {key: key for key in cls_.__dict__}
 
-        return cls_.__repr_fields__
+        return cls_._repr_fields
 
 
 def odeprint(y, expr):
@@ -77,16 +77,16 @@ def check_symbol(symbol: str) -> None:
     """Checks whether the given symbol is a valid symbol"""
 
     ERROR_MESSAGE = f"""Symbol '{symbol}' is not a valid symbol. The following rules apply:
-    
+
     (1) The first character must be a letter
     (2) The remaining characters can be letters and numbers
     (3) The symbol cannot end with an underscore
     (4) The symbol can contain at most one underscore followed by letters and numbers
-    
+
     These are valid symbols:
-    
+
     k1, k_12, k_max, k_max1
-    
+
     """
 
     # Convert to string to use string methods

catalax/tools/simulation.py

@@ -57,7 +57,7 @@ class Config:
     parameters: List[str]
     stoich_mat: jax.Array
     dt0: float = 0.1
-    solver: AbstractSolver = Tsit5
+    solver: Any = Tsit5
     rtol: float = 1e-5
     atol: float = 1e-5
     max_steps: int = 64**4