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

vivarium_public_health/population/mortality.py

@@ -44,25 +44,81 @@ on mortality is calculated, by subtracting the raw unmodeled csmr before adding
 back the modified unmodeled csmr.
 
 """
-from typing import List, Optional, Union
+
+from typing import Any, Dict, List, Optional, Union
 
 import pandas as pd
 from vivarium import Component
 from vivarium.framework.engine import Builder
 from vivarium.framework.event import Event
-from vivarium.framework.lookup import LookupTable
 from vivarium.framework.population import SimulantData
 from vivarium.framework.randomness import RandomnessStream
 from vivarium.framework.values import Pipeline, list_combiner, union_post_processor
 
+from vivarium_public_health.utilities import get_lookup_columns
+
 
 class Mortality(Component):
-    CONFIGURATION_DEFAULTS = {"unmodeled_causes": []}
+    """
+    This is the mortality component which models sources of mortality for a model.
+    THe component models all cause mortality and allows for disease models to contribute
+    cause specific mortality. Data used by this class should be supplied in the artifact
+    and is configurable in the configuration to build lookup tables. For instance, let's
+    say we want to use sex and hair color to build a lookup table for all cause mortality.
+
+    .. code-block:: yaml
+
+       configuration:
+           mortality:
+               all_cause_mortality_rate:
+                   categorical_columns: ["sex", "hair_color"]
+
+    Similarly, we can do the same thing for unmodeled causes. Here is an example:
+
+    .. code-block:: yaml
+
+       configuration:
+           mortality:
+               unmodeled_cause_specific_mortality_rate:
+                   unmodeled_causes: ["maternal_disorders", maternal_hemorrhage]
+                   categorical_columns: ["sex", "hair_color"]
+
+    Or if we wanted to make the data a scalar value for all cause mortality rate we could
+    configure that as well.
+
+    .. code-block:: yaml
+
+       configuration:
+           mortality:
+               all_cause_mortality_rate:
+                   value: 0.01
+
+    """
 
     ##############
     # Properties #
     ##############
 
+    @property
+    def configuration_defaults(self) -> Dict[str, Any]:
+        return {
+            "mortality": {
+                "data_sources": {
+                    "all_cause_mortality_rate": "cause.all_causes.cause_specific_mortality_rate",
+                    "unmodeled_cause_specific_mortality_rate": "self::load_unmodeled_csmr",
+                    "life_expectancy": "population.theoretical_minimum_risk_life_expectancy",
+                },
+                "unmodeled_causes": [],
+            },
+        }
+
+    @property
+    def standard_lookup_tables(self) -> List[str]:
+        return [
+            "all_cause_mortality_rate",
+            "life_expectancy",
+        ]
+
     @property
     def columns_created(self) -> List[str]:
         return [self.cause_of_death_column_name, self.years_of_life_lost_column_name]
@@ -97,14 +153,10 @@ class Mortality(Component):
         self.clock = builder.time.clock()
 
         self.cause_specific_mortality_rate = self.get_cause_specific_mortality_rate(builder)
-        self.mortality_rate = self.get_mortality_rate(builder)
 
-        self.all_cause_mortality_rate = self.get_all_cause_mortality_rate(builder)
-        self.life_expectancy = self.get_life_expectancy(builder)
-
-        self._raw_unmodeled_csmr = self.get_raw_unmodeled_csmr(builder)
         self.unmodeled_csmr = self.get_unmodeled_csmr(builder)
         self.unmodeled_csmr_paf = self.get_unmodeled_csmr_paf(builder)
+        self.mortality_rate = self.get_mortality_rate(builder)
 
     #################
     # Setup methods #
@@ -120,90 +172,37 @@ class Mortality(Component):
         )
 
     def get_mortality_rate(self, builder: Builder) -> Pipeline:
+        required_columns = get_lookup_columns(
+            [
+                self.lookup_tables["all_cause_mortality_rate"],
+                self.lookup_tables["unmodeled_cause_specific_mortality_rate"],
+            ],
+        )
         return builder.value.register_rate_producer(
             self.mortality_rate_pipeline_name,
             source=self.calculate_mortality_rate,
-            requires_columns=["age", "sex"],
-        )
-
-    # noinspection PyMethodMayBeStatic
-    def get_all_cause_mortality_rate(self, builder: Builder) -> Union[LookupTable, Pipeline]:
-        """
-        Load all cause mortality rate data and build a lookup table or pipeline.
-
-        Parameters
-        ----------
-        builder
-            Interface to access simulation managers.
-
-        Returns
-        -------
-        Union[LookupTable, Pipeline]
-            A lookup table or pipeline returning the all cause mortality rate.
-        """
-        acmr_data = builder.data.load("cause.all_causes.cause_specific_mortality_rate")
-        return builder.lookup.build_table(
-            acmr_data, key_columns=["sex"], parameter_columns=["age", "year"]
+            requires_columns=required_columns,
         )
 
-    # noinspection PyMethodMayBeStatic
-    def get_life_expectancy(self, builder: Builder) -> Union[LookupTable, Pipeline]:
-        """
-        Load life expectancy data and build a lookup table or pipeline.
-
-        Parameters
-        ----------
-        builder
-            Interface to access simulation managers.
-
-        Returns
-        -------
-        Union[LookupTable, Pipeline]
-            A lookup table or pipeline returning the life expectancy.
-        """
-        life_expectancy_data = builder.data.load(
-            "population.theoretical_minimum_risk_life_expectancy"
-        )
-        return builder.lookup.build_table(life_expectancy_data, parameter_columns=["age"])
-
-    # noinspection PyMethodMayBeStatic
-    def get_raw_unmodeled_csmr(self, builder: Builder) -> Union[LookupTable, Pipeline]:
-        """
-        Load unmodeled cause specific mortality rate data and build a lookup
-        table or pipeline.
-
-        Parameters
-        ----------
-        builder
-            Interface to access simulation managers.
-
-        Returns
-        -------
-        Union[LookupTable, Pipeline]
-            A lookup table or pipeline returning the unmodeled csmr.
-        """
-        unmodeled_causes = builder.configuration.unmodeled_causes
+    def load_unmodeled_csmr(self, builder: Builder) -> Union[float, pd.DataFrame]:
+        # todo validate that all data have the same columns
         raw_csmr = 0.0
-        for idx, cause in enumerate(unmodeled_causes):
+        for idx, cause in enumerate(builder.configuration[self.name].unmodeled_causes):
             csmr = f"cause.{cause}.cause_specific_mortality_rate"
             if 0 == idx:
                 raw_csmr = builder.data.load(csmr)
             else:
                 raw_csmr.loc[:, "value"] += builder.data.load(csmr).value
-
-        additional_parameters = (
-            {"key_columns": ["sex"], "parameter_columns": ["age", "year"]}
-            if unmodeled_causes
-            else {}
-        )
-
-        return builder.lookup.build_table(raw_csmr, **additional_parameters)
+        return raw_csmr
 
     def get_unmodeled_csmr(self, builder: Builder) -> Pipeline:
+        required_columns = get_lookup_columns(
+            [self.lookup_tables["unmodeled_cause_specific_mortality_rate"]]
+        )
         return builder.value.register_value_producer(
             self.unmodeled_csmr_pipeline_name,
             source=self.get_unmodeled_csmr_source,
-            requires_columns=["age", "sex"],
+            requires_columns=required_columns,
         )
 
     def get_unmodeled_csmr_paf(self, builder: Builder) -> Pipeline:
@@ -248,7 +247,9 @@ class Mortality(Component):
             )
             pop.loc[deaths, "alive"] = "dead"
             pop.loc[deaths, "exit_time"] = event.time
-            pop.loc[deaths, "years_of_life_lost"] = self.life_expectancy(deaths)
+            pop.loc[deaths, "years_of_life_lost"] = self.lookup_tables["life_expectancy"](
+                deaths
+            )
             pop.loc[deaths, "cause_of_death"] = cause_of_death
             self.population_view.update(pop)
 
@@ -257,9 +258,11 @@ class Mortality(Component):
     ##################################
 
     def calculate_mortality_rate(self, index: pd.Index) -> pd.DataFrame:
-        acmr = self.all_cause_mortality_rate(index)
+        acmr = self.lookup_tables["all_cause_mortality_rate"](index)
         modeled_csmr = self.cause_specific_mortality_rate(index)
-        unmodeled_csmr_raw = self._raw_unmodeled_csmr(index)
+        unmodeled_csmr_raw = self.lookup_tables["unmodeled_cause_specific_mortality_rate"](
+            index
+        )
         unmodeled_csmr = self.unmodeled_csmr(index)
         cause_deleted_mortality_rate = (
             acmr - modeled_csmr - unmodeled_csmr_raw + unmodeled_csmr
@@ -267,6 +270,6 @@ class Mortality(Component):
         return pd.DataFrame({"other_causes": cause_deleted_mortality_rate})
 
     def get_unmodeled_csmr_source(self, index: pd.Index) -> pd.Series:
-        raw_csmr = self._raw_unmodeled_csmr(index)
+        raw_csmr = self.lookup_tables["unmodeled_cause_specific_mortality_rate"](index)
         paf = self.unmodeled_csmr_paf(index)
         return raw_csmr * (1 - paf)

vivarium_public_health/{metrics → results}/__init__.py

@@ -1,5 +1,7 @@
+from .columns import COLUMNS
 from .disability import DisabilityObserver
 from .disease import DiseaseObserver
 from .mortality import MortalityObserver
+from .observer import PublicHealthObserver
 from .risk import CategoricalRiskObserver
 from .stratification import ResultsStratifier

vivarium_public_health/results/columns.py

@@ -0,0 +1,22 @@
+from typing import NamedTuple
+
+from vivarium.framework.results import VALUE_COLUMN
+
+
+class __Columns(NamedTuple):
+    """column names"""
+
+    VALUE: str = VALUE_COLUMN
+    MEASURE: str = "measure"
+    TRANSITION: str = "transition"
+    STATE: str = "state"
+    ENTITY_TYPE: str = "entity_type"
+    SUB_ENTITY: str = "sub_entity"
+    ENTITY: str = "entity"
+
+    @property
+    def name(self) -> str:
+        return "columns"
+
+
+COLUMNS = __Columns()

vivarium_public_health/results/disability.py

@@ -0,0 +1,187 @@
+"""
+===================
+Disability Observer
+===================
+
+This module contains tools for observing years lived with disability (YLDs)
+in the simulation.
+
+"""
+
+from typing import Any, List, Union
+
+import pandas as pd
+from layered_config_tree import LayeredConfigTree
+from loguru import logger
+from pandas.api.types import CategoricalDtype
+from vivarium.framework.engine import Builder
+from vivarium.framework.values import Pipeline, list_combiner, union_post_processor
+
+from vivarium_public_health.disease import DiseaseState, RiskAttributableDisease
+from vivarium_public_health.results.columns import COLUMNS
+from vivarium_public_health.results.observer import PublicHealthObserver
+from vivarium_public_health.results.simple_cause import SimpleCause
+from vivarium_public_health.utilities import to_years
+
+
+class DisabilityObserver(PublicHealthObserver):
+    """Counts years lived with disability.
+
+    By default, this counts both aggregate and cause-specific years lived
+    with disability over the full course of the simulation.
+
+    In the model specification, your configuration for this component should
+    be specified as, e.g.:
+
+    .. code-block:: yaml
+
+        configuration:
+            observers:
+                disability:
+                    exclude:
+                        - "sex"
+                    include:
+                        - "sample_stratification"
+    """
+
+    ##############
+    # Properties #
+    ##############
+
+    @property
+    def disability_classes(self) -> list[type]:
+        """The classes to be considered for causes of disability."""
+        return [DiseaseState, RiskAttributableDisease]
+
+    #####################
+    # Lifecycle methods #
+    #####################
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.disability_weight_pipeline_name = "all_causes.disability_weight"
+
+    #################
+    # Setup methods #
+    #################
+
+    def setup(self, builder: Builder) -> None:
+        self.step_size = pd.Timedelta(days=builder.configuration.time.step_size)
+        self.disability_weight = self.get_disability_weight_pipeline(builder)
+        self.set_causes_of_disability(builder)
+
+    def set_causes_of_disability(self, builder: Builder) -> None:
+        """Set the causes of disability to be observed by removing any excluded
+        via the model spec from the list of all disability class causes. We implement
+        exclusions here because disabilities are unique in that they are not
+        registered stratifications and so cannot be excluded during the stratification
+        call like other categories.
+        """
+        causes_of_disability = builder.components.get_components_by_type(
+            self.disability_classes
+        )
+        # Convert to SimpleCause instances and add on all_causes
+        causes_of_disability = [
+            SimpleCause.create_from_disease_state(cause) for cause in causes_of_disability
+        ] + [SimpleCause("all_causes", "all_causes", "cause")]
+
+        excluded_causes = (
+            builder.configuration.stratification.excluded_categories.to_dict().get(
+                "disability", []
+            )
+        )
+
+        # Handle exclusions that don't exist in the list of causes
+        cause_names = [cause.state_id for cause in causes_of_disability]
+        unknown_exclusions = set(excluded_causes) - set(cause_names)
+        if len(unknown_exclusions) > 0:
+            raise ValueError(
+                f"Excluded 'disability' causes {unknown_exclusions} not found in "
+                f"expected categories categories: {cause_names}"
+            )
+
+        # Drop excluded causes
+        if excluded_causes:
+            logger.debug(
+                f"'disability' has category exclusion requests: {excluded_causes}\n"
+                "Removing these from the allowable categories."
+            )
+        self.causes_of_disability = [
+            cause for cause in causes_of_disability if cause.state_id not in excluded_causes
+        ]
+
+    def get_configuration(self, builder: Builder) -> LayeredConfigTree:
+        return builder.configuration.stratification.disability
+
+    def register_observations(self, builder: Builder) -> None:
+        cause_pipelines = [
+            f"{cause.state_id}.disability_weight" for cause in self.causes_of_disability
+        ]
+        self.register_adding_observation(
+            builder=builder,
+            name="ylds",
+            pop_filter='tracked == True and alive == "alive"',
+            when="time_step__prepare",
+            requires_columns=["alive"],
+            requires_values=cause_pipelines,
+            additional_stratifications=self.configuration.include,
+            excluded_stratifications=self.configuration.exclude,
+            aggregator_sources=cause_pipelines,
+            aggregator=self.disability_weight_aggregator,
+        )
+
+    def get_disability_weight_pipeline(self, builder: Builder) -> Pipeline:
+        return builder.value.register_value_producer(
+            self.disability_weight_pipeline_name,
+            source=lambda index: [pd.Series(0.0, index=index)],
+            preferred_combiner=list_combiner,
+            preferred_post_processor=union_post_processor,
+        )
+
+    ###############
+    # Aggregators #
+    ###############
+
+    def disability_weight_aggregator(self, dw: pd.DataFrame) -> Union[float, pd.Series]:
+        aggregated_dw = (dw * to_years(self.step_size)).sum().squeeze()
+        if isinstance(aggregated_dw, pd.Series):
+            aggregated_dw.index.name = "cause_of_disability"
+        return aggregated_dw
+
+    ##############################
+    # Results formatting methods #
+    ##############################
+
+    def format(self, measure: str, results: pd.DataFrame) -> pd.DataFrame:
+        """Format results. Note that ylds are unique in that we
+        can't stratify by cause of disability (because there can be multiple at
+        once), and so the results here are actually wide by disability weight
+        pipeline name.
+        """
+
+        # Drop the unused 'value' column and rename the pipeline names to causes
+        results = results.drop(columns=["value"]).rename(
+            columns={col: col.replace(".disability_weight", "") for col in results.columns},
+        )
+        # Get desired index names prior to stacking
+        idx_names = list(results.index.names) + [COLUMNS.SUB_ENTITY]
+        results = pd.DataFrame(results.stack(), columns=[COLUMNS.VALUE])
+        # Name the new index level
+        results.index.set_names(idx_names, inplace=True)
+        results = results.reset_index()
+        results[COLUMNS.SUB_ENTITY] = results[COLUMNS.SUB_ENTITY].astype(CategoricalDtype())
+        return results
+
+    def get_entity_type_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        entity_type_map = {
+            cause.state_id: cause.cause_type for cause in self.causes_of_disability
+        }
+        return results[COLUMNS.SUB_ENTITY].map(entity_type_map).astype(CategoricalDtype())
+
+    def get_entity_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        entity_map = {cause.state_id: cause.model for cause in self.causes_of_disability}
+        return results[COLUMNS.SUB_ENTITY].map(entity_map).astype(CategoricalDtype())
+
+    def get_sub_entity_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        # The sub-entity col was created in the 'format' method
+        return results[COLUMNS.SUB_ENTITY]

vivarium_public_health/results/disease.py

@@ -0,0 +1,222 @@
+"""
+================
+Disease Observer
+================
+
+This module contains tools for observing disease incidence and prevalence
+in the simulation.
+
+"""
+
+from typing import Any, Dict, List
+
+import pandas as pd
+from layered_config_tree import LayeredConfigTree
+from vivarium.framework.engine import Builder
+from vivarium.framework.event import Event
+from vivarium.framework.population import SimulantData
+
+from vivarium_public_health.results.columns import COLUMNS
+from vivarium_public_health.results.observer import PublicHealthObserver
+from vivarium_public_health.utilities import to_years
+
+
+class DiseaseObserver(PublicHealthObserver):
+    """Observes disease counts and person time for a cause.
+
+    By default, this observer computes aggregate disease state person time and
+    counts of disease events over the full course of the simulation. It can be
+    configured to add or remove stratification groups to the default groups
+    defined by a ResultsStratifier.
+
+    In the model specification, your configuration for this component should
+    be specified as, e.g.:
+
+    .. code-block:: yaml
+
+        configuration:
+            stratification:
+                cause_name:
+                    exclude:
+                        - "sex"
+                    include:
+                        - "sample_stratification"
+    """
+
+    ##############
+    # Properties #
+    ##############
+
+    @property
+    def configuration_defaults(self) -> Dict[str, Any]:
+        return {
+            "stratification": {
+                self.disease: super().configuration_defaults["stratification"][
+                    self.get_configuration_name()
+                ]
+            }
+        }
+
+    @property
+    def columns_created(self) -> List[str]:
+        return [self.previous_state_column_name]
+
+    @property
+    def columns_required(self) -> List[str]:
+        return [self.disease]
+
+    @property
+    def initialization_requirements(self) -> Dict[str, List[str]]:
+        return {
+            "requires_columns": [self.disease],
+        }
+
+    #####################
+    # Lifecycle methods #
+    #####################
+
+    def __init__(self, disease: str) -> None:
+        super().__init__()
+        self.disease = disease
+        self.previous_state_column_name = f"previous_{self.disease}"
+
+    #################
+    # Setup methods #
+    #################
+
+    def setup(self, builder: Builder) -> None:
+        self.step_size = builder.time.step_size()
+        self.disease_model = builder.components.get_component(f"disease_model.{self.disease}")
+        self.entity_type = self.disease_model.cause_type
+        self.entity = self.disease_model.cause
+        self.transition_stratification_name = f"transition_{self.disease}"
+
+    def get_configuration(self, builder: Builder) -> LayeredConfigTree:
+        return builder.configuration.stratification[self.disease]
+
+    def register_observations(self, builder: Builder) -> None:
+
+        self.register_disease_state_stratification(builder)
+        self.register_transition_stratification(builder)
+
+        pop_filter = 'alive == "alive" and tracked==True'
+        self.register_person_time_observation(builder, pop_filter)
+        self.register_transition_count_observation(builder, pop_filter)
+
+    def register_disease_state_stratification(self, builder: Builder) -> None:
+        builder.results.register_stratification(
+            self.disease,
+            [state.state_id for state in self.disease_model.states],
+            requires_columns=[self.disease],
+        )
+
+    def register_transition_stratification(self, builder: Builder) -> None:
+        transitions = [
+            str(transition) for transition in self.disease_model.transition_names
+        ] + ["no_transition"]
+        # manually append 'no_transition' as an excluded transition
+        excluded_categories = (
+            builder.configuration.stratification.excluded_categories.to_dict().get(
+                self.transition_stratification_name, []
+            )
+        ) + ["no_transition"]
+        builder.results.register_stratification(
+            self.transition_stratification_name,
+            categories=transitions,
+            excluded_categories=excluded_categories,
+            mapper=self.map_transitions,
+            requires_columns=[self.disease, self.previous_state_column_name],
+            is_vectorized=True,
+        )
+
+    def register_person_time_observation(self, builder: Builder, pop_filter: str) -> None:
+        self.register_adding_observation(
+            builder=builder,
+            name=f"person_time_{self.disease}",
+            pop_filter=pop_filter,
+            when="time_step__prepare",
+            requires_columns=["alive", self.disease],
+            additional_stratifications=self.configuration.include + [self.disease],
+            excluded_stratifications=self.configuration.exclude,
+            aggregator=self.aggregate_state_person_time,
+        )
+
+    def register_transition_count_observation(
+        self, builder: Builder, pop_filter: str
+    ) -> None:
+        self.register_adding_observation(
+            builder=builder,
+            name=f"transition_count_{self.disease}",
+            pop_filter=pop_filter,
+            requires_columns=[
+                self.previous_state_column_name,
+                self.disease,
+            ],
+            additional_stratifications=self.configuration.include
+            + [self.transition_stratification_name],
+            excluded_stratifications=self.configuration.exclude,
+        )
+
+    def map_transitions(self, df: pd.DataFrame) -> pd.Series:
+        transitions = pd.Series(index=df.index, dtype=str)
+        transition_mask = df[self.previous_state_column_name] != df[self.disease]
+        transitions[~transition_mask] = "no_transition"
+        transitions[transition_mask] = (
+            df[self.previous_state_column_name].astype(str)
+            + "_to_"
+            + df[self.disease].astype(str)
+        )
+        return transitions
+
+    ########################
+    # Event-driven methods #
+    ########################
+
+    def on_initialize_simulants(self, pop_data: SimulantData) -> None:
+        """Initialize the previous state column to the current state"""
+        pop = self.population_view.subview([self.disease]).get(pop_data.index)
+        pop[self.previous_state_column_name] = pop[self.disease]
+        self.population_view.update(pop)
+
+    def on_time_step_prepare(self, event: Event) -> None:
+        # This enables tracking of transitions between states
+        prior_state_pop = self.population_view.get(event.index)
+        prior_state_pop[self.previous_state_column_name] = prior_state_pop[self.disease]
+        self.population_view.update(prior_state_pop)
+
+    ###############
+    # Aggregators #
+    ###############
+
+    def aggregate_state_person_time(self, x: pd.DataFrame) -> float:
+        return len(x) * to_years(self.step_size())
+
+    ##############################
+    # Results formatting methods #
+    ##############################
+
+    def format(self, measure: str, results: pd.DataFrame) -> pd.DataFrame:
+        results = results.reset_index()
+        if "transition_count_" in measure:
+            sub_entity = self.transition_stratification_name
+        if "person_time_" in measure:
+            sub_entity = self.disease
+        results.rename(columns={sub_entity: COLUMNS.SUB_ENTITY}, inplace=True)
+        return results
+
+    def get_measure_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        if "transition_count_" in measure:
+            measure_name = "transition_count"
+        if "person_time_" in measure:
+            measure_name = "person_time"
+        return pd.Series(measure_name, index=results.index)
+
+    def get_entity_type_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        return pd.Series(self.entity_type, index=results.index)
+
+    def get_entity_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        return pd.Series(self.entity, index=results.index)
+
+    def get_sub_entity_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        # The sub-entity col was created in the 'format' method
+        return results[COLUMNS.SUB_ENTITY]