vivarium-public-health 2.3.2__py3-none-any.whl → 3.0.0__py3-none-any.whl

vivarium_public_health/_version.py

@@ -1 +1 @@
-__version__ = "2.3.2"
+__version__ = "3.0.0"

vivarium_public_health/disease/model.py

@@ -8,7 +8,8 @@ function is to provide coordination across a set of disease states and
 transitions at simulation initialization and during transitions.
 
 """
-from typing import Callable, Dict, List, Optional
+
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import numpy as np
 import pandas as pd
@@ -27,10 +28,21 @@ class DiseaseModelError(VivariumError):
 
 
 class DiseaseModel(Machine):
+
     ##############
     # Properties #
     ##############
 
+    @property
+    def configuration_defaults(self) -> Dict[str, Any]:
+        return {
+            f"{self.name}": {
+                "data_sources": {
+                    "cause_specific_mortality_rate": "self::load_cause_specific_mortality_rate"
+                },
+            },
+        }
+
     @property
     def columns_created(self) -> List[str]:
         return [self.state_column]
@@ -88,13 +100,6 @@ class DiseaseModel(Machine):
 
         self.configuration_age_start = builder.configuration.population.initialization_age_min
         self.configuration_age_end = builder.configuration.population.initialization_age_max
-
-        cause_specific_mortality_rate = self.load_cause_specific_mortality_rate_data(builder)
-        self.cause_specific_mortality_rate = builder.lookup.build_table(
-            cause_specific_mortality_rate,
-            key_columns=["sex"],
-            parameter_columns=["age", "year"],
-        )
         builder.value.register_value_modifier(
             "cause_specific_mortality_rate",
             self.adjust_cause_specific_mortality_rate,
@@ -106,11 +111,13 @@ class DiseaseModel(Machine):
     # Setup methods #
     #################
 
-    def load_cause_specific_mortality_rate_data(self, builder):
+    def load_cause_specific_mortality_rate(
+        self, builder: Builder
+    ) -> Union[float, pd.DataFrame]:
         if "cause_specific_mortality_rate" not in self._get_data_functions:
             only_morbid = builder.data.load(f"cause.{self.cause}.restrictions")["yld_only"]
             if only_morbid:
-                csmr_data = 0
+                csmr_data = 0.0
             else:
                 csmr_data = builder.data.load(
                     f"{self.cause_type}.{self.cause}.cause_specific_mortality_rate"
@@ -130,8 +137,6 @@ class DiseaseModel(Machine):
 
         assert self.initial_state in {s.state_id for s in self.states}
 
-        # FIXME: this is a hack to figure out whether or not we're at the simulation start based on the fact that the
-        #  fertility components create this user data
         if pop_data.user_data["sim_state"] == "setup":  # simulation start
             if self.configuration_age_start != self.configuration_age_end != 0:
                 state_names, weights_bins = self.get_state_weights(
@@ -186,7 +191,7 @@ class DiseaseModel(Machine):
     ##################################
 
     def adjust_cause_specific_mortality_rate(self, index, rate):
-        return rate + self.cause_specific_mortality_rate(index)
+        return rate + self.lookup_tables["cause_specific_mortality_rate"](index)
 
     ####################
     # Helper functions #
@@ -198,18 +203,15 @@ class DiseaseModel(Machine):
             raise DiseaseModelError("Disease model must have exactly one SusceptibleState.")
         return susceptible_states[0].state_id
 
-    def get_state_weights(self, pop_index, prevalence_type):
-        states = [
-            s
-            for s in self.states
-            if hasattr(s, f"{prevalence_type}")
-            and getattr(s, f"{prevalence_type}") is not None
-        ]
+    def get_state_weights(
+        self, pop_index: pd.Index, prevalence_type: str
+    ) -> Tuple[List[str], Union[np.ndarray, None]]:
+        states = [state for state in self.states if state.lookup_tables.get(prevalence_type)]
 
         if not states:
             return states, None
 
-        weights = [getattr(s, f"{prevalence_type}")(pop_index) for s in states]
+        weights = [state.lookup_tables.get(prevalence_type)(pop_index) for state in states]
         for w in weights:
             w.reset_index(inplace=True, drop=True)
         weights += ((1 - np.sum(weights, axis=0)),)

vivarium_public_health/disease/models.py

@@ -7,6 +7,7 @@ This module contains a collection of frequently used parameterizations of
 disease models.
 
 """
+
 import pandas as pd
 
 from vivarium_public_health.disease.model import DiseaseModel

vivarium_public_health/disease/special_disease.py

@@ -6,6 +6,7 @@
 This module contains frequently used, but non-standard disease models.
 
 """
+
 import re
 from collections import namedtuple
 from operator import gt, lt
@@ -13,12 +14,13 @@ from typing import Any, Dict, List, Optional
 
 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.framework.values import list_combiner, union_post_processor
 
 from vivarium_public_health.disease.transition import TransitionString
-from vivarium_public_health.utilities import EntityString, is_non_zero
+from vivarium_public_health.utilities import EntityString, get_lookup_columns, is_non_zero
 
 
 class RiskAttributableDisease(Component):
@@ -85,25 +87,29 @@ class RiskAttributableDisease(Component):
         recoverable : True
     """
 
-    CONFIGURATION_DEFAULTS = {
-        "risk_attributable_disease": {
-            "threshold": None,
-            "mortality": True,
-            "recoverable": True,
-        }
-    }
-
     ##############
     # Properties #
     ##############
 
     @property
     def name(self):
-        return f"disease_model.{self.cause.name}"
+        return f"risk_attributable_disease.{self.cause.name}"
 
     @property
     def configuration_defaults(self) -> Dict[str, Any]:
-        return {self.cause.name: self.CONFIGURATION_DEFAULTS["risk_attributable_disease"]}
+        return {
+            self.name: {
+                "data_sources": {
+                    "raw_disability_weight": f"{self.cause}.disability_weight",
+                    "cause_specific_mortality_rate": "self::load_cause_specific_mortality_rate_data",
+                    "excess_mortality_rate": "self::load_excess_mortality_rate_data",
+                    "population_attributable_fraction": 0,
+                },
+                "threshold": None,
+                "mortality": True,
+                "recoverable": True,
+            }
+        }
 
     @property
     def columns_created(self) -> List[str]:
@@ -142,6 +148,8 @@ class RiskAttributableDisease(Component):
         self.cause = EntityString(cause)
         self.risk = EntityString(risk)
         self.state_column = self.cause.name
+        self.cause_type = "risk_attributable_disease"
+        self.model = self.risk.name
         self.state_id = self.cause.name
         self.diseased_event_time_column = f"{self.cause.name}_event_time"
         self.susceptible_event_time_column = f"susceptible_to_{self.cause.name}_event_time"
@@ -157,51 +165,40 @@ class RiskAttributableDisease(Component):
 
     # noinspection PyAttributeOutsideInit
     def setup(self, builder):
-        self.recoverable = builder.configuration[self.cause.name].recoverable
+        self.recoverable = builder.configuration[self.name].recoverable
         self.adjust_state_and_transitions()
         self.clock = builder.time.clock()
 
-        disability_weight_data = builder.data.load(f"{self.cause}.disability_weight")
-        self.has_disability = is_non_zero(disability_weight_data)
-        self.base_disability_weight = builder.lookup.build_table(
-            disability_weight_data, key_columns=["sex"], parameter_columns=["age", "year"]
-        )
         self.disability_weight = builder.value.register_value_producer(
             f"{self.cause.name}.disability_weight",
             source=self.compute_disability_weight,
-            requires_columns=["age", "sex", "alive", self.cause.name],
+            requires_columns=get_lookup_columns(
+                [self.lookup_tables["raw_disability_weight"]]
+            ),
         )
         builder.value.register_value_modifier(
-            "disability_weight", modifier=self.disability_weight
-        )
-
-        cause_specific_mortality_rate = self.load_cause_specific_mortality_rate_data(builder)
-        self.cause_specific_mortality_rate = builder.lookup.build_table(
-            cause_specific_mortality_rate,
-            key_columns=["sex"],
-            parameter_columns=["age", "year"],
+            "all_causes.disability_weight", modifier=self.disability_weight
         )
         builder.value.register_value_modifier(
             "cause_specific_mortality_rate",
             self.adjust_cause_specific_mortality_rate,
-            requires_columns=["age", "sex"],
+            requires_columns=get_lookup_columns(
+                [self.lookup_tables["cause_specific_mortality_rate"]]
+            ),
         )
+        self.has_excess_mortality = is_non_zero(self.lookup_tables["excess_mortality_rate"])
 
-        excess_mortality_data = self.load_excess_mortality_rate_data(builder)
-        self.has_excess_mortality = is_non_zero(excess_mortality_data)
-        self.base_excess_mortality_rate = builder.lookup.build_table(
-            excess_mortality_data, key_columns=["sex"], parameter_columns=["age", "year"]
-        )
         self.excess_mortality_rate = builder.value.register_value_producer(
             self.excess_mortality_rate_pipeline_name,
             source=self.compute_excess_mortality_rate,
-            requires_columns=["age", "sex", "alive", self.cause.name],
+            requires_columns=get_lookup_columns(
+                [self.lookup_tables["excess_mortality_rate"]]
+            ),
             requires_values=[self.excess_mortality_rate_paf_pipeline_name],
         )
-        paf = builder.lookup.build_table(0)
         self.joint_paf = builder.value.register_value_producer(
             self.excess_mortality_rate_paf_pipeline_name,
-            source=lambda idx: [paf(idx)],
+            source=lambda idx: [self.lookup_tables["population_attributable_fraction"](idx)],
             preferred_combiner=list_combiner,
             preferred_post_processor=union_post_processor,
         )
@@ -213,7 +210,7 @@ class RiskAttributableDisease(Component):
 
         distribution = builder.data.load(f"{self.risk}.distribution")
         exposure_pipeline = builder.value.get_value(f"{self.risk.name}.exposure")
-        threshold = builder.configuration[self.cause.name].threshold
+        threshold = builder.configuration[self.name].threshold
 
         self.filter_by_exposure = self.get_exposure_filter(
             distribution, exposure_pipeline, threshold
@@ -230,7 +227,7 @@ class RiskAttributableDisease(Component):
             )
 
     def load_cause_specific_mortality_rate_data(self, builder):
-        if builder.configuration[self.cause.name].mortality:
+        if builder.configuration[self.name].mortality:
             csmr_data = builder.data.load(
                 f"cause.{self.cause.name}.cause_specific_mortality_rate"
             )
@@ -239,7 +236,7 @@ class RiskAttributableDisease(Component):
         return csmr_data
 
     def load_excess_mortality_rate_data(self, builder):
-        if builder.configuration[self.cause.name].mortality:
+        if builder.configuration[self.name].mortality:
             emr_data = builder.data.load(f"cause.{self.cause.name}.excess_mortality_rate")
         else:
             emr_data = 0
@@ -331,13 +328,15 @@ class RiskAttributableDisease(Component):
     def compute_disability_weight(self, index):
         disability_weight = pd.Series(0.0, index=index)
         with_condition = self.with_condition(index)
-        disability_weight.loc[with_condition] = self.base_disability_weight(with_condition)
+        disability_weight.loc[with_condition] = self.lookup_tables["raw_disability_weight"](
+            with_condition
+        )
         return disability_weight
 
     def compute_excess_mortality_rate(self, index):
         excess_mortality_rate = pd.Series(0.0, index=index)
         with_condition = self.with_condition(index)
-        base_excess_mort = self.base_excess_mortality_rate(with_condition)
+        base_excess_mort = self.lookup_tables["excess_mortality_rate"](with_condition)
         joint_mediated_paf = self.joint_paf(with_condition)
         excess_mortality_rate.loc[with_condition] = base_excess_mort * (
             1 - joint_mediated_paf.values
@@ -345,7 +344,7 @@ class RiskAttributableDisease(Component):
         return excess_mortality_rate
 
     def adjust_cause_specific_mortality_rate(self, index, rate):
-        return rate + self.cause_specific_mortality_rate(index)
+        return rate + self.lookup_tables["cause_specific_mortality_rate"](index)
 
     def adjust_mortality_rate(self, index, rates_df):
         """Modifies the baseline mortality rate for a simulant if they are in this state.

vivarium_public_health/disease/state.py

@@ -6,12 +6,13 @@ Disease States
 This module contains tools to manage standard disease states.
 
 """
+
 from typing import Any, Callable, Dict, List, Optional
 
 import numpy as np
 import pandas as pd
 from vivarium.framework.engine import Builder
-from vivarium.framework.lookup import LookupTable, LookupTableData
+from vivarium.framework.lookup import LookupTableData
 from vivarium.framework.population import PopulationView, SimulantData
 from vivarium.framework.randomness import RandomnessStream
 from vivarium.framework.state_machine import State, Transient, Transition, Trigger
@@ -22,7 +23,7 @@ from vivarium_public_health.disease.transition import (
     RateTransition,
     TransitionString,
 )
-from vivarium_public_health.utilities import is_non_zero
+from vivarium_public_health.utilities import get_lookup_columns, is_non_zero
 
 
 class BaseDiseaseState(State):
@@ -273,6 +274,24 @@ class RecoveredState(NonDiseasedState):
 class DiseaseState(BaseDiseaseState):
     """State representing a disease in a state machine model."""
 
+    ##############
+    # Properties #
+    ##############
+
+    @property
+    def configuration_defaults(self) -> Dict[str, Any]:
+        return {
+            f"{self.name}": {
+                "data_sources": {
+                    "prevalence": "self::load_prevalence",
+                    "birth_prevalence": "self::load_birth_prevalence",
+                    "dwell_time": "self::load_dwell_time",
+                    "disability_weight": "self::load_disability_weight",
+                    "excess_mortality_rate": "self::load_excess_mortality_rate",
+                },
+            },
+        }
+
     #####################
     # Lifecycle methods #
     #####################
@@ -331,31 +350,15 @@ class DiseaseState(BaseDiseaseState):
         super().setup(builder)
         self.clock = builder.time.clock()
 
-        prevalence_data = self.load_prevalence_data(builder)
-        self.prevalence = self.get_prevalence(builder, prevalence_data)
-
-        birth_prevalence_data = self.load_birth_prevalence_data(builder)
-        self.birth_prevalence = self.get_birth_prevalence(builder, birth_prevalence_data)
-
-        dwell_time_data = self.load_dwell_time_data(builder)
-        self.dwell_time = self.get_dwell_time_pipeline(builder, dwell_time_data)
-
-        disability_weight_data = self.load_disability_weight_data(builder)
-        self.has_disability = is_non_zero(disability_weight_data)
-        self.base_disability_weight = self.get_base_disability_weight(
-            builder, disability_weight_data
-        )
-
+        self.dwell_time = self.get_dwell_time_pipeline(builder)
         self.disability_weight = self.get_disability_weight_pipeline(builder)
 
         builder.value.register_value_modifier(
-            "disability_weight", modifier=self.disability_weight
+            "all_causes.disability_weight", modifier=self.disability_weight
         )
 
-        excess_mortality_data = self.load_excess_mortality_rate_data(builder)
-        self.has_excess_mortality = is_non_zero(excess_mortality_data)
-        self.base_excess_mortality_rate = self.get_base_excess_mortality_rate(
-            builder, excess_mortality_data
+        self.has_excess_mortality = is_non_zero(
+            self.lookup_tables["excess_mortality_rate"].data
         )
         self.excess_mortality_rate = self.get_excess_mortality_rate_pipeline(builder)
         self.joint_paf = self.get_joint_paf(builder)
@@ -372,62 +375,19 @@ class DiseaseState(BaseDiseaseState):
     # Setup methods #
     #################
 
-    def load_prevalence_data(self, builder: Builder) -> LookupTableData:
+    def load_prevalence(self, builder: Builder) -> LookupTableData:
         if "prevalence" in self._get_data_functions:
             return self._get_data_functions["prevalence"](builder, self.state_id)
         else:
             return builder.data.load(f"{self.cause_type}.{self.state_id}.prevalence")
 
-    def get_prevalence(
-        self, builder: Builder, prevalence_data: LookupTableData
-    ) -> LookupTable:
-        """Builds a LookupTable for the prevalence of this state.
-
-        Parameters
-        ----------
-        builder
-            Interface to access simulation managers.
-        prevalence_data
-            The data to use to build the LookupTable.
-
-        Returns
-        -------
-        LookupTable
-            The LookupTable for the prevalence of this state.
-        """
-        return builder.lookup.build_table(
-            prevalence_data, key_columns=["sex"], parameter_columns=["age", "year"]
-        )
-
-    def load_birth_prevalence_data(self, builder: Builder) -> LookupTableData:
+    def load_birth_prevalence(self, builder: Builder) -> LookupTableData:
         if "birth_prevalence" in self._get_data_functions:
             return self._get_data_functions["birth_prevalence"](builder, self.state_id)
         else:
             return 0
 
-    def get_birth_prevalence(
-        self, builder: Builder, birth_prevalence_data: LookupTableData
-    ) -> LookupTable:
-        """
-        Builds a LookupTable for the birth prevalence of this state.
-
-        Parameters
-        ----------
-        builder
-            Interface to access simulation managers.
-        birth_prevalence_data
-            The data to use to build the LookupTable.
-
-        Returns
-        -------
-        LookupTable
-            The LookupTable for the birth prevalence of this state.
-        """
-        return builder.lookup.build_table(
-            birth_prevalence_data, key_columns=["sex"], parameter_columns=["year"]
-        )
-
-    def load_dwell_time_data(self, builder: Builder) -> LookupTableData:
+    def load_dwell_time(self, builder: Builder) -> LookupTableData:
         if "dwell_time" in self._get_data_functions:
             dwell_time = self._get_data_functions["dwell_time"](builder, self.state_id)
         else:
@@ -442,18 +402,15 @@ class DiseaseState(BaseDiseaseState):
 
         return dwell_time
 
-    def get_dwell_time_pipeline(
-        self, builder: Builder, dwell_time_data: LookupTableData
-    ) -> Pipeline:
+    def get_dwell_time_pipeline(self, builder: Builder) -> Pipeline:
+        required_columns = get_lookup_columns([self.lookup_tables["dwell_time"]])
         return builder.value.register_value_producer(
             f"{self.state_id}.dwell_time",
-            source=builder.lookup.build_table(
-                dwell_time_data, key_columns=["sex"], parameter_columns=["age", "year"]
-            ),
-            requires_columns=["age", "sex"],
+            source=self.lookup_tables["dwell_time"],
+            requires_columns=required_columns,
         )
 
-    def load_disability_weight_data(self, builder: Builder) -> LookupTableData:
+    def load_disability_weight(self, builder: Builder) -> LookupTableData:
         if "disability_weight" in self._get_data_functions:
             disability_weight = self._get_data_functions["disability_weight"](
                 builder, self.state_id
@@ -468,36 +425,15 @@ class DiseaseState(BaseDiseaseState):
 
         return disability_weight
 
-    def get_base_disability_weight(
-        self, builder: Builder, disability_weight_data: LookupTableData
-    ) -> LookupTable:
-        """
-        Builds a LookupTable for the base disability weight of this state.
-
-        Parameters
-        ----------
-        builder
-            Interface to access simulation managers.
-        disability_weight_data
-            The data to use to build the LookupTable.
-
-        Returns
-        -------
-        LookupTable
-            The LookupTable for the disability weight of this state.
-        """
-        return builder.lookup.build_table(
-            disability_weight_data, key_columns=["sex"], parameter_columns=["age", "year"]
-        )
-
     def get_disability_weight_pipeline(self, builder: Builder) -> Pipeline:
+        lookup_columns = get_lookup_columns([self.lookup_tables["disability_weight"]])
         return builder.value.register_value_producer(
             f"{self.state_id}.disability_weight",
             source=self.compute_disability_weight,
-            requires_columns=["age", "sex", "alive", self.model],
+            requires_columns=lookup_columns + ["alive", self.model],
         )
 
-    def load_excess_mortality_rate_data(self, builder: Builder) -> LookupTableData:
+    def load_excess_mortality_rate(self, builder: Builder) -> LookupTableData:
         if "excess_mortality_rate" in self._get_data_functions:
             return self._get_data_functions["excess_mortality_rate"](builder, self.state_id)
         elif builder.data.load(f"cause.{self.model}.restrictions")["yld_only"]:
@@ -507,33 +443,12 @@ class DiseaseState(BaseDiseaseState):
                 f"{self.cause_type}.{self.state_id}.excess_mortality_rate"
             )
 
-    def get_base_excess_mortality_rate(
-        self, builder: Builder, excess_mortality_data: LookupTableData
-    ) -> LookupTable:
-        """
-        Builds a LookupTable for the base excess mortality rate of this state.
-
-        Parameters
-        ----------
-        builder
-            Interface to access simulation managers.
-        excess_mortality_data
-            The data to use to build the LookupTable.
-
-        Returns
-        -------
-        LookupTable
-            The LookupTable for the base excess mortality rate of this state.
-        """
-        return builder.lookup.build_table(
-            excess_mortality_data, key_columns=["sex"], parameter_columns=["age", "year"]
-        )
-
     def get_excess_mortality_rate_pipeline(self, builder: Builder) -> Pipeline:
+        lookup_columns = get_lookup_columns([self.lookup_tables["excess_mortality_rate"]])
         return builder.value.register_rate_producer(
             self.excess_mortality_rate_pipeline_name,
             source=self.compute_excess_mortality_rate,
-            requires_columns=["age", "sex", "alive", self.model],
+            requires_columns=lookup_columns + ["alive", self.model],
             requires_values=[self.excess_mortality_rate_paf_pipeline_name],
         )
 
@@ -618,13 +533,15 @@ class DiseaseState(BaseDiseaseState):
         """
         disability_weight = pd.Series(0.0, index=index)
         with_condition = self.with_condition(index)
-        disability_weight.loc[with_condition] = self.base_disability_weight(with_condition)
+        disability_weight.loc[with_condition] = self.lookup_tables["disability_weight"](
+            with_condition
+        )
         return disability_weight
 
     def compute_excess_mortality_rate(self, index: pd.Index) -> pd.Series:
         excess_mortality_rate = pd.Series(0.0, index=index)
         with_condition = self.with_condition(index)
-        base_excess_mort = self.base_excess_mortality_rate(with_condition)
+        base_excess_mort = self.lookup_tables["excess_mortality_rate"](with_condition)
         joint_mediated_paf = self.joint_paf(with_condition)
         excess_mortality_rate.loc[with_condition] = base_excess_mort * (
             1 - joint_mediated_paf.values