factor out Parameters class

Author: bmtgoncalves

src/epidemik/EpiModel.py

@@ -15,7 +15,6 @@
 import networkx as nx
 import numpy as np
 from numpy import linalg
-from numpy import random
 import scipy.integrate
 import pandas as pd
 import matplotlib.pyplot as plt
@@ -49,7 +48,7 @@ def __init__(self, compartments=None, seed=None, rng=None):
         self.population = None
         self.orig_comps = None
         self.demographics = False
-        self.params = {}
+        self.params = utils.Parameters()
 
         if seed is None:
             seed = int(time.time()) + os.getpid()
@@ -62,24 +61,6 @@ def __init__(self, compartments=None, seed=None, rng=None):
         if compartments is not None:
             self.transitions.add_nodes_from([comp for comp in compartments])
 
-    def _compute_rate(self, rate: Union[str, float]) -> float:
-        """
-        Compute the rate from a string
-
-        Parameters:
-        - rate: string
-            Rate of the transition
-
-        Returns:
-        float
-            The computed rate
-        """
-        if rate in self.params:
-            if isinstance(self.params[rate], float):
-                return self.params[rate]
-            else:
-                return eval(self.params[rate], {}, self.params)
-
     def add_interaction(
         self,
         source: str,
@@ -108,9 +89,9 @@ def add_interaction(
         if rate is not None:
             count = len(self.params) + 1
             rate_key = "rate" + str(count)
-            self.define_parameters(rate_key=rate)
+            self.params[rate_key]=rate
         else:
-            self.define_parameters(**rates)
+            self.params.define_parameters(**rates)
             rates = list(rates.keys())
             rate_key = rates[0]
 
@@ -137,21 +118,22 @@ def add_spontaneous(
         if rate is not None:
             count = len(self.params) + 1
             rate_key = "rate" + str(count)
-            self.define_parameters(rate_key=rate)
+            self.params[rate_key, rate]
         else:
-            self.define_parameters(**rates)
+            self.params.define_parameters(**rates)
             rates = list(rates.keys())
             rate_key = rates[0]
 
         self.transitions.add_edge(source, target, rate=rate_key)
 
-    def add_viral_generation(self, 
-                             source:str, 
-                             target:str, 
-                             source_rate:Union[float, str,  None] = None, 
-                             target_rate:Union[float, str, None] = None,
-                             **rates
-        ) -> None:
+    def add_viral_generation(
+        self,
+        source: str,
+        target: str,
+        source_rate: Union[float, str, None] = None,
+        target_rate: Union[float, str, None] = None,
+        **rates,
+    ) -> None:
         """
         Add a viral generation transition
 
@@ -169,20 +151,25 @@ def add_viral_generation(self,
         if source_rate is not None and target_rate is not None:
             count = len(self.params) + 1
             rate_key = "rate" + str(count)
-            self.define_parameters(rate_key = source_rate)
+            self.params[rate_key] = source_rate
 
-            rate_key = "rate" + str(count+1)
-            self.define_parameters(**{rate_key: target_rate})
-        else: 
-            self.define_parameters(**rates)
+            rate_key = "rate" + str(count + 1)
+            self.params[rate_key] = target_rate
+        else:
+            self.params.define_parameters(**rates)
             rates = list(rates.keys())
             source_rate = rates[0]
             target_rate = rates[1]
 
-        self.transitions.add_edge(source, target, source_rate=source_rate, target_rate=target_rate)
+        self.transitions.add_edge(
+            source, target, rate=source_rate
+        )
+        self.transitions.add_edge(
+            source, target, rate=target_rate
+        )
 
     def add_birth_rate(
-        self, comps: Union[List, None] = None, rate: Union[float, None] = None, **rates
+        self, comps: Union[List, None] = None, rate: Union[float, None] = None, fixed=False, **rates
     ) -> None:
         """
         Add a birth rate to one or more compartments
@@ -212,6 +199,7 @@ def add_birth_rate(
                 self.transitions.add_node(comp)
 
             self.transitions.nodes[comp]["birth"] = rate_key
+            self.transitions.nodes[comp]["fixed"] = fixed
 
     def add_death_rate(
         self, comps: Union[List, None] = None, rate: Union[None, float] = None, **rates
@@ -246,7 +234,12 @@ def add_death_rate(
             self.transitions.nodes[comp]["death"] = rate_key
 
     def add_vaccination(
-        self, source: str, target: str, start: int, rate: Union[None, float, str], **rates
+        self,
+        source: str,
+        target: str,
+        start: int,
+        rate: Union[None, float, str],
+        **rates,
     ) -> None:
         """
         Add a vaccination transition between two compartments
@@ -275,27 +268,6 @@ def add_vaccination(
 
         self.transitions.add_edge(source, target, rate=rate_key, start=start)
 
-    def define_parameters(self, **kwargs) -> None:
-        """
-        Define one or more parameter for the model
-
-        Parameters:
-        - kwargs: keyword arguments
-            Named parameters for the model
-
-        Returns:
-        None
-        """
-        for key, value in kwargs.items():
-            if isinstance(value, str):
-                try:
-                    # Convert floats written as strings to float
-                    value = float(value.strip())
-                except:
-                    pass
-
-            self.params[key] = value
-
     def add_age_structure(self, matrix: List, population: List) -> List[List]:
         """
         Add a vaccination transition between two compartments
@@ -384,7 +356,7 @@ def _new_cases(self, time: float, population: np.ndarray, pos: Dict) -> np.ndarr
             target = edge[1]
             trans = edge[2]
 
-            rate_val = self._compute_rate(trans["rate"]) #self.params[trans["rate"]]
+            rate_val = self.params[trans["rate"]]
             rate = rate_val * population[pos[source]]
 
             if "start" in trans and trans["start"] >= time:
@@ -412,11 +384,15 @@ def _new_cases(self, time: float, population: np.ndarray, pos: Dict) -> np.ndarr
                     comp_id = pos[comp]
 
                     if "birth" in data:
-                        births = population[comp_id] * data["birth"]
+                        if "fixed" in data:
+                            births = self.params[data["birth"]]
+                        else:
+                            births = population[comp_id] * self.params[data["birth"]]
+    
                         diff[comp_id] += births
 
                     if "death" in data:
-                        deaths = population[comp_id] * data["death"]
+                        deaths = population[comp_id] * self.params[data["death"]]
                         diff[comp_id] -= deaths
 
         return diff
@@ -543,7 +519,7 @@ def simulate(
                     source = pos[comp]
                     target = pos[node_j]
 
-                    rate = self._compute_rate(data["rate"]) # self.params[data["rate"]]
+                    rate = self.params[data["rate"]]  # self.params[data["rate"]]
 
                     if "start" in data and data["start"] >= t:
                         continue
@@ -561,7 +537,7 @@ def simulate(
 
                 prob[source] = 1 - np.sum(prob)
 
-                delta = random.multinomial(pop[source], prob)
+                delta = self.rng.multinomial(pop[source], prob)
                 delta[source] = 0
 
                 changes = np.sum(delta)
@@ -580,11 +556,15 @@ def simulate(
                     comp_id = pos[comp]
 
                     if "birth" in data:
-                        births = self.rng.binomial(pop[comp_id], self._compute_rate(data["birth"]))
+                        births = self.rng.binomial(
+                            pop[comp_id], self.params[data["birth"]]
+                        )
                         new_pop[comp_id] += births
 
                     if "death" in data:
-                        deaths = self.rng.binomial(pop[comp_id], self._compute_rate(data["death"]))
+                        deaths = self.rng.binomial(
+                            pop[comp_id], self.params[data["death"]]
+                        )
                         new_pop[comp_id] -= deaths
 
             values.append(new_pop)
@@ -886,7 +866,7 @@ def R0(self) -> Union[float, None]:
 
         try:
             for node_i, node_j, data in self.transitions.edges(data=True):
-                rate = self._compute_rate(data["rate"]) # self.params[data["rate"]]
+                rate = self.params[data["rate"]]  # self.params[data["rate"]]
 
                 if "agent" in data:
                     target = pos[node_j]

src/epidemik/utils.py

@@ -10,6 +10,7 @@
     "NotInitialized",
     "get_cache_directory",
     "NotImplementedError",
+    "Parameters",
 ]
 
 
@@ -32,6 +33,62 @@ class NotImplementedError(Exception):
 EPI_COLORS["D"] = "#8b8b8b"
 
 
+class Parameters(dict):
+    def __init__(self):
+        super().__init__()
+        self.globals = {"__builtins__": None}
+
+    def __setitem__(self, key, value):
+        self.define_parameters(**{key: value})
+
+    def __getitem__(self, key):
+        return self.compute_parameter(key)
+
+    def define_parameters(self, **kwargs) -> None:
+        """
+        Define one or more parameter for the model
+
+        Parameters:
+        - kwargs: keyword arguments
+            Named parameters for the model
+
+        Returns:
+        None
+        """
+        for key, value in kwargs.items():
+            if isinstance(value, str):
+                try:
+                    # Convert floats written as strings to float
+                    value = float(value.strip())
+                except:
+                    pass
+            
+            super().__setitem__(key, value)
+
+    def compute_parameter(self, param: Union[str]) -> float:
+        """
+        Compute the rate from a string
+
+        Parameters:
+        - rate: string
+            Rate of the transition
+
+        Returns:
+        float
+            The computed rate
+        """
+        import logging
+
+        if param in self.keys():
+            if isinstance(super().__getitem__(param), (int, float)):
+                return super().__getitem__(param)
+            else:
+                try:
+                    return eval(super().__getitem__(param), self.globals, self)
+                except Exception as e:
+                    logging.error(f"Error computing parameter {param}: {e}")
+                    return None
+
 def get_cache_directory():
     """
     Return the location of the cache directory for the current platform.