move to named parameters

Author: bmtgoncalves

src/epidemik/EpiModel.py

@@ -40,6 +40,7 @@ def __init__(self, compartments=None, seed=None, rng=None):
         self.population = None
         self.orig_comps = None
         self.demographics = False
+        self.params = {}
 
         if seed is None:
             seed = int(time.time()) + os.getpid()
@@ -48,11 +49,15 @@ def __init__(self, compartments=None, seed=None, rng=None):
             self.rng = np.random.default_rng(seed=seed)
         else:
             self.rng = rng
-                
+
         if compartments is not None:
             self.transitions.add_nodes_from([comp for comp in compartments])
 
-    def add_interaction(self, source: str, target: str, agent: str, rate: float) -> None:  
+    def add_interaction(self, 
+                        source: str, 
+                        target: str, 
+                        agent: str, 
+                        **rates) -> None:  
         """
         Add an interaction between two compartments
         
@@ -63,15 +68,19 @@ def add_interaction(self, source: str, target: str, agent: str, rate: float) ->
             Name of the target compartment
         - agent: string
             Name of the agent
-        - rate: float
-            Rate of the interaction
+        - params: string
+            Named parameters for the interaction
         
         Returns:
         None
-        """      
-        self.transitions.add_edge(source, target, agent=agent, rate=rate)        
+        """
+
+        self.params.update(rates)
+        rates = list(rates.keys())
+
+        self.transitions.add_edge(source, target, agent=agent, rate=rates[0])        
 
-    def add_spontaneous(self, source: str, target: str, rate: float) -> None:
+    def add_spontaneous(self, source: str, target: str, **rates) -> None:
         """
         Add a spontaneous transition between two compartments
         
@@ -86,7 +95,11 @@ def add_spontaneous(self, source: str, target: str, rate: float) -> None:
         Returns:
         None
         """
-        self.transitions.add_edge(source, target, rate=rate)
+
+        self.params.update(rates)
+        rates = list(rates.keys())
+
+        self.transitions.add_edge(source, target, rate=rates[0])
 
     def add_birth_rate(self, rate: float, comps: Union[List, None] = None) -> None:
         """
@@ -224,7 +237,8 @@ def _new_cases(self, time: float, population: np.ndarray,  pos: Dict) -> np.ndar
             target = edge[1]
             trans = edge[2]
 
-            rate = trans['rate']*population[pos[source]]
+            rate_val = self.params[trans['rate']]
+            rate = rate_val*population[pos[source]]
 
             if 'start' in trans and trans['start'] >= time:
                 continue
@@ -370,7 +384,7 @@ def simulate(self, timesteps: int, t_min: int = 1, seasonality: Union[np.ndarray
                     source = pos[comp]
                     target = pos[node_j]
 
-                    rate = data['rate']
+                    rate = self.params[data['rate']]
 
                     if 'start' in data and data['start'] >= t:
                         continue
@@ -505,21 +519,24 @@ def __repr__(self) -> str:
               (self.transitions.number_of_nodes(), 
                self.transitions.number_of_edges())
 
+        text += "Parameters:\n"
+        for rate, value in self.params.items():
+            text += "%s= %f\n" % (rate, value)
+        text += "\n\nTransitions:\n"
+
         for edge in self.transitions.edges(data=True):
             source = edge[0]
             target = edge[1]
             trans = edge[2]
-            
-            rate = trans['rate']
-
+        
             if 'agent' in trans:
                 agent = trans['agent']
-                text += "%s + %s = %s %f\n" % (source, agent, target, rate)
+                text += "%s + %s = %s %s\n" % (source, agent, target, trans['rate'])
             elif 'start' in trans:
                 start = trans['start']
-                text+="%s -> %s %f starting at %s days\n" % (source, target, rate, start)
+                text+="%s -> %s %s starting at %s days\n" % (source, target, rate, start)
             else:
-                text+="%s -> %s %f\n" % (source, target, rate)
+                text+="%s -> %s %s\n" % (source, target, rate)
 
         R0 = self.R0()
 
@@ -664,7 +681,7 @@ def R0(self) -> Union[float, None]:
 
         try:
             for node_i, node_j, data in self.transitions.edges(data=True):
-                rate = data['rate']
+                rate = self.params[data['rate']]
 
                 if "agent" in data:
                     target = pos[node_j]

src/epidemik/MetaEpiModel.py

@@ -70,7 +70,7 @@ def __repr__(self):
         text = "Metapopulation model with %u populations\n\nThe disease is defined by an %s" % (self.travel_graph.shape[0], model_text)
         return text 
 
-    def add_interaction(self, source, target, agent, rate):  
+    def add_interaction(self, source, target, agent, **rates):  
         """
         Add an interaction between two compartments_
         
@@ -88,9 +88,9 @@ def add_interaction(self, source, target, agent, rate):
         None
         """    
         for state in self.models:  
-            self.models[state].add_interaction(source, target, agent, rate)       
+            self.models[state].add_interaction(source, target, agent,  **rates)       
 
-    def add_spontaneous(self, source, target, rate):
+    def add_spontaneous(self, source, target,  **rates):
         """
         Add a spontaneous transition between two compartments_
         
@@ -106,9 +106,9 @@ def add_spontaneous(self, source, target, rate):
         None
         """
         for state in self.models:  
-            self.models[state].add_spontaneous(source, target, rate)      
+            self.models[state].add_spontaneous(source, target,  **rates)      
 
-    def add_vaccination(self, source, target, rate, start):
+    def add_vaccination(self, source, target, start,  **rates):
         """
         Add a vaccination transition between two compartments_
         
@@ -126,7 +126,7 @@ def add_vaccination(self, source, target, rate, start):
         None
         """
         for state in self.models:  
-            self.models[state].add_vaccination(source, target, rate, start)
+            self.models[state].add_vaccination(source, target, start, **rates)
 
     def R0(self):
         key = list(self.models.keys())[0]

src/epidemik/NetworkEpiModel.py

@@ -20,14 +20,17 @@ def __init__(self, network, compartments=None):
         self.kavg_ = 2*network.number_of_edges()/network.number_of_nodes() 
         self.spontaneous = {}
         self.interactions = {}
+        self.params = {}
 
     def integrate(self, timesteps, **kwargs):
         raise NotImplementedError("Network Models don't support numerical integration")
 
-    def add_interaction(self, source, target, agent, rate, rescale=False):
+    def add_interaction(self, source, target, agent, rescale=False, **rates):
         if rescale:
             rate /= self.kavg_
 
+        self.params.update(rates)
+        rate = list(rates.keys())[0]
         super(NetworkEpiModel, self).add_interaction(source, target, agent=agent, rate=rate)
 
         if source not in self.interactions:
@@ -36,9 +39,13 @@ def add_interaction(self, source, target, agent, rate, rescale=False):
         if target not in self.interactions[source]:
             self.interactions[source] = {}
 
+
         self.interactions[source][agent] = {'target': target, 'rate': rate}
 
-    def add_spontaneous(self, source, target, rate):
+    def add_spontaneous(self, source, target, **rates):
+        self.params.update(rates)
+        rate = list(rates.keys())[0]
+
         super(NetworkEpiModel, self).add_spontaneous(source, target, rate=rate)
         if source not in self.spontaneous:
             self.spontaneous[source] = {}
@@ -97,7 +104,8 @@ def simulate(self, timesteps, seeds, **kwargs):
                         if state_j in infections[state_i]:
                             prob = self.rng.random()
 
-                            if prob < infections[state_i][state_j]['rate']:
+                            rate = self.params[infections[state_i][state_j]['rate']]
+                            if prob < rate:
                                 new_state = infections[state_i][state_j]['target']
                                 population[t, node_j] = new_state
 
@@ -110,7 +118,8 @@ def simulate(self, timesteps, seeds, **kwargs):
                     prob = np.zeros(len(pos))
 
                     for target in self.spontaneous[state_i]:
-                        prob[pos[target]] = self.spontaneous[state_i][target]
+                        rate = self.params[self.spontaneous[state_i]['target']]
+                        prob[pos[target]] = rate
 
                     prob[pos[state_i]] = 1-np.sum(prob)
 

tests/tests_EpiModel.py

@@ -1,13 +1,14 @@
 import unittest
 from epidemik import EpiModel
+import logging
 
 class EpiModelTestCase(unittest.TestCase):
 	def setUp(self):
 		self.SIR = EpiModel()
 		self.beta = 0.3
 		self.mu = 0.1
-		self.SIR.add_interaction('S', 'I', 'I', self.beta)
-		self.SIR.add_spontaneous('I', 'R', self.mu)
+		self.SIR.add_interaction('S', 'I', 'I', beta=self.beta)
+		self.SIR.add_spontaneous('I', 'R', mu=self.mu)
 
 	def test_R0(self):
 		self.assertEqual(self.SIR.R0(), 3.0, 'incorrect R0')
@@ -21,10 +22,12 @@ def test_edges(self):
 
 			if "agent" in data:
 				self.assertEqual(data['agent'], 'I')
-				self.assertEqual(data['rate'], self.beta)
+				self.assertEqual(data['rate'], 'beta')
+				self.assertEqual(self.SIR.params['beta'], self.beta)
 				self.assertEqual(edge[0], 'S')
 				self.assertEqual(edge[1], 'I')
 			else:
-				self.assertEqual(data['rate'], self.mu)
+				self.assertEqual(data['rate'], 'mu')
+				self.assertEqual(self.SIR.params['mu'], self.mu)
 				self.assertEqual(edge[0], 'I')
 				self.assertEqual(edge[1], 'R')