Merge pull request #999 from amath-idm/add-docs

Fixes to init_immunity

CHANGELOG.rst

@@ -24,6 +24,18 @@ These are the major improvements we are currently working on. If there is a spec
 Latest versions (3.0.x)
 ~~~~~~~~~~~~~~~~~~~~~~~
 
+
+Version 3.0.1 (2021-04-16)
+--------------------------
+- Immunity and vaccine parameters have been updated.
+- The ``People`` class has been updated to remove parameters that were copied into attributes; thus there is no longer both ``people.pars['pop_size']`` and ``people.pop_size``; only the former. Recommended practice is to use ``len(people)`` to get the number of people.
+- Loaded population files can now be used with more than one strain; arrays will be resized automatically. If there is a mismatch in the number of people, this will *not* be automatically resized.
+- A bug was fixed with the ``rescale`` argument to ``cv.strain()`` not having any effect.
+- Dead people are no longer eligible to be vaccinated.
+- *Regression information*: Any user scripts that call ``sim.people.pop_size`` should be updated to call ``len(sim.people)`` (preferred), or ``sim.n``, ``sim['pop_size']``, or ``sim.people.pars['pop_size']``.
+- *GitHub info*: PR `999 <https://github.com/amath-idm/covasim/pull/999>`__
+
+
 Version 3.0.0 (2021-04-13)
 --------------------------
 This version introduces fully featured vaccines, variants, and immunity. **Note:** These new features are still under development; please use with caution and email us at [email protected] if you have any questions or issues. We expect there to be several more releases over the next few weeks as we refine these new features.
@@ -34,6 +46,22 @@ Highlights
 - **Multi-strain modeling**: Model functionality has been extended to allow for modeling of multiple different co-circulating strains with different properties. This means you can now do e.g. ``b117 = cv.strain('b117', days=1, n_imports=20)`` followed by ``sim = cv.Sim(strains=b117)`` to import strain B117. Further examples are contained in ``tests/test_immunity.py`` and in Tutorial 8.
 - **New methods for vaccine modeling**: A new ``cv.vaccinate()`` intervention has been added, which allows more flexible modeling of vaccinations. Vaccines, like natural infections, are assumed to boost agents' immunity.
 - **Consistency**: By default, results from Covasim 3.0.0 should exactly match Covasim 2.1.2. To use the new features, you will need to manually specify ``cv.Sim(use_waning=True)``.
+- **Still TLDR?** Here's a quick showcase of the new features:
+
+.. code-block:: python
+
+    import covasim as cv
+
+    pars = dict(
+        use_waning    = True,  # Use the new immunity features
+        n_days        = 180,   # Set the days, as before
+        n_agents      = 50e3,  # New alias for pop_size
+        scaled_pop    = 200e3, # New alternative to specifying pop_scale
+        strains       = cv.strain('b117', days=20, n_imports=20), # Introduce B117
+        interventions = cv.vaccinate('astrazeneca', days=80), # Create a vaccine
+    )
+
+    cv.Sim(pars).run().plot('strain') # Create, run, and plot strain results
 
 Immunity-related parameter changes
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

README.rst

@@ -36,13 +36,13 @@ Covasim has been used for analyses in over a dozen countries, both to inform pol
 
 2. **Determining the optimal strategy for reopening schools, the impact of test and trace interventions, and the risk of occurrence of a second COVID-19 epidemic wave in the UK: a modelling study**. Panovska-Griffiths J, Kerr CC, Stuart RM, Mistry D, Klein DJ, Viner R, Bonnell C (2020-08-03). *Lancet Child and Adolescent Health* S2352-4642(20) 30250-9. doi: https://doi.org/10.1016/S2352-4642(20)30250-9.
 
-3. **Modelling the impact of reducing control measures on the COVID-19 pandemic in a low transmission setting**. Scott N, Palmer A, Delport D, Abeysuriya RG, Stuart RM, Kerr CC, Mistry D, Klein DJ, Sacks-Davis R, Heath K, Hainsworth S, Pedrana A, Stoove M, Wilson DP, Hellard M (in press; accepted 2020-09-02). *Medical Journal of Australia* [`Preprint <https://www.mja.com.au/journal/2020/modelling-impact-reducing-control-measures-covid-19-pandemic-low-transmission-setting>`__]; doi: https://doi.org/10.1101/2020.06.11.20127027.
+3. **Estimating and mitigating the risk of COVID-19 epidemic rebound associated with reopening of international borders in Vietnam: a modelling study**. Pham QD, Stuart RM, Nguyen TV, Luong QC, Tran DQ, Phan LT, Dang TQ, Tran DN, Mistry D, Klein DJ, Abeysuriya RG, Oron AP, Kerr CC (2021-04-12). *Lancet Global Health* S2214-109X(21) 00103-0; doi: https://doi.org/10.1016/S2214-109X(21)00103-0.
 
-4. **Lessons learned from Vietnam's COVID-19 response: the role of adaptive behaviour change and testing in epidemic control**. Pham QD, Stuart RM, Nguyen TV, Luong QC, Tran DQ, Phan LT, Dang TQ, Tran DN, Mistry D, Klein DJ, Abeysuriya RG, Oron AP, Kerr CC (in press; accepted 2021-02-25). *Lancet Global Health*; doi: https://doi.org/10.1101/2020.12.18.20248454.
+4. **Modelling the impact of reducing control measures on the COVID-19 pandemic in a low transmission setting**. Scott N, Palmer A, Delport D, Abeysuriya RG, Stuart RM, Kerr CC, Mistry D, Klein DJ, Sacks-Davis R, Heath K, Hainsworth S, Pedrana A, Stoove M, Wilson DP, Hellard M (in press; accepted 2020-09-02). *Medical Journal of Australia* [`Preprint <https://www.mja.com.au/journal/2020/modelling-impact-reducing-control-measures-covid-19-pandemic-low-transmission-setting>`__]; doi: https://doi.org/10.1101/2020.06.11.20127027.
 
-5. **The role of masks, testing and contact tracing in preventing COVID-19 resurgences: a case study from New South Wales, Australia**. Stuart RM, Abeysuriya RG, Kerr CC, Mistry D, Klein DJ, Gray R, Hellard M, Scott N (in press; accepted 2021-03-19). *BMJ Open* doi: https://doi.org/10.1101/2020.09.02.20186742.
+5. **The role of masks, testing and contact tracing in preventing COVID-19 resurgences: a case study from New South Wales, Australia**. Stuart RM, Abeysuriya RG, Kerr CC, Mistry D, Klein DJ, Gray R, Hellard M, Scott N (in press; accepted 2021-03-19). *BMJ Open*; doi: https://doi.org/10.1101/2020.09.02.20186742.
 
-6. **The potential contribution of face coverings to the control of SARS-CoV-2 transmission in schools and broader society in the UK: a modelling study**. Panovska-Griffiths J, Kerr CC, Waites W, Stuart RM, Mistry D, Foster D, Klein DJ, Viner R, Bonnell C (in press; accepted 2021-04-08). *Nature Scientific Reports* doi: https://doi.org/10.1101/2020.09.28.20202937.
+6. **The potential contribution of face coverings to the control of SARS-CoV-2 transmission in schools and broader society in the UK: a modelling study**. Panovska-Griffiths J, Kerr CC, Waites W, Stuart RM, Mistry D, Foster D, Klein DJ, Viner R, Bonnell C (in press; accepted 2021-04-08). *Nature Scientific Reports*; doi: https://doi.org/10.1101/2020.09.28.20202937.
 
 7. **Schools are not islands: Balancing COVID-19 risk and educational benefits using structural and temporal countermeasures**. Cohen JA, Mistry D, Kerr CC, Klein DJ (under review; posted 2020-09-10). *medRxiv* 2020.09.08.20190942; doi: https://doi.org/10.1101/2020.09.08.20190942.
 

covasim/base.py

@@ -858,7 +858,7 @@ def __setitem__(self, key, value):
 
     def __len__(self):
         ''' This is just a scalar, but validate() and _resize_arrays() make sure it's right '''
-        return self.pop_size
+        return self.pars['pop_size']
 
 
     def __iter__(self):
@@ -883,7 +883,7 @@ def __add__(self, people2):
                 raise NotImplementedError(errormsg)
 
         # Validate
-        newpeople.pop_size += people2.pop_size
+        newpeople.pars['pop_size'] += people2.pars['pop_size']
         newpeople.validate()
 
         # Reassign UIDs so they're unique
@@ -892,6 +892,12 @@ def __add__(self, people2):
         return newpeople
 
 
+    def __radd__(self, people2):
+        ''' Allows sum() to work correctly '''
+        if not people2: return self
+        else:           return self.__add__(people2)
+
+
     def _brief(self):
         '''
         Return a one-line description of the people -- used internally and by repr();
@@ -1029,11 +1035,16 @@ def validate(self, die=True, verbose=False):
 
         # Check that the length of each array is consistent
         expected_len = len(self)
+        expected_strains = self.pars['n_strains']
         for key in self.keys():
-            actual_len = len(self[key])
-            # check if it's 2d
-            if self[key].ndim > 1:
-                actual_len = len(self[key][0])
+            if self[key].ndim == 1:
+                actual_len = len(self[key])
+            else: # If it's 2D, strains need to be checked separately
+                actual_strains, actual_len = self[key].shape
+                if actual_strains != expected_strains:
+                    if verbose:
+                        print(f'Resizing "{key}" from {actual_strains} to {expected_strains}')
+                    self._resize_arrays(keys=key, new_size=(expected_strains, expected_len))
             if actual_len != expected_len: # pragma: no cover
                 if die:
                     errormsg = f'Length of key "{key}" did not match population size ({actual_len} vs. {expected_len})'
@@ -1050,16 +1061,22 @@ def validate(self, die=True, verbose=False):
         return
 
 
-    def _resize_arrays(self, pop_size=None, keys=None):
+    def _resize_arrays(self, new_size=None, keys=None):
         ''' Resize arrays if any mismatches are found '''
-        if pop_size is None:
-            pop_size = len(self)
-        self.pop_size = pop_size
+
+        # Handle None or tuple input (representing strains and pop_size)
+        if new_size is None:
+            new_size = len(self)
+        pop_size = new_size if not isinstance(new_size, tuple) else new_size[1]
+        self.pars['pop_size'] = pop_size
+
+        # Reset sizes
         if keys is None:
             keys = self.keys()
         keys = sc.promotetolist(keys)
         for key in keys:
-            self[key].resize(pop_size, refcheck=False)
+            self[key].resize(new_size, refcheck=False) # Don't worry about cross-references to the arrays
+
         return
 
 
@@ -1113,7 +1130,7 @@ def from_people(self, people, resize=True):
         # Handle population size
         pop_size = len(people)
         if resize:
-            self._resize_arrays(pop_size=pop_size)
+            self._resize_arrays(new_size=pop_size)
 
         # Iterate over people -- slow!
         for p,person in enumerate(people):

covasim/defaults.py

@@ -190,7 +190,7 @@ def __init__(self):
 
 # Parameters that can vary by strain
 strain_pars = [
-    'rel_imm',
+    'rel_imm_strain',
     'rel_beta',
     'rel_symp_prob',
     'rel_severe_prob',

covasim/immunity.py

@@ -38,6 +38,7 @@ class strain(sc.prettyobj):
     def __init__(self, strain, days, label=None, n_imports=1, rescale=True):
         self.days = days # Handle inputs
         self.n_imports = int(n_imports)
+        self.rescale   = rescale
         self.index     = None # Index of the strain in the sim; set later
         self.label     = None # Strain label (used as a dict key)
         self.p         = None # This is where the parameters will be stored
@@ -117,7 +118,8 @@ def apply(self, sim):
         ''' Introduce new infections with this strain '''
         for ind in cvi.find_day(self.days, sim.t, interv=self, sim=sim): # Time to introduce strain
             susceptible_inds = cvu.true(sim.people.susceptible)
-            n_imports = sc.randround(self.n_imports/sim.rescale_vec[sim.t]) # Round stochastically to the nearest number of imports
+            rescale_factor = sim.rescale_vec[sim.t] if self.rescale else 1.0
+            n_imports = sc.randround(self.n_imports/rescale_factor) # Round stochastically to the nearest number of imports
             importation_inds = np.random.choice(susceptible_inds, n_imports)
             sim.people.infect(inds=importation_inds, layer='importation', strain=self.index)
         return
@@ -251,29 +253,29 @@ def init_immunity(sim, create=False):
     # Pull out all of the circulating strains for cross-immunity
     ns       = sim['n_strains']
     immunity = {}
-    rel_imms = {}
-    strain_labels = sim['strain_map'].values()
-    for label in strain_labels:
-        rel_imms[label] = sim['strain_pars'][label]['rel_imm']
 
     # If immunity values have been provided, process them
     if sim['immunity'] is None or create:
-        # Initialize immunity
+
+        # Firstly, initialize immunity matrix with defaults. These are then overwitten with strain-specific values below
         for ax in cvd.immunity_axes:
             if ax == 'sus':  # Susceptibility matrix is of size sim['n_strains']*sim['n_strains']
-                immunity[ax] = np.full((ns, ns), sim['cross_immunity'], dtype=cvd.default_float)  # Default for off-diagnonals
-                np.fill_diagonal(immunity[ax], 1.0)  # Default for own-immunity
+                immunity[ax] = np.ones((ns, ns), dtype=cvd.default_float)  # Fill with defaults
             else:  # Progression and transmission are matrices of scalars of size sim['n_strains']
-                immunity[ax] = np.ones(ns, dtype=cvd.default_float)
+                immunity[ax] = np.ones(ns, dtype=cvd.default_float)  # Fill with defaults
 
+        # Next, overwrite these defaults with any known immunity values about specific strains
         default_cross_immunity = cvpar.get_cross_immunity()
         for i in range(ns):
+            label_i = sim['strain_map'][i]
             for j in range(ns):
-                if i != j:
-                    label_i = sim['strain_map'][i]
+                if i != j: # Populate cross-immunity
                     label_j = sim['strain_map'][j]
                     if label_i in default_cross_immunity and label_j in default_cross_immunity:
                         immunity['sus'][j][i] = default_cross_immunity[label_j][label_i]
+                else: # Populate own-immunity
+                    immunity['sus'][i, i] = sim['strain_pars'][label_i]['rel_imm_strain']
+
         sim['immunity'] = immunity
 
     # Next, precompute the NAb kinetics and store these for access during the sim
@@ -464,4 +466,4 @@ def linear_decay(length, init_val, slope):
 def linear_growth(length, slope):
     ''' Calculate linear growth '''
     t = np.arange(length, dtype=cvd.default_int)
-    return (slope * t)
+    return (slope * t)