Merge pull request #253 from SABS-R3-Epidemiology/waning-immunity

Initial waning immunity implementation

pyEpiabm/pyEpiabm/core/person.py

@@ -51,6 +51,7 @@ def __init__(self, microcell, age_group=None):
         self.next_infection_status = None
         self.time_of_status_change = None
         self.infection_start_time = None
+        self.time_of_recovery = None
         self.care_home_resident = False
         self.key_worker = False
         self.date_positive = None
@@ -258,3 +259,14 @@ def set_id(self, id: str):
             raise ValueError(f"Duplicate id: {id}.")
 
         self.id = id
+
+    def set_time_of_recovery(self, time: float):
+        """Records the time at which a person enters the Recovered compartment.
+
+        Parameters
+        ----------
+        time : float
+            Current simulation time
+
+        """
+        self.time_of_recovery = time

pyEpiabm/pyEpiabm/routine/simulation.py

@@ -43,6 +43,8 @@ def configure(self,
             * `initial_infected_number`: The initial number of infected \
                individuals in the population
             * `simulation_seed`:  Random seed for reproducible simulations
+            * `include_waning`: Boolean to determine whether immunity waning \
+                is included in the simulation
 
         file_params Contains:
             * `output_file`: String for the name of the output .csv file
@@ -104,6 +106,11 @@ def configure(self,
 
         Parameters.instance().use_ages = self.age_stratified
 
+        self.include_waning = sim_params["include_waning"] \
+            if "include_waning" in sim_params else False
+
+        Parameters.instance().use_waning_immunity = self.include_waning
+
         # If random seed is specified in parameters, set this in numpy
         if "simulation_seed" in self.sim_params:
             Simulation.set_random_seed(self.sim_params["simulation_seed"])

pyEpiabm/pyEpiabm/sweep/host_progression_sweep.py

@@ -43,16 +43,16 @@ def __init__(self):
 
         self.number_of_states = len(InfectionStatus)
         assert self.state_transition_matrix.shape == \
-            (self.number_of_states, self.number_of_states), \
-            'Matrix dimensions must match number of infection states'
+               (self.number_of_states, self.number_of_states), \
+               'Matrix dimensions must match number of infection states'
 
         # Instantiate transmission time matrix
         time_matrix_object = TransitionTimeMatrix()
-        self.transition_time_matrix =\
+        self.transition_time_matrix = \
             time_matrix_object.create_transition_time_matrix()
         # Instantiate parameters to be used in update transition time
         # method
-        self.latent_to_symptom_delay =\
+        self.latent_to_symptom_delay = \
             pe.Parameters.instance().latent_to_sympt_delay
         # Defining the length of the model time step (in days, can be a
         # fraction of day as well).
@@ -69,7 +69,7 @@ def __init__(self):
         # Extreme case where model time step would be too small
         max_inf_steps = 2550
         # Define number of time steps a person is infectious:
-        num_infectious_ts =\
+        num_infectious_ts = \
             int(np.ceil(infectious_period / self.model_time_step))
         if num_infectious_ts >= max_inf_steps:
             raise ValueError('Number of timesteps in infectious period exceeds'
@@ -86,11 +86,11 @@ def __init__(self):
             associated_inf_value = int(np.floor(t))
             t -= associated_inf_value
             if associated_inf_value < inf_prof_resolution:
-                infectiousness_prog[i] =\
+                infectiousness_prog[i] = \
                     (infectious_profile[associated_inf_value] * (1 - t)
                      + infectious_profile[associated_inf_value + 1] * t)
             else:  # limit case where we define infectiousness to 0
-                infectiousness_prog[i] =\
+                infectiousness_prog[i] = \
                     infectious_profile[inf_prof_resolution]
         # Scaling
         scaling_param = inf_prof_average
@@ -147,11 +147,14 @@ def update_next_infection_status(self, person: Person):
             Instance of person class with infection status attributes
 
         """
-        if person.infection_status in [InfectionStatus.Recovered,
-                                       InfectionStatus.Dead,
+        if person.infection_status in [InfectionStatus.Dead,
                                        InfectionStatus.Vaccinated]:
             person.next_infection_status = None
             return
+        elif (person.infection_status == InfectionStatus.Recovered and not
+              Parameters.instance().use_waning_immunity):
+            person.next_infection_status = None
+            return
         elif (person.care_home_resident and
               person.infection_status == InfectionStatus.InfectICU):
             person.next_infection_status = InfectionStatus.Dead
@@ -175,7 +178,7 @@ def update_next_infection_status(self, person: Person):
                                  ' probabilities')
 
         next_infection_status_number = random.choices(outcomes, weights)[0]
-        next_infection_status =\
+        next_infection_status = \
             InfectionStatus(next_infection_status_number)
 
         person.next_infection_status = next_infection_status
@@ -203,20 +206,30 @@ def update_time_status_change(self, person: Person, time: float):
         if person.infection_status == InfectionStatus.Susceptible:
             raise ValueError("Method should not be used to infect people")
 
-        if person.infection_status in [InfectionStatus.Recovered,
-                                       InfectionStatus.Dead,
+        if person.infection_status in [InfectionStatus.Dead,
                                        InfectionStatus.Vaccinated]:
             transition_time = np.inf
+        elif (person.infection_status == InfectionStatus.Recovered and not
+              Parameters.instance().use_waning_immunity):
+            transition_time = np.inf
         else:
             row_index = person.infection_status.name
             column_index = person.next_infection_status.name
-            transition_time_icdf_object =\
-                self.transition_time_matrix.loc[row_index, column_index]
             # Checks for susceptible to exposed case
             # where transition time is zero
             try:
-                transition_time =\
-                    transition_time_icdf_object.icdf_choose_noexp()
+                if person.infection_status != InfectionStatus.Recovered:
+                    transition_time_icdf_object = \
+                        self.transition_time_matrix.loc[row_index,
+                                                        column_index]
+                    transition_time = \
+                        transition_time_icdf_object.icdf_choose_noexp()
+                else:
+                    # If someone is recovered, then their transition time
+                    # will be equal to 1 when waning immunity is turned on.
+                    # This means that everyone spends exactly 1 day in the
+                    # Recovered compartment with waning immunity
+                    transition_time = 1
             except AttributeError as e:
                 if "object has no attribute 'icdf_choose_noexp'" in str(e):
                     transition_time = transition_time_icdf_object
@@ -300,13 +313,18 @@ def __call__(self, time: float):
                 while person.time_of_status_change <= time:
                     person.update_status(person.next_infection_status)
                     if person.infection_status in \
-                            [InfectionStatus.InfectASympt,
-                             InfectionStatus.InfectMild,
-                             InfectionStatus.InfectGP]:
+                        [InfectionStatus.InfectASympt,
+                         InfectionStatus.InfectMild,
+                         InfectionStatus.InfectGP]:
                         self.set_infectiousness(person, time)
                         if not person.is_symptomatic():
                             asympt_or_uninf_people.append((cell, person))
                     self.update_next_infection_status(person)
+                    if person.infection_status == InfectionStatus.Susceptible:
+                        person.time_of_status_change = None
+                        break
+                    elif person.infection_status == InfectionStatus.Recovered:
+                        person.set_time_of_recovery(time)
                     self.update_time_status_change(person, time)
                     self.sympt_testing_queue(cell, person)
                 self._updates_infectiousness(person, time)
@@ -331,9 +349,9 @@ def sympt_testing_queue(self, cell, person: Person):
 
         """
         if hasattr(Parameters.instance(), 'intervention_params'):
-            if 'disease_testing' in Parameters.instance().\
+            if 'disease_testing' in Parameters.instance(). \
               intervention_params.keys():
-                testing_params = Parameters.instance().\
+                testing_params = Parameters.instance(). \
                     intervention_params['disease_testing']
                 r = random.random()
                 type_r = random.random()
@@ -380,9 +398,9 @@ def asympt_uninf_testing_queue(self, person_list: list, time):
 
         """
         if hasattr(Parameters.instance(), 'intervention_params'):
-            if 'disease_testing' in Parameters.instance().\
-               intervention_params.keys():
-                testing_params = Parameters.instance().\
+            if 'disease_testing' in Parameters.instance(). \
+              intervention_params.keys():
+                testing_params = Parameters.instance(). \
                     intervention_params['disease_testing']
                 for item in person_list:
                     cell = item[0]

pyEpiabm/pyEpiabm/sweep/initial_demographics_sweep.py

@@ -55,7 +55,7 @@ def __init__(self, dem_file_params: typing.Dict):
         file_name = "demographics.csv"
         self.titles = ["id"]
         if self.age_output:
-            self.titles.append("age_group")
+            self.titles.append("age")
         if self.spatial_output:
             self.titles += ["location_x", "location_y"]
         self.titles.append("kw_or_chr")
@@ -78,7 +78,7 @@ def __call__(self, *args):
         for cell in self._population.cells:
             for person in cell.persons:
                 data = {"id": person.id,
-                        "age_group": person.age_group
+                        "age": person.age
                         if self.age_output else None,
                         "location_x": cell.location[0]
                         if self.spatial_output else None,

pyEpiabm/pyEpiabm/sweep/transition_matrices.py

@@ -91,7 +91,10 @@ def create_state_transition_matrix(self, coeff: typing.
                                                           "to_icurecov"]
         matrix.loc['InfectICU', 'Dead'] = coeff["prob_icu_to_death"]
         matrix.loc['InfectICURecov', 'Recovered'] = 1
-        matrix.loc['Recovered', 'Recovered'] = 1
+        if pe.core.Parameters.instance().use_waning_immunity:
+            matrix.loc['Recovered', 'Susceptible'] = 1
+        else:
+            matrix.loc['Recovered', 'Recovered'] = 1
         matrix.loc['Dead', 'Dead'] = 1
         matrix.loc['Vaccinated', 'Vaccinated'] = 1
 

pyEpiabm/pyEpiabm/tests/test_func/test_sim_functional.py

@@ -26,7 +26,8 @@ def setUp(self) -> None:
                            "place_number": 2}
         self.sim_params = {"simulation_start_time": 0,
                            "simulation_end_time": 100,
-                           "initial_infected_number": 0}
+                           "initial_infected_number": 0,
+                           "include_waning": False}
 
         self.file_params = {"output_file": "output.csv",
                             "output_dir": "test_folder/integration_tests",
@@ -153,6 +154,32 @@ def test_total_infection(self, *mocks):
                               "test_folder/integration_tests")
         mocks[2].assert_called_with(folder)  # Mock for mkdir()
 
+    def test_waning_compartments(self, *mocks):
+        """Basic functional test to ensure everyone is not recovered at the
+        end of the simulation.
+        """
+        with patch('pyEpiabm.Parameters.instance') as mock_param:
+            mock_param.return_value.use_waning_immunity = 1.0
+            mock_param.return_value.asympt_infect_period = 14
+            mock_param.return_value.time_steps_per_day = 1
+            self.sim_params["initial_infected_number"] = 5
+            self.sim_params["include_waning"] = True
+            pop = TestSimFunctional.toy_simulation(self.pop_params,
+                                                   self.sim_params,
+                                                   self.file_params)
+
+            # Test not all individuals have entered Recovered or Dead at end
+            recov_dead_state_count = 0
+
+            for cell in pop.cells:
+                cell_data = cell.compartment_counter.retrieve()
+                for status in [InfectionStatus.Recovered,
+                               InfectionStatus.Dead]:
+                    recov_dead_state_count += cell_data[status]
+
+            self.assertNotEqual(np.sum(recov_dead_state_count),
+                                self.pop_params["population_size"])
+
     def test_no_infection(self, *mocks):
         """Basic functional test to ensure noone is infected when there are
         no initial cases in the entire population

pyEpiabm/pyEpiabm/tests/test_unit/test_core/test_person.py

@@ -127,6 +127,10 @@ def test_set_id(self):
         self.microcell.persons.append(new_person)
         self.assertRaises(ValueError, new_person.set_id, "1.1.1.1")
 
+    def test_set_time_of_recovery(self):
+        self.person.set_time_of_recovery(time=5)
+        self.assertTrue(self.person.time_of_recovery, 5)
+
 
 if __name__ == '__main__':
     unittest.main()

pyEpiabm/pyEpiabm/tests/test_unit/test_routine/test_simulation.py

@@ -23,7 +23,8 @@ def setUpClass(cls) -> None:
         pe.Parameters.instance().time_steps_per_day = 1
         cls.sim_params = {"simulation_start_time": 0,
                           "simulation_end_time": 1,
-                          "initial_infected_number": 0}
+                          "initial_infected_number": 0,
+                          "include_waning": True}
 
         cls.mock_output_dir = "pyEpiabm/pyEpiabm/tests/test_output/mock"
         cls.file_params = {"output_file": "test_file.csv",
@@ -67,6 +68,7 @@ def test_configure(self, mock_mkdir):
             self.assertEqual(test_sim.infectiousness_output, False)
             self.assertEqual(test_sim.ih_status_writer, None)
             self.assertEqual(test_sim.ih_infectiousness_writer, None)
+            self.assertEqual(test_sim.include_waning, True)
             self.assertEqual(test_sim.compress, False)
 
             del test_sim.writer