Sbachmei/mic 6736/vivarium 4 interventions (#130)

Author: stevebachmeier

src/vivarium_gates_nutrition_optimization/components/__init__.py

@@ -2,7 +2,8 @@
 from .intervention import MaternalInterventions
 from .maternal_bmi import MaternalBMIExposure
 from .maternal_disorders import MaternalDisorders, MaternalHemorrhage
-from .morbidity import BackgroundMorbidity
+
+# from .morbidity import BackgroundMorbidity
 from .mortality import MaternalMortality
 from .observers import (
     AnemiaObserver,

src/vivarium_gates_nutrition_optimization/components/disease.py

@@ -60,49 +60,52 @@ def _probability(self, index) -> pd.Series:
         return self.population_view.get_attributes(index, self.pipeline_name)
 
 
-class ParturitionExclusionState(DiseaseState):
-
-    # #####################
-    # # Lifecycle methods #
-    # #####################
-
-    def get_disability_weight_pipeline(self, builder: Builder) -> None:
-        builder.value.register_attribute_producer(
-            f"{self.state_id}.disability_weight",
-            source=self.compute_disability_weight,
-            required_resources=[
-                "is_alive",
-                "pregnancy",
-                self.model,
-                self.lookup_tables["disability_weight"],
-            ],
-        )
-
-    ##################################
-    # Pipeline sources and modifiers #
-    ##################################
-
-    def compute_disability_weight(self, index: pd.Index) -> pd.Series:
-        disability_weight = pd.Series(0, index=index)
-        raw_disability_weight = pd.Series(0, index=index)
-        with_condition = self.with_condition(index)
-        # FIXME: this is broken (self.base_disability_weight is undefined)
-        raw_disability_weight.loc[with_condition] = self.base_disability_weight(
-            with_condition
-        )
-
-        dw_map = {
-            models.NOT_PREGNANT_STATE_NAME: raw_disability_weight,
-            models.PREGNANT_STATE_NAME: raw_disability_weight,
-            ## Pause YLD accumulation during the parturition state
-            models.PARTURITION_STATE_NAME: pd.Series(0, index=index),
-            models.POSTPARTUM_STATE_NAME: raw_disability_weight,
-        }
-
-        pop = self.population_view.get(index)
-        alive = pop["is_alive"] == True
-        for state, dw in dw_map.items():
-            in_state = alive & (pop["pregnancy"] == state)
-            disability_weight[in_state] = dw.loc[in_state]
-
-        return disability_weight
+# NOTE: This component is only used by Morbidity which is an exploratory component
+#   and not fully functional. At the very least, this class has not been updated
+#   to work with the vivarium v4.0 / vph v5.0. Commenting it out.
+# class ParturitionExclusionState(DiseaseState):
+
+#     # #####################
+#     # # Lifecycle methods #
+#     # #####################
+
+#     def get_disability_weight_pipeline(self, builder: Builder) -> None:
+#         builder.value.register_attribute_producer(
+#             f"{self.state_id}.disability_weight",
+#             source=self.compute_disability_weight,
+#             required_resources=[
+#                 "is_alive",
+#                 "pregnancy",
+#                 self.model,
+#                 self.lookup_tables["disability_weight"],
+#             ],
+#         )
+
+#     ##################################
+#     # Pipeline sources and modifiers #
+#     ##################################
+
+#     def compute_disability_weight(self, index: pd.Index) -> pd.Series:
+#         disability_weight = pd.Series(0, index=index)
+#         raw_disability_weight = pd.Series(0, index=index)
+#         with_condition = self.with_condition(index)
+#         # FIXME: this is broken (self.base_disability_weight is undefined)
+#         raw_disability_weight.loc[with_condition] = self.base_disability_weight(
+#             with_condition
+#         )
+
+#         dw_map = {
+#             models.NOT_PREGNANT_STATE_NAME: raw_disability_weight,
+#             models.PREGNANT_STATE_NAME: raw_disability_weight,
+#             ## Pause YLD accumulation during the parturition state
+#             models.PARTURITION_STATE_NAME: pd.Series(0, index=index),
+#             models.POSTPARTUM_STATE_NAME: raw_disability_weight,
+#         }
+
+#         pop = self.population_view.get(index)
+#         alive = pop["is_alive"] == True
+#         for state, dw in dw_map.items():
+#             in_state = alive & (pop["pregnancy"] == state)
+#             disability_weight[in_state] = dw.loc[in_state]
+
+#         return disability_weight

src/vivarium_gates_nutrition_optimization/components/intervention.py

@@ -25,18 +25,6 @@ class MaternalInterventions(Component):
         }
     }
 
-    @property
-    def columns_created(self) -> List[str]:
-        return ["intervention"]
-
-    @property
-    def columns_required(self) -> List[str]:
-        return ["maternal_bmi_anemia_category", "tracked"]
-
-    @property
-    def initialization_requirements(self) -> Dict[str, List[str]]:
-        return {"requires_streams": [self.name], "requires_columns": self.columns_required}
-
     # noinspection PyAttributeOutsideInit
     def setup(self, builder: Builder) -> None:
         self.clock = builder.time.clock()
@@ -57,42 +45,46 @@ def setup(self, builder: Builder) -> None:
             data_keys.MATERNAL_INTERVENTIONS.IFA_EFFECT_SIZE
         ).value[0]
 
-        builder.value.register_value_modifier(
+        builder.value.register_attribute_modifier(
             "hemoglobin.exposure",
             self.update_exposure,
-            requires_columns=self.columns_created,
+            required_resources=["intervention"],
         )
 
-        builder.value.register_value_modifier(
+        builder.value.register_attribute_modifier(
             "birth_outcome_probabilities",
             self.adjust_stillbirth_probability,
-            requires_columns=self.columns_created,
+            required_resources=["intervention"],
         )
 
-    def on_initialize_simulants(self, pop_data: SimulantData) -> None:
-        pop = self.population_view.subview(["maternal_bmi_anemia_category"]).get(
-            pop_data.index
+        builder.population.register_initializer(
+            self.initialize_intervention,
+            "intervention",
+            required_resources=[self.randomness, "maternal_bmi_anemia_category"],
+        )
+
+    def initialize_intervention(self, pop_data: SimulantData) -> None:
+        categories = self.population_view.get_attributes(
+            pop_data.index, "maternal_bmi_anemia_category"
         )
 
         if self.scenario == "ifa":
             pop_update = pd.DataFrame(
                 {"intervention": "ifa"},
-                index=pop.index,
+                index=categories.index,
             )
         else:
             pop_update = pd.DataFrame(
                 {"intervention": None},
-                index=pop.index,
+                index=categories.index,
             )
             baseline_ifa = self.randomness.choice(
-                pop.index,
+                categories.index,
                 choices=[models.IFA_SUPPLEMENTATION, models.NO_TREATMENT],
                 p=[self.ifa_coverage, RESIDUAL_CHOICE],
                 additional_key="baseline_ifa",
             )
-            low_bmi = pop["maternal_bmi_anemia_category"].isin(
-                [models.LOW_BMI_NON_ANEMIC, models.LOW_BMI_ANEMIC]
-            )
+            low_bmi = categories.isin([models.LOW_BMI_NON_ANEMIC, models.LOW_BMI_ANEMIC])
             coverage = data_values.INTERVENTION_SCENARIO_COVERAGE.loc[self.scenario]
             pop_update["intervention"] = np.where(
                 low_bmi, coverage["low_bmi"], coverage["normal_bmi"]
@@ -107,24 +99,24 @@ def update_exposure(self, index, exposure):
         if self.clock() - self.start_date >= timedelta(
             days=data_values.DURATIONS.INTERVENTION_DELAY_DAYS
         ):
-            pop = self.population_view.get(index)
-            exposure.loc[pop["intervention"] == models.NO_TREATMENT] -= (
+            intervention = self.population_view.get_attributes(index, "intervention")
+            exposure.loc[intervention == models.NO_TREATMENT] -= (
                 self.ifa_coverage * self.ifa_effect_size
             )
-            exposure.loc[pop["intervention"] != models.NO_TREATMENT] += (
+            exposure.loc[intervention != models.NO_TREATMENT] += (
                 1 - self.ifa_coverage
             ) * self.ifa_effect_size
 
         return exposure
 
     def adjust_stillbirth_probability(self, index, birth_outcome_probabilities):
-        pop = self.population_view.subview(["intervention"]).get(index)
+        pop = self.population_view.get_attributes(index, "intervention")
         rrs = {
             models.MMS_SUPPLEMENTATION: self.mms_stillbirth_rr,
             models.BEP_SUPPLEMENTATION: self.bep_stillbirth_rr,
         }
         for intervention, rr in rrs.items():
-            on_treatment = pop["intervention"] == intervention
+            on_treatment = pop == intervention
             # Add spare probability onto live births first
             birth_outcome_probabilities.loc[
                 on_treatment, models.LIVE_BIRTH_OUTCOME

src/vivarium_gates_nutrition_optimization/components/morbidity.py

@@ -1,27 +1,30 @@
-from vivarium_public_health.disease import DiseaseModel, SusceptibleState
+# from vivarium_public_health.disease import DiseaseModel, SusceptibleState
 
-from vivarium_gates_nutrition_optimization.components.disease import (
-    ParturitionExclusionState,
-)
+# from vivarium_gates_nutrition_optimization.components.disease import (
+#     ParturitionExclusionState,
+# )
 
+# NOTE: I have not updated this component to work with the new lookup table
+#   configuration work so this component is not used in the model spec because
+#   it is currently broken. This component was an exploratory component and may
+#   be revisited in the future - albrja
+# NOTE: This has also not been updated to work with the vivarium v4.0 / vph v5.0.
+#   Commenting it out. - sbachmei
 
-def BackgroundMorbidity():
-    # NOTE: I have not updated this component to work with the new lookup table
-    # configuration work so this component is not used in the model spec because
-    # it is currently broken. This component was an exploratory component and may
-    # be revisited in the future - albrja
-    cause = "other_causes"
-    susceptible = SusceptibleState(cause)
-    with_condition = ParturitionExclusionState(
-        cause,
-        allow_self_transition=True,
-        get_data_functions={
-            "prevalence": lambda *_: 1.0,
-            "excess_mortality_rate": lambda *_: 0.0,
-        },
-    )
-    return DiseaseModel(
-        cause,
-        states=[susceptible, with_condition],
-        get_data_functions={"cause_specific_mortality_rate": lambda *_: 0.0},
-    )
+
+# def BackgroundMorbidity():
+#     cause = "other_causes"
+#     susceptible = SusceptibleState(cause)
+#     with_condition = ParturitionExclusionState(
+#         cause,
+#         allow_self_transition=True,
+#         get_data_functions={
+#             "prevalence": lambda *_: 1.0,
+#             "excess_mortality_rate": lambda *_: 0.0,
+#         },
+#     )
+#     return DiseaseModel(
+#         cause,
+#         states=[susceptible, with_condition],
+#         get_data_functions={"cause_specific_mortality_rate": lambda *_: 0.0},
+#     )

src/vivarium_gates_nutrition_optimization/components/observers.py

@@ -130,7 +130,7 @@ def register_observations(self, builder: Builder) -> None:
             name=f"person_time_maternal_bmi_anemia",
             pop_filter=f"is_alive == True",
             when="time_step__prepare",
-            requires_attributes=["is_alive", "maternal_bmi_anemia_category"],
+            requires_attributes=["maternal_bmi_anemia_category"],
             additional_stratifications=builder.configuration.stratification.maternal_bmi.include,
             excluded_stratifications=builder.configuration.stratification.maternal_bmi.exclude,
             aggregator=partial(aggregate_state_person_time, builder.time.step_size()()),
@@ -176,11 +176,11 @@ def register_observations(self, builder: Builder) -> None:
             builder=builder,
             name="intervention_count",
             pop_filter=(
-                'alive == "alive" and tracked == True and '
+                "is_alive == True and "
                 f'event_time > "{intervention_date}" and '
                 f'event_time <= "{intervention_date + builder.time.step_size()()}"'
             ),
-            requires_columns=["alive", "intervention", "event_time"],
+            requires_attributes=["intervention"],
             additional_stratifications=builder.configuration.stratification.maternal_intervention.include,
             excluded_stratifications=builder.configuration.stratification.maternal_intervention.exclude,
         )
@@ -225,7 +225,6 @@ def register_observations(self, builder: Builder) -> None:
         self.register_adding_observation(
             builder=builder,
             name=f"pregnancy_outcome_count",
-            pop_filter="",
             requires_attributes=["pregnancy_outcome"],
             additional_stratifications=builder.configuration.stratification.pregnancy_outcome.include,
             excluded_stratifications=builder.configuration.stratification.pregnancy_outcome.exclude,
@@ -296,7 +295,7 @@ def register_observations(self, builder: Builder) -> None:
                 f"and previous_pregnancy == '{models.PREGNANT_STATE_NAME}' "
                 f"and pregnancy == '{models.PARTURITION_STATE_NAME}'"
             ),
-            requires_columns=list(self.COL_MAPPING),
+            requires_attributes=list(self.COL_MAPPING),
             results_formatter=self.format,
         )
 

src/vivarium_gates_nutrition_optimization/components/pregnancy.py

@@ -134,13 +134,6 @@ def sample_full_term_durations(self, full_term_pop: pd.Index) -> pd.DataFrame:
         )
         return child_status
 
-    def get_dwell_time_pipeline(self, builder: Builder) -> Pipeline:
-        return builder.value.register_value_producer(
-            f"{self.state_id}.dwell_time",
-            source=lambda index: self.population_view.get(index)["pregnancy_duration"],
-            requires_columns=["age", "sex", "pregnancy_outcome"],
-        )
-
     def get_initial_event_times(self, pop_data: SimulantData) -> pd.DataFrame:
         """Overwrite the BaseDiseaseState method"""
         return pd.DataFrame(

src/vivarium_gates_nutrition_optimization/model_specifications/model_spec.yaml

@@ -16,16 +16,16 @@ components:
             - Hemoglobin() 
             - Anemia()
             - MaternalBMIExposure()
-    #         - MaternalInterventions() # run 3
+            - MaternalInterventions()
             - ResultsStratifier()
             - PregnancyObserver()
             - PregnancyOutcomeObserver()
-    #         - BirthObserver() # run 4
+            - BirthObserver()
             - MaternalMortalityObserver()
             - DisabilityObserver()
             - AnemiaObserver()
             - MaternalBMIObserver()
-    #         - MaternalInterventionObserver() # run 3
+            - MaternalInterventionObserver()
 
 configuration:
     input_data:
@@ -35,7 +35,12 @@ configuration:
         order: 0
         extrapolate: True
     intervention:
-        scenario: "baseline"
+        scenario: baseline
+        # scenario: zero_coverage
+        # scenario: mms
+        # scenario: ifa
+        # scenario: targeted_bep_ifa
+        # scenario: targeted_bep_mms
     randomness:
         map_size: 1_000_000
         key_columns: ['entrance_time', 'age']