vivarium-public-health 2.3.3__py3-none-any.whl → 3.0.1__py3-none-any.whl

vivarium_public_health/disease/transition.py

@@ -6,7 +6,8 @@ Disease Transitions
 This module contains tools to model transitions between disease states.
 
 """
-from typing import TYPE_CHECKING, Callable, Dict
+
+from typing import TYPE_CHECKING, Any, Callable, Dict, Union
 
 import pandas as pd
 from vivarium.framework.engine import Builder
@@ -14,6 +15,8 @@ from vivarium.framework.state_machine import Transition, Trigger
 from vivarium.framework.utilities import rate_to_probability
 from vivarium.framework.values import list_combiner, union_post_processor
 
+from vivarium_public_health.utilities import get_lookup_columns
+
 if TYPE_CHECKING:
     from vivarium_public_health.disease import BaseDiseaseState
 
@@ -25,8 +28,43 @@ class TransitionString(str):
         obj.from_state, obj.to_state = value.split("_TO_")
         return obj
 
+    def __getnewargs__(self):
+        return (self.from_state + "_TO_" + self.to_state,)
+
 
 class RateTransition(Transition):
+
+    ##############
+    # Properties #
+    ##############
+
+    @property
+    def configuration_defaults(self) -> Dict[str, Any]:
+        return {
+            f"{self.name}": {
+                "data_sources": {
+                    "transition_rate": "self::load_transition_rate",
+                },
+            },
+        }
+
+    @property
+    def transition_rate_pipeline_name(self) -> str:
+        if "incidence_rate" in self._get_data_functions:
+            pipeline_name = f"{self.output_state.state_id}.incidence_rate"
+        elif "remission_rate" in self._get_data_functions:
+            pipeline_name = f"{self.input_state.state_id}.remission_rate"
+        elif "transition_rate" in self._get_data_functions:
+            pipeline_name = (
+                f"{self.input_state.state_id}_to_{self.output_state.state_id}.transition_rate"
+            )
+        else:
+            raise ValueError(
+                "Cannot determine rate_transition pipeline name: "
+                "no valid data functions supplied."
+            )
+        return pipeline_name
+
     #####################
     # Lifecycle methods #
     #####################
@@ -47,19 +85,16 @@ class RateTransition(Transition):
 
     # noinspection PyAttributeOutsideInit
     def setup(self, builder: Builder) -> None:
-        rate_data, pipeline_name = self.load_transition_rate_data(builder)
-        self.base_rate = builder.lookup.build_table(
-            rate_data, key_columns=["sex"], parameter_columns=["age", "year"]
-        )
+        lookup_columns = get_lookup_columns([self.lookup_tables["transition_rate"]])
         self.transition_rate = builder.value.register_rate_producer(
-            pipeline_name,
+            self.transition_rate_pipeline_name,
             source=self.compute_transition_rate,
-            requires_columns=["age", "sex", "alive"],
-            requires_values=[f"{pipeline_name}.paf"],
+            requires_columns=lookup_columns + ["alive"],
+            requires_values=[f"{self.transition_rate_pipeline_name}.paf"],
         )
         paf = builder.lookup.build_table(0)
         self.joint_paf = builder.value.register_value_producer(
-            f"{pipeline_name}.paf",
+            f"{self.transition_rate_pipeline_name}.paf",
             source=lambda index: [paf(index)],
             preferred_combiner=list_combiner,
             preferred_post_processor=union_post_processor,
@@ -71,27 +106,22 @@ class RateTransition(Transition):
     # Setup methods #
     #################
 
-    def load_transition_rate_data(self, builder):
+    def load_transition_rate(self, builder: Builder) -> Union[float, pd.DataFrame]:
         if "incidence_rate" in self._get_data_functions:
             rate_data = self._get_data_functions["incidence_rate"](
                 builder, self.output_state.state_id
             )
-            pipeline_name = f"{self.output_state.state_id}.incidence_rate"
         elif "remission_rate" in self._get_data_functions:
             rate_data = self._get_data_functions["remission_rate"](
                 builder, self.input_state.state_id
             )
-            pipeline_name = f"{self.input_state.state_id}.remission_rate"
         elif "transition_rate" in self._get_data_functions:
             rate_data = self._get_data_functions["transition_rate"](
                 builder, self.input_state.state_id, self.output_state.state_id
             )
-            pipeline_name = (
-                f"{self.input_state.state_id}_to_{self.output_state.state_id}.transition_rate"
-            )
         else:
             raise ValueError("No valid data functions supplied.")
-        return rate_data, pipeline_name
+        return rate_data
 
     ##################################
     # Pipeline sources and modifiers #
@@ -100,7 +130,7 @@ class RateTransition(Transition):
     def compute_transition_rate(self, index: pd.Index) -> pd.Series:
         transition_rate = pd.Series(0.0, index=index)
         living = self.population_view.get(index, query='alive == "alive"').index
-        base_rates = self.base_rate(living)
+        base_rates = self.lookup_tables["transition_rate"](living)
         joint_paf = self.joint_paf(living)
         transition_rate.loc[living] = base_rates * (1 - joint_paf)
         return transition_rate
@@ -114,6 +144,21 @@ class RateTransition(Transition):
 
 
 class ProportionTransition(Transition):
+
+    ##############
+    # Properties #
+    ##############
+
+    @property
+    def configuration_defaults(self) -> Dict[str, Any]:
+        return {
+            f"{self.name}": {
+                "data_sources": {
+                    "proportion": "self::load_proportion",
+                },
+            },
+        }
+
     #####################
     # Lifecycle methods #
     #####################
@@ -132,16 +177,14 @@ class ProportionTransition(Transition):
             get_data_functions if get_data_functions is not None else {}
         )
 
-    # noinspection PyAttributeOutsideInit
-    def setup(self, builder):
-        super().setup(builder)
-        get_proportion_func = self._get_data_functions.get("proportion", None)
-        if get_proportion_func is None:
+    #################
+    # Setup methods #
+    #################
+
+    def load_proportion(self, builder: Builder) -> Union[float, pd.DataFrame]:
+        if "proportion" not in self._get_data_functions:
             raise ValueError("Must supply a proportion function")
-        self._proportion_data = get_proportion_func(builder, self.output_state.state_id)
-        self.proportion = builder.lookup.build_table(
-            self._proportion_data, key_columns=["sex"], parameter_columns=["age", "year"]
-        )
+        return self._get_data_functions["proportion"](builder, self.output_state.state_id)
 
     def _probability(self, index):
-        return self.proportion(index)
+        return self.lookup_tables["proportion"](index)

vivarium_public_health/mslt/delay.py

@@ -7,6 +7,7 @@ This module contains tools to represent delayed effects in a multi-state
 lifetable simulation.
 
 """
+
 from typing import Any, Dict, List, Optional
 
 import numpy as np

vivarium_public_health/mslt/disease.py

@@ -7,6 +7,7 @@ This module contains tools for modeling diseases in multi-state lifetable
 simulations.
 
 """
+
 from typing import Any, Dict, List, Optional
 
 import numpy as np

vivarium_public_health/mslt/intervention.py

@@ -7,6 +7,7 @@ This module contains tools for modeling interventions in multi-state lifetable
 simulations.
 
 """
+
 from typing import Any, Dict
 
 from vivarium import Component

vivarium_public_health/mslt/magic_wand_components.py

@@ -7,6 +7,7 @@ This module contains tools for making crude adjustments to rates in
 multi-state lifetable simulations.
 
 """
+
 from typing import Any, Dict
 
 from vivarium import Component

vivarium_public_health/mslt/observer.py

@@ -7,6 +7,7 @@ This module contains tools for recording various outputs of interest in
 multi-state lifetable simulations.
 
 """
+
 from typing import List, Optional
 
 import pandas as pd

vivarium_public_health/mslt/population.py

@@ -7,6 +7,7 @@ This module contains tools for modeling the core demography in
 multi-state lifetable simulations.
 
 """
+
 from typing import List, Optional
 
 import numpy as np

vivarium_public_health/plugins/parser.py

@@ -8,12 +8,14 @@ Component Configuration Parsers in this module are specialized implementations o
 that can parse configurations of components specific to the Vivarium Public
 Health package.
 """
+
 from importlib import import_module
 from typing import Any, Callable, Dict, List, Optional, Union
 
 import pandas as pd
+from layered_config_tree import LayeredConfigTree
 from pkg_resources import resource_filename
-from vivarium import Component, ConfigTree
+from vivarium import Component
 from vivarium.framework.components import ComponentConfigurationParser
 from vivarium.framework.components.parser import ParsingError
 from vivarium.framework.engine import Builder
@@ -84,7 +86,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
     default value will be used. The default triggered value is 'NOT_TRIGGERED'.
     """
 
-    def parse_component_config(self, component_config: ConfigTree) -> List[Component]:
+    def parse_component_config(self, component_config: LayeredConfigTree) -> List[Component]:
         """
         Parses the component configuration and returns a list of components.
 
@@ -137,7 +139,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
         Parameters
         ----------
         component_config
-            A ConfigTree defining the components to initialize.
+            A LayeredConfigTree defining the components to initialize.
 
         Returns
         -------
@@ -158,7 +160,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
                     source = f"{package}::{config_file}"
                     config_file = resource_filename(package, config_file)
 
-                    external_config = ConfigTree(config_file)
+                    external_config = LayeredConfigTree(config_file)
                     component_config.update(
                         external_config, layer="model_override", source=source
                     )
@@ -185,7 +187,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
     # Configuration methods #
     #########################
 
-    def _add_default_config_layer(self, causes_config: ConfigTree) -> None:
+    def _add_default_config_layer(self, causes_config: LayeredConfigTree) -> None:
         """
         Adds a default layer to the provided configuration that specifies
         default values for the cause model configuration.
@@ -193,7 +195,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
         Parameters
         ----------
         causes_config
-            A ConfigTree defining the cause model configurations
+            A LayeredConfigTree defining the cause model configurations
 
         Returns
         -------
@@ -223,7 +225,9 @@ class CausesConfigurationParser(ComponentConfigurationParser):
     # Cause model creation methods #
     ################################
 
-    def _get_cause_model_components(self, causes_config: ConfigTree) -> List[Component]:
+    def _get_cause_model_components(
+        self, causes_config: LayeredConfigTree
+    ) -> List[Component]:
         """
         Parses the cause model configuration and returns a list of
         `DiseaseModel` components.
@@ -231,7 +235,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
         Parameters
         ----------
         causes_config
-            A ConfigTree defining the cause model components to initialize
+            A LayeredConfigTree defining the cause model components to initialize
 
         Returns
         -------
@@ -241,6 +245,11 @@ class CausesConfigurationParser(ComponentConfigurationParser):
         cause_models = []
 
         for cause_name, cause_config in causes_config.items():
+            data_sources = None
+            if "data_sources" in cause_config:
+                data_sources_config = cause_config.data_sources
+                data_sources = self._get_data_sources(data_sources_config)
+
             states: Dict[str, BaseDiseaseState] = {
                 state_name: self._get_state(state_name, state_config, cause_name)
                 for state_name, state_config in cause_config.states.items()
@@ -256,14 +265,17 @@ class CausesConfigurationParser(ComponentConfigurationParser):
             model_type = import_by_path(cause_config.model_type)
             initial_state = states.get(cause_config.initial_state, None)
             model = model_type(
-                cause_name, initial_state=initial_state, states=list(states.values())
+                cause_name,
+                initial_state=initial_state,
+                states=list(states.values()),
+                get_data_functions=data_sources,
             )
             cause_models.append(model)
 
         return cause_models
 
     def _get_state(
-        self, state_name: str, state_config: ConfigTree, cause_name: str
+        self, state_name: str, state_config: LayeredConfigTree, cause_name: str
     ) -> BaseDiseaseState:
         """
         Parses a state configuration and returns an initialized `BaseDiseaseState`
@@ -274,7 +286,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
         state_name
             The name of the state to initialize
         state_config
-            A ConfigTree defining the state to initialize
+            A LayeredConfigTree defining the state to initialize
         cause_name
             The name of the cause to which the state belongs
 
@@ -296,10 +308,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
             state_kwargs["cleanup_function"] = lambda *x: x
         if "data_sources" in state_config:
             data_sources_config = state_config.data_sources
-            state_kwargs["get_data_functions"] = {
-                name: self._get_data_source(name, data_sources_config[name])
-                for name in data_sources_config.keys()
-            }
+            state_kwargs["get_data_functions"] = self._get_data_sources(data_sources_config)
 
         if state_config.state_type is not None:
             state_type = import_by_path(state_config.state_type)
@@ -319,7 +328,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
         self,
         source_state: BaseDiseaseState,
         sink_state: BaseDiseaseState,
-        transition_config: ConfigTree,
+        transition_config: LayeredConfigTree,
     ) -> None:
         """
         Adds a transition between two states.
@@ -331,7 +340,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
         sink_state
             The state the transition ends at
         transition_config
-            A `ConfigTree` defining the transition to add
+            A `LayeredConfigTree` defining the transition to add
 
         Returns
         -------
@@ -340,10 +349,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
         triggered = Trigger[transition_config.triggered]
         if "data_sources" in transition_config:
             data_sources_config = transition_config.data_sources
-            data_sources = {
-                name: self._get_data_source(name, data_sources_config[name])
-                for name in data_sources_config.keys()
-            }
+            data_sources = self._get_data_sources(data_sources_config)
         else:
             data_sources = None
 
@@ -363,6 +369,25 @@ class CausesConfigurationParser(ComponentConfigurationParser):
                 f" provided for transition '{transition_config}'."
             )
 
+    def _get_data_sources(
+        self, config: LayeredConfigTree
+    ) -> Dict[str, Callable[[Builder, Any], Any]]:
+        """
+        Parses a data sources configuration and returns a dictionary of data
+        sources.
+
+        Parameters
+        ----------
+        config
+            A LayeredConfigTree defining the data sources to initialize
+
+        Returns
+        -------
+        Dict[str, Callable[[Builder, Any], Any]]
+            A dictionary of data source getters
+        """
+        return {name: self._get_data_source(name, config[name]) for name in config.keys()}
+
     @staticmethod
     def _get_data_source(
         name: str, source: Union[str, float]
@@ -408,7 +433,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
     # Validation methods #
     ######################
 
-    _CAUSE_KEYS = {"model_type", "initial_state", "states", "transitions"}
+    _CAUSE_KEYS = {"model_type", "initial_state", "states", "transitions", "data_sources"}
     _STATE_KEYS = {
         "state_type",
         "cause_type",
@@ -420,6 +445,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
     }
 
     _DATA_SOURCE_KEYS = {
+        "cause": {"cause_specific_mortality_rate"},
         "state": {
             "prevalence",
             "birth_prevalence",
@@ -438,14 +464,14 @@ class CausesConfigurationParser(ComponentConfigurationParser):
     _TRANSITION_TYPE_KEYS = {"rate", "proportion", "dwell_time"}
 
     @staticmethod
-    def _validate_external_configuration(external_configuration: ConfigTree) -> None:
+    def _validate_external_configuration(external_configuration: LayeredConfigTree) -> None:
         """
         Validates the external configuration.
 
         Parameters
         ----------
         external_configuration
-            A ConfigTree defining the external configuration
+            A LayeredConfigTree defining the external configuration
 
         Returns
         -------
@@ -477,14 +503,14 @@ class CausesConfigurationParser(ComponentConfigurationParser):
         if error_messages:
             raise CausesParsingErrors(error_messages)
 
-    def _validate_causes_config(self, causes_config: ConfigTree) -> None:
+    def _validate_causes_config(self, causes_config: LayeredConfigTree) -> None:
         """
         Validates the cause model configuration.
 
         Parameters
         ----------
         causes_config
-            A ConfigTree defining the cause model configurations
+            A LayeredConfigTree defining the cause model configurations
 
         Returns
         -------
@@ -512,7 +538,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
         cause_name
             The name of the cause to validate
         cause_config
-            A ConfigTree defining the cause to validate
+            A LayeredConfigTree defining the cause to validate
 
         Returns
         -------
@@ -572,6 +598,10 @@ class CausesConfigurationParser(ComponentConfigurationParser):
                     cause_name, transition_name, transition_config, states_config
                 )
 
+        error_messages += self._validate_data_sources(
+            cause_config, cause_name, "cause", cause_name
+        )
+
         return error_messages
 
     def _validate_state(
@@ -587,7 +617,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
         state_name
             The name of the state to validate
         state_config
-            A ConfigTree defining the state to validate
+            A LayeredConfigTree defining the state to validate
 
         Returns
         -------
@@ -682,9 +712,9 @@ class CausesConfigurationParser(ComponentConfigurationParser):
         transition_name
             The name of the transition to validate
         transition_config
-            A ConfigTree defining the transition to validate
+            A LayeredConfigTree defining the transition to validate
         states_config
-            A ConfigTree defining the states for the cause
+            A LayeredConfigTree defining the states for the cause
 
         Returns
         -------
@@ -824,7 +854,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
         Parameters
         ----------
         config
-            A ConfigTree defining the configuration to validate
+            A LayeredConfigTree defining the configuration to validate
         cause_name
             The name of the cause to which the configuration belongs
         config_type

vivarium_public_health/population/add_new_birth_cohorts.py

@@ -6,6 +6,7 @@ Fertility Models
 This module contains several different models of fertility.
 
 """
+
 from typing import Dict, List, Optional
 
 import numpy as np
@@ -16,9 +17,7 @@ from vivarium.framework.event import Event
 from vivarium.framework.population import SimulantData
 
 from vivarium_public_health import utilities
-from vivarium_public_health.population.data_transformations import (
-    get_live_births_per_year,
-)
+from vivarium_public_health.population.data_transformations import get_live_births_per_year
 
 # TODO: Incorporate better data into gestational model (probably as a separate component)
 PREGNANCY_DURATION = pd.Timedelta(days=9 * utilities.DAYS_PER_MONTH)

vivarium_public_health/population/base_population.py

@@ -7,13 +7,14 @@ This module contains tools for sampling and assigning core demographic
 characteristics to simulants.
 
 """
+
 from typing import Callable, Dict, Iterable, List
 
 import numpy as np
 import pandas as pd
+from layered_config_tree.exceptions import ConfigurationKeyError
 from loguru import logger
 from vivarium import Component
-from vivarium.config_tree import ConfigurationKeyError
 from vivarium.framework.engine import Builder
 from vivarium.framework.event import Event
 from vivarium.framework.population import SimulantData