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

vivarium_public_health/results/mortality.py

@@ -0,0 +1,186 @@
+"""
+==================
+Mortality Observer
+==================
+
+This module contains tools for observing cause-specific and
+excess mortality in the simulation, including "other causes".
+
+"""
+
+from typing import Any, Dict, List
+
+import pandas as pd
+from layered_config_tree import LayeredConfigTree
+from pandas.api.types import CategoricalDtype
+from vivarium.framework.engine import Builder
+
+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
+
+
+class MortalityObserver(PublicHealthObserver):
+    """An observer for cause-specific deaths and ylls (including "other causes").
+
+    By default, this counts cause-specific deaths and years of life lost over
+    the full course of the simulation. It can be configured to add or remove
+    stratification groups to the default groups defined by a
+    :class:ResultsStratifier. The aggregate configuration key can be set to
+    True to aggregate all deaths and ylls into a single observation and remove
+    the stratification by cause of death to improve runtime.
+
+    In the model specification, your configuration for this component should
+    be specified as, e.g.:
+
+    .. code-block:: yaml
+
+        configuration:
+            stratification:
+                mortality:
+                    exclude:
+                        - "sex"
+                    include:
+                        - "sample_stratification"
+
+    This observer needs to access the has_excess_mortality attribute of the causes
+    we're observing, but this attribute gets defined in the setup of the cause models.
+    As a result, the model specification should list this observer after causes.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.required_death_columns = ["alive", "exit_time", "cause_of_death"]
+        self.required_yll_columns = [
+            "alive",
+            "cause_of_death",
+            "exit_time",
+            "years_of_life_lost",
+        ]
+
+    ##############
+    # Properties #
+    ##############
+
+    @property
+    def mortality_classes(self) -> list[type]:
+        return [DiseaseState, RiskAttributableDisease]
+
+    @property
+    def configuration_defaults(self) -> Dict[str, Any]:
+        """
+        A dictionary containing the defaults for any configurations managed by
+        this component.
+        """
+        config_defaults = super().configuration_defaults
+        config_defaults["stratification"][self.get_configuration_name()]["aggregate"] = False
+        return config_defaults
+
+    @property
+    def columns_required(self) -> List[str]:
+        return [
+            "alive",
+            "years_of_life_lost",
+            "cause_of_death",
+            "exit_time",
+        ]
+
+    #################
+    # Setup methods #
+    #################
+
+    def setup(self, builder: Builder) -> None:
+        self.clock = builder.time.clock()
+        self.set_causes_of_death(builder)
+
+    def set_causes_of_death(self, builder: Builder) -> None:
+        causes_of_death = [
+            cause
+            for cause in builder.components.get_components_by_type(
+                tuple(self.mortality_classes)
+            )
+            if cause.has_excess_mortality
+        ]
+
+        # Convert to SimpleCauses and add on other_causes and not_dead
+        self.causes_of_death = [
+            SimpleCause.create_from_disease_state(cause) for cause in causes_of_death
+        ] + [
+            SimpleCause("not_dead", "not_dead", "cause"),
+            SimpleCause("other_causes", "other_causes", "cause"),
+        ]
+
+    def get_configuration(self, builder: Builder) -> LayeredConfigTree:
+        return builder.configuration.stratification[self.get_configuration_name()]
+
+    def register_observations(self, builder: Builder) -> None:
+        pop_filter = 'alive == "dead" and tracked == True'
+        additional_stratifications = self.configuration.include
+        if not self.configuration.aggregate:
+            # manually append 'not_dead' as an excluded cause
+            excluded_categories = (
+                builder.configuration.stratification.excluded_categories.to_dict().get(
+                    "cause_of_death", []
+                )
+            ) + ["not_dead"]
+            builder.results.register_stratification(
+                "cause_of_death",
+                [cause.state_id for cause in self.causes_of_death],
+                excluded_categories=excluded_categories,
+                requires_columns=["cause_of_death"],
+            )
+            additional_stratifications += ["cause_of_death"]
+        self.register_adding_observation(
+            builder=builder,
+            name="deaths",
+            pop_filter=pop_filter,
+            requires_columns=self.required_death_columns,
+            additional_stratifications=additional_stratifications,
+            excluded_stratifications=self.configuration.exclude,
+            aggregator=self.count_deaths,
+        )
+        self.register_adding_observation(
+            builder=builder,
+            name="ylls",
+            pop_filter=pop_filter,
+            requires_columns=self.required_yll_columns,
+            additional_stratifications=additional_stratifications,
+            excluded_stratifications=self.configuration.exclude,
+            aggregator=self.calculate_ylls,
+        )
+
+    ###############
+    # Aggregators #
+    ###############
+
+    def count_deaths(self, x: pd.DataFrame) -> float:
+        died_of_cause = x["exit_time"] > self.clock()
+        return sum(died_of_cause)
+
+    def calculate_ylls(self, x: pd.DataFrame) -> float:
+        died_of_cause = x["exit_time"] > self.clock()
+        return x.loc[died_of_cause, "years_of_life_lost"].sum()
+
+    ##############################
+    # Results formatting methods #
+    ##############################
+
+    def format(self, measure: str, results: pd.DataFrame) -> pd.DataFrame:
+        results = results.reset_index()
+        if self.configuration.aggregate:
+            results[COLUMNS.ENTITY] = "all_causes"
+        else:
+            results.rename(columns={"cause_of_death": COLUMNS.ENTITY}, inplace=True)
+        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_death}
+        return results[COLUMNS.ENTITY].map(entity_type_map).astype(CategoricalDtype())
+
+    def get_entity_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        # The entity col was created in the 'format' method
+        return results[COLUMNS.ENTITY]
+
+    def get_sub_entity_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        return results[COLUMNS.ENTITY]

vivarium_public_health/results/observer.py

@@ -0,0 +1,78 @@
+from typing import Callable, List, Optional, Union
+
+import pandas as pd
+from vivarium.framework.engine import Builder
+from vivarium.framework.results import Observer
+
+from vivarium_public_health.results.columns import COLUMNS
+
+
+class PublicHealthObserver(Observer):
+    """A convenience class for typical public health observers. It provides
+    an entry point for registering the most common observation type
+    as well as standardized results formatting methods to overwrite as necessary.
+    """
+
+    def register_adding_observation(
+        self,
+        builder: Builder,
+        name,
+        pop_filter,
+        when: str = "collect_metrics",
+        requires_columns: List[str] = [],
+        requires_values: List[str] = [],
+        additional_stratifications: List[str] = [],
+        excluded_stratifications: List[str] = [],
+        aggregator_sources: Optional[List[str]] = None,
+        aggregator: Callable[[pd.DataFrame], Union[float, pd.Series]] = len,
+    ):
+        builder.results.register_adding_observation(
+            name=name,
+            pop_filter=pop_filter,
+            when=when,
+            requires_columns=requires_columns,
+            requires_values=requires_values,
+            results_formatter=self.format_results,
+            additional_stratifications=additional_stratifications,
+            excluded_stratifications=excluded_stratifications,
+            aggregator_sources=aggregator_sources,
+            aggregator=aggregator,
+        )
+
+    def format_results(self, measure: str, results: pd.DataFrame) -> pd.DataFrame:
+        """Top-level results formatter that calls standard sub-methods to be
+        overwritten as necessary.
+        """
+
+        results = self.format(measure, results)
+        results[COLUMNS.MEASURE] = self.get_measure_column(measure, results)
+        results[COLUMNS.ENTITY_TYPE] = self.get_entity_type_column(measure, results)
+        results[COLUMNS.ENTITY] = self.get_entity_column(measure, results)
+        results[COLUMNS.SUB_ENTITY] = self.get_sub_entity_column(measure, results)
+
+        ordered_columns = [
+            COLUMNS.MEASURE,
+            COLUMNS.ENTITY_TYPE,
+            COLUMNS.ENTITY,
+            COLUMNS.SUB_ENTITY,
+        ]
+        ordered_columns += [
+            c for c in results.columns if c not in ordered_columns + [COLUMNS.VALUE]
+        ]
+        ordered_columns += [COLUMNS.VALUE]
+        return results[ordered_columns]
+
+    def format(self, measure: str, results: pd.DataFrame) -> pd.DataFrame:
+        return results
+
+    def get_measure_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        return pd.Series(measure, index=results.index)
+
+    def get_entity_type_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        return pd.Series(None, index=results.index)
+
+    def get_entity_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        return pd.Series(None, index=results.index)
+
+    def get_sub_entity_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        return pd.Series(None, index=results.index)

vivarium_public_health/results/risk.py

@@ -0,0 +1,138 @@
+"""
+==============
+Risk Observers
+==============
+
+This module contains tools for observing risk exposure during the simulation.
+
+"""
+
+from typing import Any, Dict, List, Optional
+
+import pandas as pd
+from layered_config_tree import LayeredConfigTree
+from vivarium.framework.engine import Builder
+
+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 CategoricalRiskObserver(PublicHealthObserver):
+    """An observer for a categorical risk factor.
+
+    Observes category person time for a risk factor.
+
+    By default, this observer computes aggregate categorical person time
+    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:
+                risk_name:
+                    exclude:
+                        - "sex"
+                    include:
+                        - "sample_stratification"
+    """
+
+    ##############
+    # Properties #
+    ##############
+
+    @property
+    def configuration_defaults(self) -> Dict[str, Any]:
+        """
+        A dictionary containing the defaults for any configurations managed by
+        this component.
+        """
+        return {
+            "stratification": {
+                f"{self.risk}": super().configuration_defaults["stratification"][
+                    self.get_configuration_name()
+                ]
+            }
+        }
+
+    @property
+    def columns_required(self) -> Optional[List[str]]:
+        return ["alive"]
+
+    #####################
+    # Lifecycle methods #
+    #####################
+
+    def __init__(self, risk: str) -> None:
+        """
+        Parameters
+        ----------
+        risk: name of a risk
+
+        """
+        super().__init__()
+        self.risk = risk
+        self.exposure_pipeline_name = f"{self.risk}.exposure"
+
+    #################
+    # Setup methods #
+    #################
+
+    def setup(self, builder: Builder) -> None:
+        self.step_size = builder.time.step_size()
+        self.categories = builder.data.load(f"risk_factor.{self.risk}.categories")
+
+    def get_configuration(self, builder: Builder) -> LayeredConfigTree:
+        return builder.configuration.stratification[self.risk]
+
+    def register_observations(self, builder: Builder) -> None:
+        builder.results.register_stratification(
+            f"{self.risk}",
+            list(self.categories.keys()),
+            requires_values=[self.exposure_pipeline_name],
+        )
+        self.register_adding_observation(
+            builder=builder,
+            name=f"person_time_{self.risk}",
+            pop_filter=f'alive == "alive" and tracked==True',
+            when="time_step__prepare",
+            requires_columns=["alive"],
+            requires_values=[self.exposure_pipeline_name],
+            additional_stratifications=self.configuration.include + [self.risk],
+            excluded_stratifications=self.configuration.exclude,
+            aggregator=self.aggregate_risk_category_person_time,
+        )
+
+    ###############
+    # Aggregators #
+    ###############
+
+    def aggregate_risk_category_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()
+        results.rename(columns={self.risk: COLUMNS.SUB_ENTITY}, inplace=True)
+        return results
+
+    def get_measure_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        return pd.Series("person_time", index=results.index)
+
+    def get_entity_type_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        return pd.Series("rei", index=results.index)
+
+    def get_entity_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        return pd.Series(self.risk, 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]

vivarium_public_health/results/simple_cause.py

@@ -0,0 +1,18 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class SimpleCause:
+    """A simple dataclass to represent the bare minimum information needed
+    for observers, e.g. 'all_causes' as a cause of disability. It also
+    includes a class method to convert a provided disease state into a
+    ``SimpleCause`` instance.
+    """
+
+    state_id: str
+    model: str
+    cause_type: str
+
+    @classmethod
+    def create_from_disease_state(cls, disease_state: type) -> "SimpleCause":
+        return cls(disease_state.state_id, disease_state.model, disease_state.cause_type)

vivarium_public_health/{metrics → results}/stratification.py

@@ -7,6 +7,8 @@ This module contains tools for stratifying observed quantities
 by specified characteristics through the vivarium results interface.
 """
 
+from __future__ import annotations
+
 import pandas as pd
 from vivarium import Component
 from vivarium.framework.engine import Builder
@@ -33,14 +35,14 @@ class ResultsStratifier(Component):
         builder.results.register_stratification(
             "age_group",
             self.age_bins["age_group_name"].to_list(),
-            self.map_age_groups,
+            mapper=self.map_age_groups,
             is_vectorized=True,
             requires_columns=["age"],
         )
         builder.results.register_stratification(
             "current_year",
             [str(year) for year in range(self.start_year, self.end_year + 1)],
-            self.map_year,
+            mapper=self.map_year,
             is_vectorized=True,
             requires_columns=["current_time"],
         )
@@ -49,7 +51,7 @@ class ResultsStratifier(Component):
         # builder.results.register_stratification(
         #     "event_year",
         #     [str(year) for year in range(self.start_year, self.end_year + 1)],
-        #     self.map_year,
+        #     mapper=self.map_year,
         #     is_vectorized=True,
         #     requires_columns=["event_time"],
         # )
@@ -66,7 +68,7 @@ class ResultsStratifier(Component):
         # builder.results.register_stratification(
         #     "exit_year",
         #     [str(year) for year in range(self.start_year, self.end_year + 1)] + ["nan"],
-        #     self.map_year,
+        #     mapper=self.map_year,
         #     is_vectorized=True,
         #     requires_columns=["exit_time"],
         # )
@@ -78,13 +80,13 @@ class ResultsStratifier(Component):
     # Mappers #
     ###########
 
-    def map_age_groups(self, pop: pd.DataFrame) -> pd.Series:
+    def map_age_groups(self, pop: pd.DataFrame) -> pd.Series[str]:
         """Map age with age group name strings
 
         Parameters
         ----------
         pop
-            A DataFrame with one column, an age to be mapped to an age group name string
+            A pd.DataFrame with one column, an age to be mapped to an age group name string
 
         Returns
         ------
@@ -97,13 +99,13 @@ class ResultsStratifier(Component):
         return age_group
 
     @staticmethod
-    def map_year(pop: pd.DataFrame) -> pd.Series:
+    def map_year(pop: pd.DataFrame) -> pd.Series[str]:
         """Map datetime with year
 
         Parameters
         ----------
         pop
-            A DataFrame with one column, a datetime to be mapped to year
+            A pd.DataFrame with one column, a datetime to be mapped to year
 
         Returns
         ------

vivarium_public_health/risks/__init__.py

@@ -1,6 +1,5 @@
 from .base_risk import Risk
-from .distributions import get_distribution
-from .effect import RiskEffect
+from .effect import NonLogLinearRiskEffect, RiskEffect
 from .implementations.low_birth_weight_and_short_gestation import (
     LBWSGDistribution,
     LBWSGRisk,