co: more-readable update_params_for_interventions() fnc

diff --git src/tlo/methods/contraception.py src/tlo/methods/contraception.py
index 94f2c7108..bc118d4e5 100644
--- src/tlo/methods/contraception.py
+++ src/tlo/methods/contraception.py
@@ -611,121 +611,37 @@ class Contraception(Module):

         processed_params = self.processed_params

-        def expand_to_age_years(values_by_age_groups, ages_by_year):
-            _d = dict(zip(['15-19', '20-24', '25-29', '30-34', '35-39', '40-44', '45-49'], values_by_age_groups))
-            return np.array(
-                [_d[self.sim.modules['Demography'].AGE_RANGE_LOOKUP[_age_year]] for _age_year in ages_by_year]
-            )
+        def contraception_initiation_with_interv(p_start_per_month_without_interv):
+            """Increase the probabilities of a woman starting modern contraceptives due to Pop intervention being
+            applied."""
+            # TODO: remove the keys before intervention year
+            p_start_per_month_with_interv = {}
+            for year, age_method_df in p_start_per_month_without_interv.items():
+                p_start_per_month_with_interv[year] = age_method_df * self.parameters['Interventions_Pop'].loc[0]
+            return p_start_per_month_with_interv

-        def time_age_trend_in_initiation():
-            """The age-specific effect of calendar year on the probability of starting use of contraceptive
-            (multiplicative effect). Values are chosen to induce a trend in age-specific fertility consistent with
-             the WPP estimates."""
+        def contraception_initiation_after_birth_with_interv(p_start_after_birth_without_interv):
+            """Increase the probabilities of a woman starting modern contraceptives following giving birth due to PPFP
+            intervention being applied."""
+            # Exclude prob of 'not_using'
+            p_start_after_birth_with_interv = p_start_after_birth_without_interv.copy().drop('not_using')

-            _years = np.arange(2010, 2101)
-            _ages = np.arange(15, 50)
+            # Apply PPFP intervention multipliers (ie increase probs of modern methods)
+            p_start_after_birth_with_interv = \
+                p_start_after_birth_with_interv.mul(self.parameters['Interventions_PPFP'].loc[0])

-            _init_over_time = np.exp(+0.05 * np.minimum(2020 - 2010, (_years - 2010))) * np.maximum(1.0, np.exp(
-                +0.01 * (_years - 2020)))
-            _init_over_time_modification_by_age = 1.0 / expand_to_age_years([1.0, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5], _ages)
-            _init = np.outer(_init_over_time, _init_over_time_modification_by_age)
+            # Return reduced prob of 'not_using'
+            p_start_after_birth_with_interv = pd.Series((1.0 - p_start_after_birth_with_interv.sum()),
+                                                        index=['not_using']).append(p_start_after_birth_with_interv)

-            return pd.DataFrame(index=_years, columns=_ages, data=_init)
+            return p_start_after_birth_with_interv

-        def avoid_sterilization_below30(probs):
-            """Prevent women below 30 years having female sterilization and adjust the probability for women 30 and over
-            to preserve the overall probability of initiating sterilization."""
-            # Input 'probs' must include probs for all methods including 'not_using'
-            assert set(probs.index) == set(self.all_contraception_states)
-
-            # Prevent women below 30 years having 'female_sterilization'
-            probs_below30 = probs.copy()
-            probs_below30['female_sterilization'] = 0.0
-            # Scale so that the probability of all outcomes sum to 1.0
-            probs_below30 = probs_below30 / probs_below30.sum()
-            assert np.isclose(1.0, probs_below30.sum())
-
-            # Increase prob of 'female_sterilization' in older women accordingly
-            probs_30plus = probs.copy()
-            probs_30plus['female_sterilization'] = \
-                probs.loc['female_sterilization'] / \
-                (self.n_female3049_in_2010 / self.n_female1549_in_2010)
-            # Scale so that the probability of all outcomes sum to 1.0
-            probs_30plus = probs_30plus / probs_30plus.sum()
-            assert np.isclose(1.0, probs_30plus.sum())
-
-            return probs_below30, probs_30plus
-
-        def contraception_initiation_with_interv():
-            """Generate the probability per month of a woman initiating onto each contraceptive, by the age (in whole
-             years) if FP interventions are applied."""
-
-            # Probability of initiation by method per month (average over all ages)
-            p_init_by_method = self.parameters['Initiation_ByMethod'].loc[0]
-
-            # Prevent women below 30 years having 'female_sterilization' while preserving the overall probability of
-            # 'female_sterilization' initiation
-            p_init_by_method_below30, p_init_by_method_30plus = avoid_sterilization_below30(p_init_by_method)
-
-            # Effect of age
-            age_effect = 1.0 + self.parameters['Initiation_ByAge'].set_index('age')['r_init1_age'].rename_axis(
-                "age_years")
-
-            # Year effect
-            year_effect = time_age_trend_in_initiation()
-
-            def apply_intervention_age_year_effects(probs_below30, probs_30plus):
-                # Apply Pop intervention
-                probs_by_method_below30 = \
-                    probs_below30.copy().drop('not_using').mul(self.parameters['Interventions_Pop'].loc[0])
-                probs_by_method_30plus = \
-                    probs_30plus.copy().drop('not_using').mul(self.parameters['Interventions_Pop'].loc[0])
-                # Assemble into age-specific data-frame:
-                p_init = dict()
-                for year in year_effect.index:
-
-                    p_init_this_year = dict()
-                    for a in age_effect.index:
-                        if a < 30:
-                            p_init_this_year[a] = probs_by_method_below30 * age_effect.at[a] * year_effect.at[year, a]
-                        else:
-                            p_init_this_year[a] = probs_by_method_30plus * age_effect.at[a] * year_effect.at[year, a]
-                    p_init_this_year_df = pd.DataFrame.from_dict(p_init_this_year, orient='index')
-
-                    # Check correct format of age/method data-frame
-                    assert set(p_init_this_year_df.columns) == set(self.all_contraception_states - {'not_using'})
-                    assert (p_init_this_year_df.index == range(15, 50)).all()
-                    assert (p_init_this_year_df >= 0.0).all().all()
-
-                    p_init[year] = p_init_this_year_df
-
-                return p_init
-
-            return apply_intervention_age_year_effects(p_init_by_method_below30, p_init_by_method_30plus)
-
-        def contraception_initiation_after_birth_with_interv():
-            """Get the probability of a woman starting a contraceptive following giving birth if FP interventions are
-            applied. Avoid sterilization in women below 30 years old."""
-
-            # Get initiation probabilities of contraception methods after birth from read-in Excel sheet
-            p_start_after_birth = self.parameters['Initiation_AfterBirth'].loc[0].drop('not_using')
-
-            # Apply PPFP intervention multipliers
-            p_start_after_birth = p_start_after_birth.mul(self.parameters['Interventions_PPFP'].loc[0])
-
-            # Add 'not_using' to initiation probabilities of contraception methods after birth
-            p_start_after_birth = pd.concat(
-                (
-                    pd.Series((1.0 - p_start_after_birth.sum()), index=['not_using']),
-                    p_start_after_birth
-                )
-            )
-
-            return avoid_sterilization_below30(p_start_after_birth)
-
-        processed_params['p_start_per_month'] = contraception_initiation_with_interv()
-        processed_params['p_start_after_birth_below30'], processed_params['p_start_after_birth_30plus'] =\
-            contraception_initiation_after_birth_with_interv()
+        processed_params['p_start_per_month'] = \
+            contraception_initiation_with_interv(processed_params['p_start_per_month'])
+        processed_params['p_start_after_birth_below30'] = \
+            contraception_initiation_after_birth_with_interv(processed_params['p_start_after_birth_below30'])
+        processed_params['p_start_after_birth_30plus'] = \
+            contraception_initiation_after_birth_with_interv(processed_params['p_start_after_birth_30plus'])

         return processed_params

@@ -750,7 +666,6 @@ class Contraception(Module):
         get_items_from_pkg = self.sim.modules['HealthSystem'].get_item_codes_from_package_name

         _cons_codes = dict()
-        # items for each method that requires an HSI to switch to
         _cons_codes['pill'] = get_items_from_pkg('Pill')
         _cons_codes['male_condom'] = get_items_from_pkg('Male condom')
         _cons_codes['other_modern'] = get_items_from_pkg('Female Condom')
@@ -760,7 +675,6 @@ class Contraception(Module):
         _cons_codes['implant'] = get_items_from_pkg('Implant')
         _cons_codes['female_sterilization'] = get_items_from_pkg('Female sterilization')
         assert set(_cons_codes.keys()) == set(self.states_that_may_require_HSI_to_switch_to)
-        # items used when initiating a modern reliable method after not using or switching from non-reliable method
         _cons_codes['co_initiation'] = get_items_from_pkg('Contraception initiation')

         return _cons_codes

Author: EvaJanouskova

src/tlo/methods/contraception.py

@@ -611,121 +611,37 @@ def update_params_for_interventions(self):
 
         processed_params = self.processed_params
 
-        def expand_to_age_years(values_by_age_groups, ages_by_year):
-            _d = dict(zip(['15-19', '20-24', '25-29', '30-34', '35-39', '40-44', '45-49'], values_by_age_groups))
-            return np.array(
-                [_d[self.sim.modules['Demography'].AGE_RANGE_LOOKUP[_age_year]] for _age_year in ages_by_year]
-            )
-
-        def time_age_trend_in_initiation():
-            """The age-specific effect of calendar year on the probability of starting use of contraceptive
-            (multiplicative effect). Values are chosen to induce a trend in age-specific fertility consistent with
-             the WPP estimates."""
-
-            _years = np.arange(2010, 2101)
-            _ages = np.arange(15, 50)
-
-            _init_over_time = np.exp(+0.05 * np.minimum(2020 - 2010, (_years - 2010))) * np.maximum(1.0, np.exp(
-                +0.01 * (_years - 2020)))
-            _init_over_time_modification_by_age = 1.0 / expand_to_age_years([1.0, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5], _ages)
-            _init = np.outer(_init_over_time, _init_over_time_modification_by_age)
-
-            return pd.DataFrame(index=_years, columns=_ages, data=_init)
-
-        def avoid_sterilization_below30(probs):
-            """Prevent women below 30 years having female sterilization and adjust the probability for women 30 and over
-            to preserve the overall probability of initiating sterilization."""
-            # Input 'probs' must include probs for all methods including 'not_using'
-            assert set(probs.index) == set(self.all_contraception_states)
-
-            # Prevent women below 30 years having 'female_sterilization'
-            probs_below30 = probs.copy()
-            probs_below30['female_sterilization'] = 0.0
-            # Scale so that the probability of all outcomes sum to 1.0
-            probs_below30 = probs_below30 / probs_below30.sum()
-            assert np.isclose(1.0, probs_below30.sum())
-
-            # Increase prob of 'female_sterilization' in older women accordingly
-            probs_30plus = probs.copy()
-            probs_30plus['female_sterilization'] = \
-                probs.loc['female_sterilization'] / \
-                (self.n_female3049_in_2010 / self.n_female1549_in_2010)
-            # Scale so that the probability of all outcomes sum to 1.0
-            probs_30plus = probs_30plus / probs_30plus.sum()
-            assert np.isclose(1.0, probs_30plus.sum())
-
-            return probs_below30, probs_30plus
-
-        def contraception_initiation_with_interv():
-            """Generate the probability per month of a woman initiating onto each contraceptive, by the age (in whole
-             years) if FP interventions are applied."""
-
-            # Probability of initiation by method per month (average over all ages)
-            p_init_by_method = self.parameters['Initiation_ByMethod'].loc[0]
-
-            # Prevent women below 30 years having 'female_sterilization' while preserving the overall probability of
-            # 'female_sterilization' initiation
-            p_init_by_method_below30, p_init_by_method_30plus = avoid_sterilization_below30(p_init_by_method)
-
-            # Effect of age
-            age_effect = 1.0 + self.parameters['Initiation_ByAge'].set_index('age')['r_init1_age'].rename_axis(
-                "age_years")
-
-            # Year effect
-            year_effect = time_age_trend_in_initiation()
-
-            def apply_intervention_age_year_effects(probs_below30, probs_30plus):
-                # Apply Pop intervention
-                probs_by_method_below30 = \
-                    probs_below30.copy().drop('not_using').mul(self.parameters['Interventions_Pop'].loc[0])
-                probs_by_method_30plus = \
-                    probs_30plus.copy().drop('not_using').mul(self.parameters['Interventions_Pop'].loc[0])
-                # Assemble into age-specific data-frame:
-                p_init = dict()
-                for year in year_effect.index:
-
-                    p_init_this_year = dict()
-                    for a in age_effect.index:
-                        if a < 30:
-                            p_init_this_year[a] = probs_by_method_below30 * age_effect.at[a] * year_effect.at[year, a]
-                        else:
-                            p_init_this_year[a] = probs_by_method_30plus * age_effect.at[a] * year_effect.at[year, a]
-                    p_init_this_year_df = pd.DataFrame.from_dict(p_init_this_year, orient='index')
-
-                    # Check correct format of age/method data-frame
-                    assert set(p_init_this_year_df.columns) == set(self.all_contraception_states - {'not_using'})
-                    assert (p_init_this_year_df.index == range(15, 50)).all()
-                    assert (p_init_this_year_df >= 0.0).all().all()
-
-                    p_init[year] = p_init_this_year_df
-
-                return p_init
-
-            return apply_intervention_age_year_effects(p_init_by_method_below30, p_init_by_method_30plus)
-
-        def contraception_initiation_after_birth_with_interv():
-            """Get the probability of a woman starting a contraceptive following giving birth if FP interventions are
-            applied. Avoid sterilization in women below 30 years old."""
-
-            # Get initiation probabilities of contraception methods after birth from read-in Excel sheet
-            p_start_after_birth = self.parameters['Initiation_AfterBirth'].loc[0].drop('not_using')
-
-            # Apply PPFP intervention multipliers
-            p_start_after_birth = p_start_after_birth.mul(self.parameters['Interventions_PPFP'].loc[0])
-
-            # Add 'not_using' to initiation probabilities of contraception methods after birth
-            p_start_after_birth = pd.concat(
-                (
-                    pd.Series((1.0 - p_start_after_birth.sum()), index=['not_using']),
-                    p_start_after_birth
-                )
-            )
-
-            return avoid_sterilization_below30(p_start_after_birth)
-
-        processed_params['p_start_per_month'] = contraception_initiation_with_interv()
-        processed_params['p_start_after_birth_below30'], processed_params['p_start_after_birth_30plus'] =\
-            contraception_initiation_after_birth_with_interv()
+        def contraception_initiation_with_interv(p_start_per_month_without_interv):
+            """Increase the probabilities of a woman starting modern contraceptives due to Pop intervention being
+            applied."""
+            # TODO: remove the keys before intervention year
+            p_start_per_month_with_interv = {}
+            for year, age_method_df in p_start_per_month_without_interv.items():
+                p_start_per_month_with_interv[year] = age_method_df * self.parameters['Interventions_Pop'].loc[0]
+            return p_start_per_month_with_interv
+
+        def contraception_initiation_after_birth_with_interv(p_start_after_birth_without_interv):
+            """Increase the probabilities of a woman starting modern contraceptives following giving birth due to PPFP
+            intervention being applied."""
+            # Exclude prob of 'not_using'
+            p_start_after_birth_with_interv = p_start_after_birth_without_interv.copy().drop('not_using')
+
+            # Apply PPFP intervention multipliers (ie increase probs of modern methods)
+            p_start_after_birth_with_interv = \
+                p_start_after_birth_with_interv.mul(self.parameters['Interventions_PPFP'].loc[0])
+
+            # Return reduced prob of 'not_using'
+            p_start_after_birth_with_interv = pd.Series((1.0 - p_start_after_birth_with_interv.sum()),
+                                                        index=['not_using']).append(p_start_after_birth_with_interv)
+
+            return p_start_after_birth_with_interv
+
+        processed_params['p_start_per_month'] = \
+            contraception_initiation_with_interv(processed_params['p_start_per_month'])
+        processed_params['p_start_after_birth_below30'] = \
+            contraception_initiation_after_birth_with_interv(processed_params['p_start_after_birth_below30'])
+        processed_params['p_start_after_birth_30plus'] = \
+            contraception_initiation_after_birth_with_interv(processed_params['p_start_after_birth_30plus'])
 
         return processed_params
 
@@ -750,7 +666,6 @@ def get_item_code_for_each_contraceptive(self):
         get_items_from_pkg = self.sim.modules['HealthSystem'].get_item_codes_from_package_name
 
         _cons_codes = dict()
-        # items for each method that requires an HSI to switch to
         _cons_codes['pill'] = get_items_from_pkg('Pill')
         _cons_codes['male_condom'] = get_items_from_pkg('Male condom')
         _cons_codes['other_modern'] = get_items_from_pkg('Female Condom')
@@ -760,7 +675,6 @@ def get_item_code_for_each_contraceptive(self):
         _cons_codes['implant'] = get_items_from_pkg('Implant')
         _cons_codes['female_sterilization'] = get_items_from_pkg('Female sterilization')
         assert set(_cons_codes.keys()) == set(self.states_that_may_require_HSI_to_switch_to)
-        # items used when initiating a modern reliable method after not using or switching from non-reliable method
         _cons_codes['co_initiation'] = get_items_from_pkg('Contraception initiation')
 
         return _cons_codes