Added MCV2. Moved RI and maternal immunity code to separate 'modules'.

nnmm/maternal_immunity.py

@@ -0,0 +1,17 @@
+import numpy as np
+import numba as nb
+
+def add( model, istart, iend ):
+    # enable this after adding susceptibility property to the population (see cells below)
+    model.population.susceptibility[istart:iend] = 0 # newborns have maternal immunity
+    model.population.susceptibility_timer[istart:iend] = int(0.5*365) # 6 months
+
+# Define the function to decrement susceptibility_timer and update susceptibility
+@nb.njit((nb.uint32, nb.uint8[:], nb.uint8[:]), parallel=True)
+def _update_susceptibility_based_on_sus_timer(count, susceptibility_timer, susceptibility):
+    for i in nb.prange(count):
+        if susceptibility_timer[i] > 0:
+            susceptibility_timer[i] -= 1
+            if susceptibility_timer[i] == 0:
+                susceptibility[i] = 1
+

nnmm/measles.py

@@ -176,6 +176,8 @@ def propagate_population(model, tick):
 model.nodes.add_vector_property("births", (model.params.ticks + 364) // 365)    # births per year
 
 # Adding ri_timer here since it's referred to in do_births.
+import ri
+import maternal_immunity as mi
 model.population.add_scalar_property("ri_timer", np.uint16)
 
 def do_births(model, tick):
@@ -195,17 +197,7 @@ def do_births(model, tick):
     model.population.dod[istart:iend] = pdsod(model.population.dob[istart:iend], max_year=100)   # make use of the fact that dob[istart:iend] is currently 0
     model.population.dob[istart:iend] = tick    # now update dob to reflect being born today
 
-    # Randomly set ri_timer for coverage fraction of agents to a value between 8.5*30.5 and 9.5*30.5 days
-    # change these numbers or parameterize as needed
-    ri_timer_values = np.random.uniform(8.5 * 30.5, 9.5 * 30.5, count_births).astype(np.uint16)
-
-    # Create a mask to select coverage fraction of agents
-    # Do coverage by node, not same for every node
-    # I don't think agents have node ids yet?
-    mask = np.random.rand(count_births) < (model.nodes.ri_coverages[model.population.nodeid[istart:iend]])
-
-    # Set ri_timer values for the selected agents
-    model.population.ri_timer[istart:iend][mask] = ri_timer_values[mask]
+    ri.add( model, count_births, istart, iend )
 
     index = istart
     nodeids = model.population.nodeid   # grab this once for efficiency
@@ -220,9 +212,7 @@ def do_births(model, tick):
         index += births
     model.nodes.population[:,tick+1] += todays_births
 
-    # enable this after adding susceptibility property to the population (see cells below)
-    model.population.susceptibility[istart:iend] = 0 # newborns have maternal immunity
-    model.population.susceptibility_timer[istart:iend] = int(0.5*365) # 6 months
+    mi.add( model, istart, iend )
 
     return
 
@@ -429,36 +419,13 @@ def initialize_susceptibility(count, dob, susceptibility):
 # In[555]:
 
 
-# Define the function to decrement ri_timer and update susceptibility
-@nb.njit((nb.uint32, nb.uint16[:], nb.uint8[:]), parallel=True)
-def _update_susceptibility_based_on_ri_timer(count, ri_timer, susceptibility):
-    for i in nb.prange(count):
-        if ri_timer[i] > 0:
-            ri_timer[i] -= 1
-            # TBD: It's perfectly possible that the individual got infected (or recovered) while this timer
-            # was counting down and we might want to abort the timer.
-            if ri_timer[i] == 0:
-                susceptibility[i] = 0
-
-# Example usage
-#update_susceptibility_based_on_ri_timer(model.population.count, model.population.ri_timer, model.population.susceptibility)
-
 def do_ri(model, tick):
-    _update_susceptibility_based_on_ri_timer(model.population.count, model.population.ri_timer, model.population.susceptibility)
+    ri._update_susceptibility_based_on_ri_timer(model.population.count, model.population.ri_timer, model.population.susceptibility)
     return
 
 
-# Define the function to decrement susceptibility_timer and update susceptibility
-@nb.njit((nb.uint32, nb.uint8[:], nb.uint8[:]), parallel=True)
-def _update_susceptibility_based_on_sus_timer(count, susceptibility_timer, susceptibility):
-    for i in nb.prange(count):
-        if susceptibility_timer[i] > 0:
-            susceptibility_timer[i] -= 1
-            if susceptibility_timer[i] == 0:
-                susceptibility[i] = 1
-
 def do_susceptibility_decay(model, tick):
-    _update_susceptibility_based_on_sus_timer(model.population.count, model.population.susceptibility_timer, model.population.susceptibility)
+    mi._update_susceptibility_based_on_sus_timer(model.population.count, model.population.susceptibility_timer, model.population.susceptibility)
     return
 
 

nnmm/ri.py

@@ -0,0 +1,42 @@
+import numpy as np
+import numba as nb
+
+def add( model, count_births, istart, iend ):
+    # Randomly set ri_timer for coverage fraction of agents to a value between 8.5*30.5 and 9.5*30.5 days
+    # change these numbers or parameterize as needed
+    ri_timer_values = np.random.uniform(8.5 * 30.5, 9.5 * 30.5, count_births).astype(np.uint16) # 9mo-ish
+    # Randomly set ri_timer for the second coverage fraction to a different range
+    ri_timer_values2 = np.random.uniform(14.5 * 30.5, 15.5 * 30.5, count_births).astype(np.uint16) # 15mo-ish
+
+    mask = np.random.rand(count_births) < (model.nodes.ri_coverages[model.population.nodeid[istart:iend]])
+
+    # Create a mask to select coverage fraction of agents
+    # Do coverage by node, not same for every node
+    # I don't think agents have node ids yet?
+    # Subdivision mask
+    subdivision_rand = np.random.rand(mask.sum())
+
+    # Create the three groups based on subdivision
+    group_85 = subdivision_rand < 0.85
+    group_14_25 = (subdivision_rand >= 0.85) & (subdivision_rand < 0.9925)
+    group_0_75 = subdivision_rand >= 0.9925
+
+    #import pdb
+    #pdb.set_trace()
+    # Apply ri_timer_values to 85% of the selected agents
+    model.population.ri_timer[istart:iend][mask][group_85] = ri_timer_values[mask][group_85]
+
+    # Apply ri_timer_values2 to 14.25% of the selected agents
+    model.population.ri_timer[istart:iend][mask][group_14_25] = ri_timer_values2[mask][group_14_25]
+
+# Define the function to decrement ri_timer and update susceptibility
+@nb.njit((nb.uint32, nb.uint16[:], nb.uint8[:]), parallel=True)
+def _update_susceptibility_based_on_ri_timer(count, ri_timer, susceptibility):
+    for i in nb.prange(count):
+        if ri_timer[i] > 0:
+            ri_timer[i] -= 1
+            # TBD: It's perfectly possible that the individual got infected (or recovered) while this timer
+            # was counting down and we might want to abort the timer.
+            if ri_timer[i] == 0:
+                susceptibility[i] = 0
+