Simplified births and deaths equality (#1628)

* simplified births elaborated to have a a mode whereby births equal deaths

* refactoring and editing

* typo fix!

* isort

* linting

* update argument in _choose_women_to_make_pregnant

* move call to sim object into apply() method

* add simple test

* move record of data back to __init__

* strengthen logic for chekcing death dates

---------

Co-authored-by: tdm32 <[email protected]>
Co-authored-by: Tara <[email protected]>

Author: tbhallett

src/tlo/methods/simplified_births.py

@@ -5,7 +5,7 @@
 
 import json
 from pathlib import Path
-from typing import Optional
+from typing import List, Optional
 
 import pandas as pd
 
@@ -85,8 +85,10 @@ class SimplifiedBirths(Module):
                                             categories=['none', 'non_exclusive', 'exclusive']),
     }
 
-    def __init__(self, name=None):
+    def __init__(self, name=None, force_one_birth_for_one_death: bool = False):
         super().__init__(name)
+        # Whether to use the _really_ simplified mode, whereby number of births is just equal to number of deaths
+        self.force_one_birth_for_one_death = force_one_birth_for_one_death
         self.asfr = dict()
 
         # Define defaults for properties:
@@ -162,8 +164,10 @@ def __init__(self, module):
         super().__init__(module, frequency=DateOffset(months=self.months_between_polls))
         self.asfr = get_medium_variant_asfr_from_wpp_resourcefile(
             dat=self.module.parameters['age_specific_fertility_rates'], months_exposure=self.months_between_polls)
+        self.date_of_last_poll = self.sim.date
 
     def apply(self, population):
+
         # Set new pregnancies:
         self.set_new_pregnancies()
 
@@ -174,26 +178,54 @@ def apply(self, population):
         self.update_breastfed_status()
 
     def set_new_pregnancies(self):
-        """Making women pregnant. Rate of doing so is based on age-specific fertility rates under assumption that every
-        pregnancy results in a birth."""
+        """Making women pregnant."""
 
-        df = self.sim.population.props  # get the population dataframe
+        def _make_pregnant(ids: List[int], df: pd.DataFrame, months_between_pregnancy_and_delivery: int) -> None:
+            """Enact the change to make the women pregnant"""
+            # updating properties for women who will get pregnant
+            df.loc[ids, 'is_pregnant'] = True
+            df.loc[ids, 'date_of_last_pregnancy'] = self.sim.date
+            df.loc[ids, 'si_date_of_last_delivery'] = \
+                self.sim.date + pd.DateOffset(months=months_between_pregnancy_and_delivery)
+
+        def _choose_women_to_make_pregnant(one_birth_for_one_death: bool, df: pd.DataFrame) -> List[int]:
+            """Choose women to make pregnant, depending on mode."""
+
+            if one_birth_for_one_death:
+                # Rate of pregnancy matches number of deaths: simple assumption and induces stable population size
+                # Non-pregnant women aged [15, 35) are selected randomly for pregnancy:
+
+                eligible_for_pregnancy = df.loc[
+                    (df.sex == 'F') & df.is_alive & ~df.is_pregnant & df.age_years.between(15, 35, inclusive='left')
+                    ]
+
+                num_of_deaths_since_last_poll = len(
+                    df.loc[df.date_of_death.between(self.date_of_last_poll, self.sim.date, inclusive='left')])
 
-        # find probability of becoming pregnant (using asfr for the year, limiting to alive, non-pregnant females)
-        prob_preg = df.loc[
-            (df.sex == 'F') & df.is_alive & ~df.is_pregnant
-            ]['age_range'].map(self.asfr[self.sim.date.year]).fillna(0)
-
-        # determine which woman will get pregnant
-        pregnant_women_ids = prob_preg.index[
-            (self.module.rng.random_sample(size=len(prob_preg)) < prob_preg)
-        ]
-
-        # updating properties for women who will get pregnant
-        df.loc[pregnant_women_ids, 'is_pregnant'] = True
-        df.loc[pregnant_women_ids, 'date_of_last_pregnancy'] = self.sim.date
-        df.loc[pregnant_women_ids, 'si_date_of_last_delivery'] = \
-            self.sim.date + pd.DateOffset(months=self.module.parameters['months_between_pregnancy_and_delivery'])
+                return self.module.rng.choice(
+                    eligible_for_pregnancy.index, size=num_of_deaths_since_last_poll, replace=False, p=None)
+
+            else:
+                # Rate of pregnancy is based on age-specific fertility rates under
+                # assumption that every pregnancy results in a birth.
+
+                # find probability of becoming pregnant
+                # (using asfr for the year, limiting to alive, non-pregnant females)
+                prob_preg = df.loc[
+                    (df.sex == 'F') & df.is_alive & ~df.is_pregnant
+                    ]['age_range'].map(self.asfr[self.sim.date.year]).fillna(0)
+
+                # determine which woman will get pregnant:
+                return prob_preg.index[
+                    (self.module.rng.random_sample(size=len(prob_preg)) < prob_preg)
+                ]
+
+        df = self.sim.population.props  # get the population dataframe
+        women_to_make_pregnant = _choose_women_to_make_pregnant(self.module.force_one_birth_for_one_death, df)
+        _make_pregnant(ids=women_to_make_pregnant,
+                       df=df,
+                       months_between_pregnancy_and_delivery=self.module.parameters['months_between_pregnancy_and_delivery'])
+        self.date_of_last_poll = self.sim.date
 
     def do_deliveries(self):
         """Checks to see if the date-of-delivery for pregnant women has been reached and implement births where

tests/test_simplified_births.py

@@ -20,13 +20,14 @@ def check_dtypes(simulation):
     assert (df.dtypes == orig.dtypes).all()
 
 
-def get_sim(seed, popsize=1000):
+def get_sim(seed, popsize=1000, **kwargs):
     sim = Simulation(start_date=Date(2010, 1, 1), seed=seed, resourcefilepath=resourcefilepath)
 
     # Register the appropriate modules
-    sim.register(demography.Demography(),
-                 simplified_births.SimplifiedBirths()
-                 )
+    sim.register(
+        demography.Demography(),
+        simplified_births.SimplifiedBirths(**kwargs)
+    )
 
     # Make the population
     sim.make_initial_population(n=popsize)
@@ -235,10 +236,11 @@ def apply(self, population):
 
 
 @pytest.mark.slow
-def test_other_modules_running_with_simplified_births_module():
+def test_other_modules_running_with_simplified_births_module(seed):
     """Run a "full simulation" using the simplified_births module and other disease modules"""
     sim = Simulation(
         start_date=Date(2010, 1, 1),
+        seed=seed,
         log_config={
             'custom_levels': {
                 '*': logging.WARNING,
@@ -254,3 +256,19 @@ def test_other_modules_running_with_simplified_births_module():
     sim.simulate(end_date=Date(2011, 12, 31))
     check_property_integrity(sim)
     check_dtypes(sim)
+
+def test_using_option_force_one_birth_for_one_death(seed):
+    """Test that when we have the option `force_one_birth_for_one_death` we see the behaviour expected."""
+
+    sim = get_sim(seed=seed, popsize=1000, force_one_birth_for_one_death=True)
+
+    # Make births happen same month as pregnancy, so that at the end of each month, pregnancies == births
+    sim.modules['SimplifiedBirths'].parameters['months_between_pregnancy_and_delivery'] = 0
+
+    # Check that in a simulation, the number of births that occurs does match the number of deaths
+    sim.simulate(end_date=Date(2015, 1, 1))
+    df = sim.population.props
+    num_births = len(df.loc[df.date_of_birth.notnull() & (df.mother_id >= 0)])
+    num_deaths = len(df.loc[df.date_of_death.notnull()])
+
+    assert num_births == num_deaths