Merge branch 'main' into waning-immunity

README.md

@@ -97,6 +97,7 @@ Configure a simulation with a number of parameters. These are split into three c
 * `output_dir`: String for the location for the output files, as a relative path
 * `status_output`: Boolean to determine whether we need a csv file containing infection status values _(Default false)_
 * `infectiousness_output`: Boolean to determine whether we need a csv file containing infectiousness (viral load) values _(Default false)_
+* `compress`: Boolean to determine whether we compress a csv file containing infection status values and/or a csv file containing infectiousness (viral load) values if they are written _(Default false)_
 
 Two lists of sweeps must also be passed to this function - the first will be executed once at the start of the simulation (i.e. to determine the initial infections in the population), while the second list will be ran at every timestep (i.e. to propagate the infection through the population).
 

pyEpiabm/pyEpiabm/output/_csv_dict_writer.py

@@ -5,6 +5,8 @@
 import csv
 import typing
 import os
+import pandas as pd
+import logging
 
 from pyEpiabm.output.abstract_reporter import AbstractReporter
 
@@ -28,6 +30,11 @@ def __init__(self, folder: str, filename: str, fieldnames: typing.List,
 
         """
         super().__init__(folder, clear_folder)
+        self.filename = filename
+        self.filepath = os.path.join(folder, filename)
+        self.filepath_without_extension = os.path.join(
+            folder, os.path.splitext(filename)[0])
+        self.fieldnames = fieldnames
 
         self.f = open(os.path.join(folder, filename), 'w')
         self.writer = csv.DictWriter(
@@ -52,3 +59,13 @@ def write(self, row: typing.Dict):
 
         """
         self.writer.writerow(row)
+        self.f.flush()
+
+    def compress(self):
+        """Compresses the csv file and deletes the unzipped csv.
+        """
+        output_filepath = f"{self.filepath_without_extension}.zip"
+        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'})
+        os.remove(self.filepath)

pyEpiabm/pyEpiabm/routine/simulation.py

@@ -62,6 +62,8 @@ def configure(self,
                a csv file containing infection status values
             * `infectiousness_output`: Boolean to determine whether we need \
                a csv file containing infectiousness (viral load) values
+            * `compress`: Boolean to determine whether we compress \
+               the infection history csv files
 
         Parameters
         ----------
@@ -87,7 +89,9 @@ def configure(self,
             the dictionary is None). These files contain the infection status
             or infectiousness of each person every time step. The EpiOS tool
             (https://github.com/SABS-R3-Epidemiology/EpiOS) samples data from
-            these files to mimic real life epidemic sampling techniques
+            these files to mimic real life epidemic sampling techniques. These
+            files can be compressed when 'compress' is True, reducing the size
+            of these files.
         """
         self.sim_params = sim_params
         self.population = population
@@ -146,6 +150,7 @@ def configure(self,
         self.infectiousness_output = False
         self.ih_status_writer = None
         self.ih_infectiousness_writer = None
+        self.compress = False
 
         if inf_history_params:
             # Setting up writer for infection history for each person. If the
@@ -155,6 +160,7 @@ def configure(self,
 
             self.infectiousness_output = inf_history_params\
                 .get("infectiousness_output")
+            self.compress = inf_history_params.get("compress", False)
             person_ids = []
             person_ids += [person.id for cell in population.cells for person
                            in cell.persons]
@@ -312,7 +318,7 @@ def write_to_ih_file(self, time, output_option: str):
         """
         if self.status_output and output_option == "status":
             ih_data = {column: 0 for column in
-                       self.ih_status_writer.writer.fieldnames}
+                       self.ih_status_writer.fieldnames}
             for cell in self.population.cells:
                 for person in cell.persons:
                     ih_data[person.id] = person.infection_status.value
@@ -322,7 +328,7 @@ def write_to_ih_file(self, time, output_option: str):
 
         if self.infectiousness_output and output_option == "infectiousness":
             infect_data = {column: 0 for column in
-                           self.ih_infectiousness_writer.writer.fieldnames}
+                           self.ih_infectiousness_writer.fieldnames}
             for cell in self.population.cells:
                 for person in cell.persons:
                     infect_data[person.id] = person.infectiousness
@@ -333,6 +339,16 @@ def write_to_ih_file(self, time, output_option: str):
     def add_writer(self, writer: AbstractReporter):
         self.writers.append(writer)
 
+    def compress_csv(self):
+        """Compresses the infection history csvs when they are written.
+
+        """
+        if self.compress and self.ih_status_writer:
+            self.ih_status_writer.compress()
+
+        if self.compress and self.ih_infectiousness_writer:
+            self.ih_infectiousness_writer.compress()
+
     @staticmethod
     def set_random_seed(seed):
         """ Set random seed for all subsequent operations. Should be used

pyEpiabm/pyEpiabm/tests/test_unit/test_output/test_csv_dict_writer.py

@@ -50,6 +50,27 @@ def test_del(self, mock_mkdir):
             m.__del__()
             fake_file.close.assert_called_once()
 
+    @patch('os.makedirs')
+    @patch('os.remove')
+    @patch('pandas.read_csv')
+    @patch('pandas.DataFrame.to_csv')
+    def test_compress(self, mock_mkdir, mock_remove, mock_read, mock_to_csv):
+        """Test the compress method of the _CsvDictWriter class.
+        """
+        mo = mock_open()
+        with patch('pyEpiabm.output._csv_dict_writer.open', mo):
+            mock_content = ['1', '2', '3']
+            m = pe.output._CsvDictWriter('mock_folder', 'mock_filename.csv',
+                                         mock_content)
+            m.compress()
+            expected_output_filepath = f"{m.filepath_without_extension}.zip"
+            self.assertEqual(expected_output_filepath,
+                             os.path.join("mock_folder", "mock_filename.zip"))
+
+            mock_read.assert_called_once_with(m.filepath)
+            mock_to_csv.assert_called_once()
+            mock_remove.assert_called_once_with(m.filepath)
+
 
 if __name__ == '__main__':
     unittest.main()

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

@@ -2,7 +2,7 @@
 import random
 import numpy as np
 import unittest
-from unittest.mock import patch, mock_open
+from unittest.mock import patch, mock_open, MagicMock
 
 import pyEpiabm as pe
 
@@ -62,13 +62,14 @@ def test_configure(self, mock_mkdir):
             self.assertEqual(len(test_sim.sweeps), 1)
             self.assertIsInstance(test_sim.population, pe.Population)
 
-            # Test that the ih writers are None, and that both status_output
-            # and infectiousness_output are False
+            # Test that the ih writers are None, that both status_output
+            # and infectiousness_output are False, and that compress is false
             self.assertEqual(test_sim.status_output, False)
             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
             del test_sim.ih_status_writer
@@ -463,13 +464,13 @@ def test_write_to_ih_file_status_infectiousness(self, mock_mkdir, time=1):
                                self.sweeps, self.sim_params, self.file_params,
                                self.inf_history_params)
             ih_data = {column: 0 for column in
-                       test_sim.ih_status_writer.writer.fieldnames}
+                       test_sim.ih_status_writer.fieldnames}
             for cell in self.test_population.cells:
                 for person in cell.persons:
                     ih_data[person.id] = person.infection_status.value
             ih_data["time"] = time
             infect_data = {column: 0 for column in
-                           test_sim.ih_infectiousness_writer.writer.fieldnames}
+                           test_sim.ih_infectiousness_writer.fieldnames}
             for cell in self.test_population.cells:
                 for person in cell.persons:
                     infect_data[person.id] = person.infectiousness
@@ -485,6 +486,71 @@ def test_write_to_ih_file_status_infectiousness(self, mock_mkdir, time=1):
         mock_mkdir.assert_called_with(
             os.path.join(os.getcwd(), self.inf_history_params["output_dir"]))
 
+    @patch('os.makedirs')
+    def test_compress_no_compression(self, mock_mkdir):
+        mo = mock_open()
+        self.inf_history_params['compress'] = False
+        self.inf_history_params['status_output'] = True
+        self.inf_history_params['infectiousness_output'] = True
+
+        with patch('pyEpiabm.output._csv_dict_writer.open', mo):
+            test_sim = pe.routine.Simulation()
+            mock_status_writer = MagicMock()
+            mock_infectiousness_writer = MagicMock()
+            test_sim.configure(self.test_population, self.initial_sweeps,
+                               self.sweeps, self.sim_params, self.file_params,
+                               self.inf_history_params)
+            test_sim.run_sweeps()
+
+            with patch.object(test_sim, 'ih_status_writer',
+                              mock_status_writer):
+                with patch.object(test_sim, 'ih_infectiousness_writer',
+                                  mock_infectiousness_writer):
+                    test_sim.compress_csv()
+                    mock_status_writer.compress.assert_not_called()
+                    mock_infectiousness_writer.compress.assert_not_called()
+            del test_sim
+
+    @patch('os.makedirs')
+    def test_compress_csv_status(self, mock_mkdir):
+        mo = mock_open()
+        self.inf_history_params['compress'] = True
+        self.inf_history_params['status_output'] = True
+        self.inf_history_params['infectiousness_output'] = False
+
+        with patch('pyEpiabm.output._csv_dict_writer.open', mo):
+            test_sim = pe.routine.Simulation()
+            mock_status_writer = MagicMock()
+            test_sim.configure(self.test_population, self.initial_sweeps,
+                               self.sweeps, self.sim_params, self.file_params,
+                               self.inf_history_params)
+            test_sim.run_sweeps()
+            with patch.object(test_sim, 'ih_status_writer',
+                              mock_status_writer):
+                test_sim.compress_csv()
+                mock_status_writer.compress.assert_called_once_with()
+            del test_sim
+
+    @patch('os.makedirs')
+    def test_compress_csv_infect(self, mock_mkdir):
+        mo = mock_open()
+        self.inf_history_params['compress'] = True
+        self.inf_history_params['status_output'] = False
+        self.inf_history_params['infectiousness_output'] = True
+
+        with patch('pyEpiabm.output._csv_dict_writer.open', mo):
+            test_sim = pe.routine.Simulation()
+            mock_infect_writer = MagicMock()
+            test_sim.configure(self.test_population, self.initial_sweeps,
+                               self.sweeps, self.sim_params, self.file_params,
+                               self.inf_history_params)
+            test_sim.run_sweeps()
+            with patch.object(test_sim, 'ih_infectiousness_writer',
+                              mock_infect_writer):
+                test_sim.compress_csv()
+                mock_infect_writer.compress.assert_called_once_with()
+            del test_sim
+
     def test_set_random_seed(self):
         pe.routine.Simulation.set_random_seed(seed=0)
         value = random.random()

python_examples/basic_infection_history_simulation/simulation_flow.py

@@ -50,7 +50,8 @@
 inf_history_params = {"output_dir": os.path.join(os.path.dirname(__file__),
                                                  "simulation_outputs"),
                       "status_output": True,
-                      "infectiousness_output": True}
+                      "infectiousness_output": True,
+                      "compress": False}  # Set to True if compression desired
 
 # Create a simulation object, configure it with the parameters given, then
 # run the simulation.
@@ -65,6 +66,7 @@
     file_params,
     inf_history_params)
 sim.run_sweeps()
+sim.compress_csv()
 
 # Need to close the writer object at the end of each simulation.
 del sim.writer