Merge remote-tracking branch 'origin/waning-immunity' into waning-imm…

…unity
SABS-R3-Epidemiology · Feb 26, 2024 · 68fb20a · 68fb20a
2 parents 7b4a8f4 + 3790ac0
commit 68fb20a
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 6 deletions.
diff --git a/pyEpiabm/pyEpiabm/output/_csv_dict_writer.py b/pyEpiabm/pyEpiabm/output/_csv_dict_writer.py
@@ -68,4 +68,5 @@ def compress(self):
         logging.info(f"Zip file created for {self.filename}")
         df = pd.read_csv(self.filepath)
         df.to_csv(output_filepath, index=False, compression={'method': 'zip'})
+        self.f.close()
         os.remove(self.filepath)
diff --git a/pyEpiabm/pyEpiabm/sweep/transition_matrices.py b/pyEpiabm/pyEpiabm/sweep/transition_matrices.py
@@ -116,7 +116,8 @@ def create_waning_transition_matrix(self, coeff: typing.
         probabilities to go from one state to another. For example, the element
         i,j in the matrix is the probability of someone with current infection
         status associated with the row i to move to the infection status
-        associated with the column j.
+        associated with the column j. Further explanation of this matrix can be
+        found at https://github.com/SABS-R3-Epidemiology/epiabm/wiki/Parameters#waning-infection-matrix
 
         Parameters
         ----------
@@ -128,7 +129,7 @@ def create_waning_transition_matrix(self, coeff: typing.
         pd.DataFrame
             Matrix in the form of a dataframe
 
-        """
+        """ # noqa
         waning_matrix =\
             StateTransitionMatrix.create_empty_state_transition_matrix()
         p = RateMultiplier(coeff["exposed_to_infect"][0],

diff --git a/pyEpiabm/pyEpiabm/tests/test_func/test_sim_functional.py b/pyEpiabm/pyEpiabm/tests/test_func/test_sim_functional.py
@@ -96,6 +96,7 @@ def test_population_conservation(self, *mocks):
         iter_num = (sim_params["simulation_end_time"]
                     - sim_params["simulation_start_time"] + 1)
 
+        pe.routine.Simulation.set_random_seed(seed=30)
         pop = TestSimFunctional.toy_simulation(pop_params,
                                                sim_params,
                                                self.file_params)
@@ -129,6 +130,7 @@ def test_total_infection(self, *mocks):
         """Basic functional test to ensure everyone is infected when the entire
         population is placed in one large household.
         """
+        pe.routine.Simulation.set_random_seed(seed=30)
         self.sim_params["initial_infected_number"] = 5
         pop = TestSimFunctional.toy_simulation(self.pop_params,
                                                self.sim_params,
@@ -158,6 +160,7 @@ def test_waning_compartments(self, *mocks):
         """Basic functional test to ensure everyone is not recovered at the
         end of the simulation.
         """
+        pe.routine.Simulation.set_random_seed(seed=30)
         self.sim_params["initial_infected_number"] = 5
         self.sim_params["include_waning"] = True
         pe.Parameters.instance().use_waning_immunity = 1.0
@@ -182,6 +185,7 @@ def test_no_infection(self, *mocks):
         """Basic functional test to ensure noone is infected when there are
         no initial cases in the entire population
         """
+        pe.routine.Simulation.set_random_seed(seed=30)
         pop = TestSimFunctional.toy_simulation(self.pop_params,
                                                self.sim_params,
                                                self.file_params)
@@ -211,6 +215,7 @@ def test_segmented_infection(self, *mocks):
         sweep_list = [pe.sweep.HouseholdSweep(), pe.sweep.QueueSweep(),
                       pe.sweep.HostProgressionSweep()]
 
+        pe.routine.Simulation.set_random_seed(seed=30)
         pe.Parameters.instance().household_size_distribution = []
         pop = TestSimFunctional.file_simulation("test_input.csv",
                                                 self.sim_params,
@@ -246,6 +251,7 @@ def test_small_cutoff(self, *mocks):
         sweep_list = [pe.sweep.HouseholdSweep(), pe.sweep.SpatialSweep(),
                       pe.sweep.QueueSweep(), pe.sweep.HostProgressionSweep()]
 
+        pe.routine.Simulation.set_random_seed(seed=30)
         pe.Parameters.instance().infection_radius = 0.5
         pop = TestSimFunctional.file_simulation("test_input.csv",
                                                 self.sim_params,
@@ -286,11 +292,12 @@ def test_small_cutoff(self, *mocks):
                                   + file_input["InfectMild"][i]))
 
     def test_waning_status_count(self, *mocks):
-        """Basic functional test to ensure more individuals enter InfectASympt
+        """Basic functional test to ensure more individuals enter Susceptible
         status given waning immunity.
         """
         # Record the number of individuals within the compartments when waning
         # immunity is active
+        pe.routine.Simulation.set_random_seed(seed=30)
         self.sim_params["initial_infected_number"] = 5
         self.sim_params["include_waning"] = True
         pe.Parameters.instance().use_waning_immunity = 1.0
@@ -314,12 +321,12 @@ def test_waning_status_count(self, *mocks):
         for cell in pop_2.cells:
             cell_data = cell.compartment_counter.retrieve()
 
-        # Compare the number of individuals in the InfectASympt compartment to
+        # Compare the number of individuals in the Susceptible compartment to
         # ensure that when waning immunity is active, the number of individuals
         # increases
         self.assertGreater(np.sum(cell_data_waning[
-            InfectionStatus.InfectASympt]), np.sum(cell_data[
-                InfectionStatus.InfectASympt]))
+            InfectionStatus.Susceptible]), np.sum(cell_data[
+                InfectionStatus.Susceptible]))
 
 
 if __name__ == '__main__':