Checking in progress - have implemented Kermack-Mckendrick testing in…

… the SIR model.

Weird bug arising - cumulative incidence is consistently loewr than cumulative change in susceptibility.  Appears that susceptibility is tracking correctly, incidence is "dropping" infetions.  Checking things in so everyone can get a look at the code

src/laser_generic/infection.py

@@ -207,6 +207,9 @@ def nb_infection_update(count, itimers, susceptibility):  # pragma: no cover
         for i in nb.prange(count):
             itimer = itimers[i]
             if itimer > 0:
+                #if np.random.random_sample()<1/itimer:
+                #    susceptibility[i] = 1
+                #    itimers[i] = 0
                 itimer -= 1
                 itimers[i] = itimer
                 if itimer == 0:

src/laser_generic/susceptibility.py

@@ -56,6 +56,7 @@ def __init__(self, model, verbose: bool = False):
         self.model = model
 
         model.population.add_scalar_property("susceptibility", dtype=np.uint8, default=1)
+        model.patches.add_vector_property("susceptibility", model.params.nticks, dtype=np.uint32)
         # self.nb_initialize_susceptibility(model.population.count, model.population.dob, model.population.susceptibility)
 
         return
@@ -79,7 +80,7 @@ def nb_set_susceptibility(istart, iend, susceptibility, value) -> None:  # pragm
 
         return
 
-    # def __call__(self, model, tick):
+    def __call__(self, model, tick):
     #     """
     #     This method allows the instance to be called as a function.
 
@@ -93,6 +94,18 @@ def nb_set_susceptibility(istart, iend, susceptibility, value) -> None:  # pragm
     #         None
     #     """
 
+        patches = model.patches
+        population = model.population
+
+        susceptible_count = patches.susceptibility[tick, :]  # we will accumulate current susceptibles into this view into the susceptibility array
+        condition = population.susceptibility[0 : population.count] >0.0
+        if len(model.patches) == 1:
+            np.add(susceptible_count, np.sum(condition), out=susceptible_count)
+        else:
+            nodeids = population.nodeid[0 : population.count]
+            np.add.at(susceptible_count, nodeids[condition], 1)
+
+
     #     return
 
     def on_birth(self, model, _tick, istart, iend):

src/laser_generic/transmission.py

@@ -128,7 +128,9 @@ def __call__(self, model, tick) -> None:
             # interestingly, not actually faster than the nb_transmission_update_SI function
             susc_inds = np.where(condition == False)[0]
             ninf = np.random.binomial(len(susc_inds), forces)
-            population.susceptibility[np.random.choice(susc_inds, ninf, replace=False)] = 0
+            myinds = np.random.choice(susc_inds, ninf, replace=False)
+            population.susceptibility[myinds] = 0
+            population.itimer[myinds] = np.maximum(np.uint8(1), np.uint8(np.ceil(np.random.exponential(model.params.inf_mean))))
 
         elif hasattr(population, "etimer"):
             Transmission.nb_transmission_update_exposed(
@@ -203,9 +205,8 @@ def nb_transmission_update_noexposed(susceptibilities, nodeids, forces, itimers,
                 force = susceptibility * forces[nodeid]  # force of infection attenuated by personal susceptibility
                 if (force > 0) and (np.random.random_sample() < force):  # draw random number < force means infection
                     susceptibilities[i] = 0  # no longer susceptible
-                    # set exposure timer for newly infected individuals to a draw from a gamma distribution, must be at least 1 day
-                    itimers[i] = np.maximum(np.uint8(1), np.uint8(np.round(np.random.exponential(inf_mean))))
-
+                    # set infectious timer for the individual to an exponential draw
+                    itimers[i] = np.maximum(np.uint8(1), np.uint8(np.ceil(np.random.exponential(inf_mean))))
                     incidence[nodeid] += 1
 
         return

src/laser_generic/utils.py

@@ -145,7 +145,6 @@ def seed_infections_randomly(model, ninfections: int = 100) -> None:
 
     return
 
-
 # def seed_infections_in_patch(model, ipatch: int, ninfections: int = 100) -> None:
 #     """
 #     Seed initial infections in a specific patch of the population at the start of the simulation.
@@ -173,3 +172,25 @@ def seed_infections_randomly(model, ninfections: int = 100) -> None:
 #             cinfections += 1
 
 #     return
+
+def set_initial_susceptibility_randomly(model, susc_frac: float = 1.0) -> None:
+    """
+    Set the population susceptibility level at the start of the simulation.
+    This function randomly selects individuals from the population and changes
+    their susceptibility to zero, according to the parameter susc_frac.
+
+    Args:
+
+        model: The simulation model containing the population and parameters.
+        susc_frac (float, optional): The fraction of individuals to keep susceptible. 
+
+    Returns:
+
+        None
+    """
+
+    # Seed initial infections in random locations at the start of the simulation
+    indices = model.prng.choice(model.population.count, int(model.population.count * (1-susc_frac)), replace=False)
+    model.population.susceptibility[indices] = 0
+
+    return