Albrja/mic 5371/maternal sepsis tree

src/vivarium_gates_mncnh/components/__init__.py

@@ -1,4 +1,5 @@
 from vivarium_gates_mncnh.components.antenatal_care import AntenatalCare
+from vivarium_gates_mncnh.components.maternal_sepsis import MaternalSepsis
 from vivarium_gates_mncnh.components.observers import (
     ANCObserver,
     BirthObserver,

src/vivarium_gates_mncnh/components/antenatal_care.py

@@ -1,18 +1,16 @@
 from __future__ import annotations
 
-from typing import Callable
-
 import numpy as np
 import pandas as pd
 import scipy.stats as stats
 from vivarium import Component
 from vivarium.framework.engine import Builder
 from vivarium.framework.event import Event
 from vivarium.framework.population import SimulantData
-from vivarium.framework.state_machine import Machine, State, TransientState
+from vivarium.framework.state_machine import State, TransientState
 from vivarium.types import ClockTime
 
-from vivarium_gates_mncnh.constants import data_keys
+from vivarium_gates_mncnh.components.tree import DecisionTreeState, TreeMachine
 from vivarium_gates_mncnh.constants.data_values import (
     ANC_RATES,
     COLUMNS,
@@ -23,30 +21,6 @@
 from vivarium_gates_mncnh.utilities import get_location
 
 
-class TreeMachine(Machine):
-    def setup(self, builder: Builder) -> None:
-        super().setup(builder)
-        self._sim_step_name = builder.time.simulation_event_name()
-
-    def on_time_step(self, event: Event) -> None:
-        if self._sim_step_name() == SIMULATION_EVENT_NAMES.PREGNANCY:
-            super().on_time_step(event)
-
-
-class ANCState(TransientState):
-    def __init__(self) -> None:
-        super().__init__("attended_antental_care")
-
-    @property
-    def columns_required(self) -> list[str]:
-        return [COLUMNS.ATTENDED_CARE_FACILITY]
-
-    def transition_side_effect(self, index: pd.Index, _event_time: ClockTime) -> None:
-        pop = self.population_view.get(index)
-        pop[COLUMNS.ATTENDED_CARE_FACILITY] = True
-        self.population_view.update(pop)
-
-
 class UltrasoundState(TransientState):
     def __init__(self, ultrasound_type: str) -> None:
         super().__init__(f"{ultrasound_type}_ultrasound")
@@ -176,9 +150,11 @@ def _determine_lbw_identification(self, pop: pd.DataFrame) -> pd.Series:
         identification[lbw_index] = draws < identification_rates
         return identification
 
-    def create_anc_decision_tree(self) -> Machine:
+    def create_anc_decision_tree(self) -> TreeMachine:
         initial_state = State("initial")
-        attended_antental_care = ANCState()
+        attended_antental_care = DecisionTreeState(
+            "attended_antental_care", COLUMNS.ATTENDED_CARE_FACILITY, True
+        )
         gets_ultrasound = TransientState("gets_ultrasound")
         standard_ultasound = UltrasoundState(ULTRASOUND_TYPES.STANDARD)
         ai_assisted_ultrasound = UltrasoundState(ULTRASOUND_TYPES.AI_ASSISTED)
@@ -224,7 +200,7 @@ def create_anc_decision_tree(self) -> Machine:
         ai_assisted_ultrasound.add_transition(output_state=end_state)
 
         return TreeMachine(
-            "anc_state",
+            COLUMNS.ANC_STATE,
             [
                 initial_state,
                 attended_antental_care,
@@ -234,4 +210,5 @@ def create_anc_decision_tree(self) -> Machine:
                 end_state,
             ],
             initial_state,
+            time_step_name=SIMULATION_EVENT_NAMES.PREGNANCY,
         )

src/vivarium_gates_mncnh/components/maternal_sepsis.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+import pandas as pd
+from vivarium import Component
+from vivarium.framework.engine import Builder
+from vivarium.framework.event import Event
+from vivarium.framework.population import SimulantData
+
+from vivarium_gates_mncnh.constants import data_keys
+from vivarium_gates_mncnh.constants.data_values import (
+    COLUMNS,
+    PREGNANCY_OUTCOMES,
+    SIMULATION_EVENT_NAMES,
+)
+from vivarium_gates_mncnh.constants.metadata import ARTIFACT_INDEX_COLUMNS
+from vivarium_gates_mncnh.utilities import get_location
+
+
+class MaternalSepsis(Component):
+    @property
+    def configuration_defaults(self) -> dict:
+        return {self.name: {"data_sources": {"incidence_risk": self.load_incidence_risk}}}
+
+    @property
+    def columns_created(self):
+        return [COLUMNS.MATERNAL_SEPSIS]
+
+    @property
+    def columns_required(self):
+        return [COLUMNS.PREGNANCY_OUTCOME]
+
+    def setup(self, builder: Builder) -> None:
+        self._sim_step_name = builder.time.simulation_event_name()
+        self.randomness = builder.randomness.get_stream(self.name)
+        self.location = get_location(builder)
+
+    def on_initialize_simulants(self, pop_data: SimulantData) -> None:
+        anc_data = pd.DataFrame(
+            {
+                COLUMNS.MATERNAL_SEPSIS: False,
+            },
+            index=pop_data.index,
+        )
+
+        self.population_view.update(anc_data)
+
+    def on_time_step(self, event: Event) -> None:
+        if self._sim_step_name() != SIMULATION_EVENT_NAMES.MATERNAL_SEPSIS:
+            return
+
+        pop = self.population_view.get(event.index)
+        full_term = pop.loc[
+            pop[COLUMNS.PREGNANCY_OUTCOME].isin(
+                [PREGNANCY_OUTCOMES.STILLBIRTH_OUTCOME, PREGNANCY_OUTCOMES.LIVE_BIRTH_OUTCOME]
+            )
+        ]
+        sepsis_risk = self.lookup_tables["incidence_risk"](full_term.index)
+        got_sepsis = self.randomness.filter_for_probability(
+            full_term.index,
+            sepsis_risk,
+            "got_sepsis_choice",
+        )
+        pop.loc[got_sepsis, COLUMNS.MATERNAL_SEPSIS] = True
+        self.population_view.update(pop)
+
+    def load_incidence_risk(self, builder: Builder) -> pd.DataFrame:
+        raw_incidence = builder.data.load(
+            data_keys.MATERNAL_SEPSIS.RAW_INCIDENCE_RATE
+        ).set_index(ARTIFACT_INDEX_COLUMNS)
+        asfr = builder.data.load(data_keys.PREGNANCY.ASFR).set_index(ARTIFACT_INDEX_COLUMNS)
+        sbr = (
+            builder.data.load(data_keys.PREGNANCY.SBR)
+            .set_index("year_start")
+            .drop(columns=["year_end"])
+            .reindex(asfr.index, level="year_start")
+        )
+        birth_rate = (sbr + 1) * asfr
+        incidence_risk = (raw_incidence / birth_rate).fillna(0.0)
+        return incidence_risk.reset_index()

src/vivarium_gates_mncnh/components/pregnancy.py

@@ -56,7 +56,6 @@ def setup(self, builder: Builder):
             Interface to several simulation tools.
         """
         self.time_step = builder.time.step_size()
-        self._sim_step_name = builder.time.simulation_event_name()
         self.randomness = builder.randomness.get_stream(self.name)
         self.birth_outcome_probabilities = builder.value.register_value_producer(
             "birth_outcome_probabilities",

src/vivarium_gates_mncnh/components/tree.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+import pandas as pd
+from vivarium.framework.engine import Builder
+from vivarium.framework.event import Event
+from vivarium.framework.state_machine import Machine, State, TransientState
+from vivarium.types import ClockTime
+
+from vivarium_gates_mncnh.constants.data_values import SIMULATION_EVENT_NAMES
+
+
+class TreeMachine(Machine):
+    def __init__(
+        self,
+        state_column: str,
+        states: list[State],
+        initial_state=None,
+        time_step_name: str = "",
+    ):
+        super().__init__(state_column, states, initial_state)
+        # Time step name where the simulants will go through the decision tree
+        self._time_step_trigger = time_step_name
+
+    def setup(self, builder: Builder) -> None:
+        super().setup(builder)
+        self._sim_step_name = builder.time.simulation_event_name()
+
+    def on_time_step(self, event: Event) -> None:
+        if self._sim_step_name() == self._time_step_trigger:
+            super().on_time_step(event)
+
+
+class DecisionTreeState(TransientState):
+    def __init__(
+        self,
+        state_id: str,
+        update_col: str,
+        update_value: str | bool,
+    ) -> None:
+        super().__init__(state_id)
+        self.update_column = update_col
+        self.update_value = update_value
+
+    @property
+    def columns_required(self) -> list[str]:
+        return [self.update_column]
+
+    def transition_side_effect(self, index: pd.Index, _event_time: ClockTime) -> None:
+        pop = self.population_view.get(index)
+        pop[self.update_column] = self.update_value
+        self.population_view.update(pop)

src/vivarium_gates_mncnh/constants/data_values.py

@@ -78,6 +78,7 @@ class _Durations(NamedTuple):
 class _SimulationEventNames(NamedTuple):
     PREGNANCY = "pregnancy"
     INTRAPARTRUM = "intrapartum"
+    MATERNAL_SEPSIS = "maternal_sepsis"
     NEONATAL = "neonatal"
 
 
@@ -149,6 +150,9 @@ class __Columns(NamedTuple):
     ULTRASOUND_TYPE = "ultrasound_type"
     STATED_GESTATIONAL_AGE = "stated_gestational_age"
     SUCCESSFUL_LBW_IDENTIFICATION = "successful_lbw_identification"
+    ANC_STATE = "anc_state"
+    SEPSIS_STATE = "sepsis_state"
+    MATERNAL_SEPSIS = "maternal_sepsis"
 
 
 COLUMNS = __Columns()

src/vivarium_gates_mncnh/model_specifications/model_spec.yaml

@@ -15,12 +15,14 @@ components:
             - ResultsStratifier()
             - BirthObserver()
             - AntenatalCare()
+            - MaternalSepsis()
+            # Add model observers below here
             - ANCObserver()
 
 configuration:
     input_data:
         input_draw_number: 0
-        artifact_path: "/mnt/team/simulation_science/pub/models/vivarium_gates_mncnh/artifacts/anc/pakistan.hdf"
+        artifact_path: "/mnt/team/simulation_science/pub/models/vivarium_gates_mncnh/artifacts/sepis/ethiopia.hdf"
     interpolation:
         order: 0
         extrapolate: True
@@ -35,7 +37,8 @@ configuration:
             day: 1
         simulation_events: 
             - 'pregnancy'
-            - 'intrapartum'  
+            - 'intrapartum' 
+            - "maternal_sepsis"
             - 'neonatal'
 
     population: