working viral model

Author: bmtgoncalves

poetry.lock

@@ -22,6 +22,7 @@ dependencies = [
  "tqdm>=4",
  "pyyaml (>=6.0.2,<7.0.0)",
  "seaborn (>=0.13.2,<0.14.0)",
+ "scikit-learn (>=1.6.1,<2.0.0)",
 ]
 [project.urls]
 Homepage = "https://github.com/DataForScience/epidemik"

pyproject.toml

@@ -67,6 +67,7 @@ def add_interaction(
         target: str,
         agent: str,
         rate: Union[None, float, str] = None,
+        norm=True,
         **rates,
     ) -> None:
         """
@@ -79,7 +80,11 @@ def add_interaction(
             Name of the target compartment
         - agent: string
             Name of the agent
-        - params: string
+        - rate: float, str, None
+            Rate of the interaction
+        - norm: bool
+            Whether to normalize the transition rate or not
+        - rates:
             Named parameters for the interaction
 
         Returns:
@@ -95,7 +100,10 @@ def add_interaction(
             rates = list(rates.keys())
             rate_key = rates[0]
 
-        self.transitions.add_edge(source, target, agent=agent, rate=rate_key)
+        if agent not in self.transitions.nodes:
+            self.transitions.add_node(agent)
+    
+        self.transitions.add_edge(source, target, agent=agent, rate=rate_key, norm=norm)
 
     def add_spontaneous(
         self, source: str, target: str, rate: Union[None, float, str] = None, **rates
@@ -162,14 +170,19 @@ def add_viral_generation(
             target_rate = rates[1]
 
         self.transitions.add_edge(
-            source, target, rate=source_rate
+            source, target, rate=source_rate, viral_source=True, viral_target=False,
         )
         self.transitions.add_edge(
-            source, target, rate=target_rate
+            source, target, rate=target_rate, viral_source=False, viral_target=True
         )
 
     def add_birth_rate(
-        self, comps: Union[List, None] = None, rate: Union[float, None] = None, fixed=False, **rates
+        self, 
+        comps: Union[List, None] = None, 
+        rate: Union[float, None] = None, 
+        fixed=False, 
+        global_rate=True,
+        **rates
     ) -> None:
         """
         Add a birth rate to one or more compartments
@@ -200,9 +213,13 @@ def add_birth_rate(
 
             self.transitions.nodes[comp]["birth"] = rate_key
             self.transitions.nodes[comp]["fixed"] = fixed
+            self.transitions.nodes[comp]["global"] = global_rate
 
     def add_death_rate(
-        self, comps: Union[List, None] = None, rate: Union[None, float] = None, **rates
+        self, 
+        comps: Union[List, None] = None, 
+        rate: Union[None, float] = None, 
+        **rates
     ) -> None:
         """
         Add a birth rate to one or more compartments
@@ -327,10 +344,10 @@ def _new_cases(self, time: float, population: np.ndarray, pos: Dict) -> np.ndarr
         Internal function used by integration routine
 
         Parameters:
-        - population: numpy array
-            Current population of each compartment
         - time: float
             Current time
+        - population: numpy array
+            Current population of each compartment
         - pos: dict
             Dictionary mapping compartment names to indices
 
@@ -364,36 +381,45 @@ def _new_cases(self, time: float, population: np.ndarray, pos: Dict) -> np.ndarr
 
             if "agent" in trans:
                 agent = trans["agent"]
+                rate *= population[pos[agent]]
 
-                if self.population is None:
-                    rate *= population[pos[agent]] / N
-                else:
-                    rate *= population[pos[agent]] / N[agent]
-
+                if trans["norm"]:
+                    if self.population is None:
+                        rate /=  N
+                    else:
+                        rate /= N[agent]
+                
                 if self.seasonality is not None:
                     curr_t = int(time) % 365
                     season = float(self.seasonality[curr_t])
                     rate *= season
 
-            diff[pos[source]] -= rate
-            diff[pos[target]] += rate
+            if "viral_source" not in trans or trans["viral_source"]:
+                diff[pos[source]] -= rate
+            # Make sure viral generations are asymetric
+            if "viral_target" not in trans or trans["viral_target"]:
+                diff[pos[target]] += rate
 
-            # Population dynamics
-            if self.demographics:
-                for comp, data in self.transitions.nodes(data=True):
-                    comp_id = pos[comp]
+        # Population dynamics
+        if self.demographics:
+            for comp, data in self.transitions.nodes(data=True):
+                comp_id = pos[comp]
 
-                    if "birth" in data:
-                        if "fixed" in data:
-                            births = self.params[data["birth"]]
+                if "birth" in data:
+                    if "fixed" in data and data["fixed"]:
+                        births = self.params[data["birth"]]
+                    else:
+                        if data["global"]:
+                            total_population = population.sum()
+                            births = total_population * 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] * self.params[data["death"]]
-                        diff[comp_id] -= deaths
+                    diff[comp_id] += births
+
+                if "death" in data:
+                    deaths = population[comp_id] * self.params[data["death"]]
+                    diff[comp_id] -= deaths
 
         return diff
 
@@ -616,7 +642,7 @@ def integrate(
 
                     population[pos[comp_age]] = n[i]
 
-        time = np.arange(t_min, t_min + timesteps, 1)
+        time = np.arange(t_min, t_min + timesteps)
 
         self.seasonality = seasonality
         values = pd.DataFrame(

src/epidemik/EpiModel.py

@@ -1,16 +1,33 @@
 import unittest
 from epidemik import EpiModel
-import logging
-
+from sklearn.linear_model import LinearRegression
+import numpy as np
 
 class EpiModelTestCase(unittest.TestCase):
     def setUp(self):
-        self.SIR = EpiModel()
+
         self.beta = 0.3
         self.mu = 0.1
+        self.birth = 0.3
+        self.death = 0.3
+        self.fixed_birth = 10
+
+        self.SIR = EpiModel()
         self.SIR.add_interaction("S", "I", "I", beta=self.beta)
         self.SIR.add_spontaneous("I", "R", mu=self.mu)
 
+        self.birth_test = EpiModel()
+        self.birth_test.add_interaction("S", "I", "I", beta=self.beta)
+        self.birth_test.add_birth_rate("S", b=self.birth)
+
+        self.fixed_birth_test = EpiModel()
+        self.fixed_birth_test.add_interaction("S", "I", "I", beta=self.beta)
+        self.fixed_birth_test.add_birth_rate("S", b=self.fixed_birth, fixed=True, global_rate=False)
+
+        self.death_test = EpiModel()
+        self.death_test.add_interaction("S", "I", "I", beta=self.beta)
+        self.death_test.add_death_rate(d=self.death)
+
     def test_R0(self):
         self.assertEqual(self.SIR.R0(), 3.0, "incorrect R0")
 
@@ -32,3 +49,48 @@ def test_edges(self):
                 self.assertEqual(self.SIR.params["mu"], self.mu)
                 self.assertEqual(edge[0], "I")
                 self.assertEqual(edge[1], "R")
+
+    def test_birth(self):
+        self.assertEqual(self.birth_test.transitions.nodes['S']["birth"], "b")
+        self.assertIn("b", self.birth_test.params.keys())
+        self.assertEqual(self.birth_test.params["b"], self.birth)
+
+    def test_birth_rate(self):
+        self.birth_test.integrate(10, S=990, I=10)
+        values = self.birth_test.values_
+        values['total'] = values.sum(axis=1)
+        values = values.reset_index()
+
+        lm = LinearRegression()
+        lm.fit(values['index'].values.reshape(-1, 1), np.log(values['total']))
+
+        self.assertAlmostEqual(lm.coef_[0], self.birth, delta=0.01)
+
+    def test_fixed_birth_rate(self):
+        self.fixed_birth_test.integrate(10, S=990, I=10)
+        values = self.fixed_birth_test.values_
+        values['total'] = values.sum(axis=1)
+        values = values.reset_index()
+
+        lm = LinearRegression()
+        lm.fit(values['index'].values.reshape(-1, 1), values['total'])
+
+        self.assertAlmostEqual(lm.coef_[0], 10, delta=0.01)
+
+    def test_death(self):
+        self.assertEqual(self.death_test.transitions.nodes['S']["death"], "d")
+        self.assertEqual(self.death_test.transitions.nodes['I']["death"], "d")
+
+        self.assertIn("d", self.death_test.params.keys())
+        self.assertEqual(self.death_test.params["d"], self.death)
+
+    def test_death_rate(self):
+        self.death_test.integrate(10, S=990, I=10)
+        values = self.death_test.values_
+        values['total'] = values.sum(axis=1)
+        values = values.reset_index()
+
+        lm = LinearRegression()
+        lm.fit(values['index'].values.reshape(-1, 1), np.log(values['total']))
+
+        self.assertAlmostEqual(lm.coef_[0], -self.death, delta=0.01)

tests/test_EpiModel.py

@@ -46,3 +46,5 @@ def test_travel(self):
     def test_integrate(self):
         with self.assertRaises(NotImplementedError) as _:
             self.SIR.integrate()
+
+

tests/test_MetaEpiModel.py

@@ -0,0 +1,23 @@
+import unittest
+from  epidemik import EpiModel, NetworkEpiModel
+import networkx as nx
+
+class NetworkEpiModelTestCase(unittest.TestCase):
+    def setUp(self):
+        self.N = 300
+        self.G_full = nx.erdos_renyi_graph(self.N, p=1.)
+        self.beta = 0.05
+        self.SI_full = NetworkEpiModel(self.G_full)
+    
+    def test_named_parameters(self):
+        self.SI_full.add_interaction("S", "I", "I", beta=self.beta)
+        self.assertIn("beta",
+            self.SI_full.params,
+            "The parameter beta should be in the params dictionary")
+
+    def test_unnamed_parameters(self):
+        self.SI_full.add_interaction("S", "I", "I", self.beta)
+        self.assertIn("rate1",
+            self.SI_full.params,
+            "The parameter rate1 should be in the params dictionary")
+