mv to use np.random.default_rng

Author: bmtgoncalves

src/epidemik/EpiModel.py

@@ -6,6 +6,8 @@
 from typing import Dict, List, Set, Union
 import warnings
 import string
+import time
+import os
 
 import networkx as nx
 import numpy as np
@@ -22,7 +24,7 @@ class EpiModel(object):
     
         Provides a way to implement and numerically integrate 
     """
-    def __init__(self, compartments=None):
+    def __init__(self, compartments=None, seed=None, rng=None):
         """
         Initialize the EpiModel object
         
@@ -38,7 +40,15 @@ def __init__(self, compartments=None):
         self.population = None
         self.orig_comps = None
         self.demographics = False
-        
+
+        if seed is None:
+            seed = int(time.time()) + os.getpid()
+
+        if rng is 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])
 
@@ -397,11 +407,11 @@ def simulate(self, timesteps: int, t_min: int = 1, seasonality: Union[np.ndarray
                     comp_id = pos[comp]
 
                     if "birth" in data:
-                        births = np.random.binomial(pop[comp_id], data["birth"])
+                        births = self.rng.binomial(pop[comp_id], data["birth"])
                         new_pop[comp_id] += births
 
                     if "death" in data:
-                        deaths = np.random.binomial(pop[comp_id], data["death"])
+                        deaths = self.rng.binomial(pop[comp_id], data["death"])
                         new_pop[comp_id] -= deaths
 
             values.append(new_pop)
@@ -438,7 +448,7 @@ def integrate(self, timesteps: int , t_min: int = 1, seasonality: Union[np.ndarr
             else:
                 total_pop = self.population.sum()
                 p = np.copy(self.population)/total_pop
-                n = np.random.multinomial(kwargs[comp], p, 1)[0]
+                n = self.rng.multinomial(kwargs[comp], p, 1)[0]
 
                 for i, age in enumerate(string.ascii_lowercase[:len(p)]):
                     comp_age = comp + '_' + age

src/epidemik/MetaEpiModel.py

@@ -22,7 +22,7 @@ class MetaEpiModel:
     
         Provides a way to implement and numerically integrate 
     """
-    def __init__(self, travel_graph, populations, population='Population'):
+    def __init__(self, travel_graph, populations, population='Population', seed=None):
         """
         Initialize the EpiModel object
         
@@ -43,8 +43,13 @@ def __init__(self, travel_graph, populations, population='Population'):
         self.prototype = None
         self.seasonality = None
 
+        if seed is None:
+            seed = int(time.time()) + os.getpid()
+
+        self.rng = np.random.default_rng(seed=seed)
+
         for state in travel_graph.index:
-            models[state] = EpiModel() 
+            models[state] = EpiModel(rng=self.rng) 
             if self.transitions is None:
                 self.transitions = models[state].transitions
                 self.prototype = models[state] 
@@ -153,7 +158,7 @@ def _run_travel(self, compartments_, travel):
         def travel_step(x, populations):
             n = populations.loc[x.name]
             p = travel.loc[x.name].values.tolist()
-            output = np.random.multinomial(n, p)
+            output = self.rng.multinomial(n, p)
 
             return pd.Series(output, index=travel.columns)
 

src/epidemik/NetworkEpiModel.py

@@ -81,21 +81,21 @@ def simulate(self, timesteps, seeds, **kwargs):
                 continue
 
             current_active = list(active_nodes)
-            np.random.shuffle(current_active)
+            self.rng.shuffle(current_active)
 
             for node_i in current_active:
                 state_i = population[t-1, node_i]
 
                 if state_i in infections:
                     # contact each neighbour to see if we infect them
                     NN = list(self.network.neighbors(node_i))
-                    np.random.shuffle(NN)
+                    self.rng.shuffle(NN)
 
                     for node_j in NN:
                         state_j = population[t-1, node_j]
 
                         if state_j in infections[state_i]:
-                            prob = np.random.random()
+                            prob = self.rng.random()
 
                             if prob < infections[state_i][state_j]['rate']:
                                 new_state = infections[state_i][state_j]['target']