vivarium-public-health 4.2.5__py3-none-any.whl → 4.3.0__py3-none-any.whl

vivarium_public_health/exposure/effect.py

@@ -0,0 +1,356 @@
+"""
+###################
+# Exposure Effect #
+###################
+
+This module contains tools for modeling the relationship between risk
+exposure models and disease models.
+
+"""
+import warnings
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from importlib import import_module
+from typing import Any
+
+import numpy as np
+import pandas as pd
+import scipy
+from layered_config_tree import ConfigurationError
+from vivarium import Component
+from vivarium.framework.engine import Builder
+from vivarium.framework.values import Pipeline
+
+from vivarium_public_health.risks.data_transformations import (
+    load_exposure_data,
+    pivot_categorical,
+)
+from vivarium_public_health.utilities import EntityString, TargetString, get_lookup_columns
+
+from .exposure import Exposure
+
+
+class ExposureEffect(Component, ABC):
+    """A component to model the effect of a risk-like factor on an affected target.
+
+    This component can source data either from builder.data or from parameters
+    supplied in the configuration.
+
+    """
+
+    def __init__(self, entity: str, target: str):
+        """
+
+        Parameters
+        ----------
+        entity
+            Type and name of exposure, supplied in the form
+            "entity_type.entity_name" where entity_type should be singular (e.g.,
+            exposure instead of exposures).
+        target
+            Type, name, and target rate of entity to be affected by risk factor,
+            supplied in the form "entity_type.entity_name.measure"
+            where entity_type should be singular (e.g., cause instead of causes).
+        """
+        super().__init__()
+        self.entity = EntityString(entity)
+        self.target = TargetString(target)
+
+        self._exposure_distribution_type = None
+        self.target_pipeline_name = f"{self.target.name}.{self.target.measure}"
+        self.target_paf_pipeline_name = f"{self.target_pipeline_name}.paf"
+
+    ###############
+    # Properties #
+    ##############
+
+    @property
+    def name(self) -> str:
+        return self.get_name(self.entity, self.target)
+
+    @abstractmethod
+    def get_name(self) -> Callable[[EntityString, TargetString], str]:
+        """Abstract property that must be implemented by subclasses to provide a naming function."""
+        raise NotImplementedError
+
+    @property
+    def configuration_defaults(self) -> dict[str, Any]:
+        """Default values for any configurations managed by this component."""
+        return {
+            self.name: {
+                "data_sources": {
+                    "relative_risk": f"{self.entity}.relative_risk",
+                    "population_attributable_fraction": f"{self.entity}.population_attributable_fraction",
+                },
+                "data_source_parameters": {
+                    "relative_risk": {},
+                },
+            }
+        }
+
+    @property
+    def is_exposure_categorical(self) -> bool:
+        return self._exposure_distribution_type in [
+            "dichotomous",
+            "ordered_polytomous",
+            "unordered_polytomous",
+        ]
+
+    #####################
+    # Lifecycle methods #
+    #####################
+
+    # noinspection PyAttributeOutsideInit
+    def setup(self, builder: Builder) -> None:
+        self.measure = self.get_exposure_callable(builder)
+
+        self._relative_risk_source = self.get_relative_risk_source(builder)
+        self.relative_risk = self.get_relative_risk_pipeline(builder)
+
+        self.register_target_modifier(builder)
+        self.register_paf_modifier(builder)
+
+    #################
+    # Setup methods #
+    #################
+
+    def setup_component(self, builder: Builder) -> None:
+        self.exposure_component = self._get_exposure_class(builder)
+        self.exposure_pipeline_name = (
+            f"{self.entity.name}.{self.exposure_component.exposure_type}"
+        )
+        super().setup_component(builder)
+
+    def build_all_lookup_tables(self, builder: Builder) -> None:
+        self._exposure_distribution_type = self.get_distribution_type(builder)
+
+        rr_data = self.load_relative_risk(builder)
+        rr_value_cols = None
+        if self.is_exposure_categorical:
+            rr_data, rr_value_cols = self.process_categorical_data(builder, rr_data)
+        self.lookup_tables["relative_risk"] = self.build_lookup_table(
+            builder, rr_data, rr_value_cols
+        )
+
+        paf_data = self.get_filtered_data(
+            builder, self.configuration.data_sources.population_attributable_fraction
+        )
+        self.lookup_tables["population_attributable_fraction"] = self.build_lookup_table(
+            builder, paf_data
+        )
+
+    def get_distribution_type(self, builder: Builder) -> str:
+        """Get the distribution type for the risk from the configuration."""
+        if self.exposure_component.distribution_type:
+            return self.exposure_component.distribution_type
+        return self.exposure_component.get_distribution_type(builder)
+
+    def load_relative_risk(
+        self,
+        builder: Builder,
+        configuration=None,
+    ) -> str | float | pd.DataFrame:
+        if configuration is None:
+            configuration = self.configuration
+
+        rr_source = configuration.data_sources.relative_risk
+        rr_dist_parameters = configuration.data_source_parameters.relative_risk.to_dict()
+
+        if isinstance(rr_source, str):
+            try:
+                distribution = getattr(import_module("scipy.stats"), rr_source)
+                rng = np.random.default_rng(builder.randomness.get_seed(self.name))
+                rr_data = distribution(**rr_dist_parameters).ppf(rng.random())
+            except AttributeError:
+                rr_data = self.get_filtered_data(builder, rr_source)
+            except TypeError:
+                raise ConfigurationError(
+                    f"Parameters {rr_dist_parameters} are not valid for distribution {rr_source}."
+                )
+        else:
+            rr_data = self.get_filtered_data(builder, rr_source)
+        return rr_data
+
+    def get_filtered_data(
+        self, builder: "Builder", data_source: str | float | pd.DataFrame
+    ) -> float | pd.DataFrame:
+        data = super().get_data(builder, data_source)
+
+        if isinstance(data, pd.DataFrame):
+            # filter data to only include the target entity and measure
+            correct_target_mask = True
+            columns_to_drop = []
+            if "affected_entity" in data.columns:
+                correct_target_mask &= data["affected_entity"] == self.target.name
+                columns_to_drop.append("affected_entity")
+            if "affected_measure" in data.columns:
+                correct_target_mask &= data["affected_measure"] == self.target.measure
+                columns_to_drop.append("affected_measure")
+            data = data[correct_target_mask].drop(columns=columns_to_drop)
+        return data
+
+    def process_categorical_data(
+        self, builder: Builder, rr_data: str | float | pd.DataFrame
+    ) -> tuple[str | float | pd.DataFrame, list[str]]:
+        if not isinstance(rr_data, pd.DataFrame):
+            exposed = builder.data.load("population.demographic_dimensions")
+            exposed[
+                "parameter"
+            ] = self.exposure_component.dichotomous_exposure_category_names.exposed
+            exposed["value"] = rr_data
+            unexposed = exposed.copy()
+            unexposed[
+                "parameter"
+            ] = self.exposure_component.dichotomous_exposure_category_names.unexposed
+            unexposed["value"] = 1
+            rr_data = pd.concat([exposed, unexposed], ignore_index=True)
+        if "parameter" in rr_data.index.names:
+            rr_data = rr_data.reset_index("parameter")
+
+        rr_value_cols = list(rr_data["parameter"].unique())
+        rr_data = pivot_categorical(builder, self.entity, rr_data, "parameter")
+        return rr_data, rr_value_cols
+
+    # todo currently this isn't being called. we need to properly set rrs if
+    #  the exposure has been rebinned
+    def rebin_relative_risk_data(
+        self, builder, relative_risk_data: pd.DataFrame
+    ) -> pd.DataFrame:
+        """Rebin relative risk data.
+
+        When the polytomous risk is rebinned, matching relative risk needs to be rebinned.
+        After rebinning, rr for both exposed and unexposed categories should be the weighted sum of relative risk
+        of the component categories where weights are relative proportions of exposure of those categories.
+        For example, if cat1, cat2, cat3 are exposed categories and cat4 is unexposed with exposure [0.1,0.2,0.3,0.4],
+        for the matching rr = [rr1, rr2, rr3, 1], rebinned rr for the rebinned cat1 should be:
+        (0.1 *rr1 + 0.2 * rr2 + 0.3* rr3) / (0.1+0.2+0.3)
+        """
+        if not self.entity in builder.configuration.to_dict():
+            return relative_risk_data
+
+        rebin_exposed_categories = set(builder.configuration[self.entity]["rebinned_exposed"])
+
+        if rebin_exposed_categories:
+            # todo make sure this works
+            exposure_data = load_exposure_data(builder, self.entity)
+            relative_risk_data = self._rebin_relative_risk_data(
+                relative_risk_data, exposure_data, rebin_exposed_categories
+            )
+
+        return relative_risk_data
+
+    def _rebin_relative_risk_data(
+        self,
+        relative_risk_data: pd.DataFrame,
+        exposure_data: pd.DataFrame,
+        rebin_exposed_categories: set,
+    ) -> pd.DataFrame:
+        cols = list(exposure_data.columns.difference(["value"]))
+
+        relative_risk_data = relative_risk_data.merge(exposure_data, on=cols)
+        relative_risk_data["value_x"] = relative_risk_data.value_x.multiply(
+            relative_risk_data.value_y
+        )
+        relative_risk_data.parameter = relative_risk_data["parameter"].map(
+            lambda p: "cat1" if p in rebin_exposed_categories else "cat2"
+        )
+        relative_risk_data = relative_risk_data.groupby(cols).sum().reset_index()
+        relative_risk_data["value"] = relative_risk_data.value_x.divide(
+            relative_risk_data.value_y
+        ).fillna(0)
+        return relative_risk_data.drop(columns=["value_x", "value_y"])
+
+    def get_exposure_callable(self, builder: Builder) -> Callable[[pd.Index], pd.Series]:
+        return builder.value.get_value(self.exposure_pipeline_name)
+
+    def adjust_target(self, index: pd.Index, target: pd.Series) -> pd.Series:
+        relative_risk = self.relative_risk(index)
+        return target * relative_risk
+
+    def get_relative_risk_source(self, builder: Builder) -> Callable[[pd.Index], pd.Series]:
+
+        if not self.is_exposure_categorical:
+            tmred = builder.data.load(f"{self.entity}.tmred")
+            tmrel = 0.5 * (tmred["min"] + tmred["max"])
+            scale = builder.data.load(f"{self.entity}.relative_risk_scalar")
+
+            def generate_relative_risk(index: pd.Index) -> pd.Series:
+                rr = self.lookup_tables["relative_risk"](index)
+                exposure = self.measure(index)
+                relative_risk = np.maximum(rr.values ** ((exposure - tmrel) / scale), 1)
+                return relative_risk
+
+        else:
+            index_columns = ["index", self.entity.name]
+
+            def generate_relative_risk(index: pd.Index) -> pd.Series:
+                rr = self.lookup_tables["relative_risk"](index)
+                exposure = self.measure(index).reset_index()
+                exposure.columns = index_columns
+                exposure = exposure.set_index(index_columns)
+
+                relative_risk = rr.stack().reset_index()
+                relative_risk.columns = index_columns + ["value"]
+                # Check if we need to remap cat1 and cat2 to exposed and unexposed categories
+                if (
+                    "cat1" in relative_risk[self.entity.name].unique()
+                    and self._exposure_distribution_type == "dichotomous"
+                ):
+                    warnings.warn(
+                        "Using 'cat1' and 'cat2' for dichotomous exposure is deprecated and will be removed in a future release. Use 'exposed' and 'unexposed' instead.",
+                        FutureWarning,
+                        stacklevel=2,
+                    )
+                    relative_risk[self.entity.name] = relative_risk[self.entity.name].replace(
+                        {
+                            "cat1": self.exposure_component.dichotomous_exposure_category_names.exposed,
+                            "cat2": self.exposure_component.dichotomous_exposure_category_names.unexposed,
+                        }
+                    )
+                relative_risk = relative_risk.set_index(index_columns)
+
+                effect = relative_risk.loc[exposure.index, "value"].droplevel(
+                    self.entity.name
+                )
+                return effect
+
+        return generate_relative_risk
+
+    def get_relative_risk_pipeline(self, builder: Builder) -> Pipeline:
+        return builder.value.register_value_producer(
+            f"{self.entity.name}_on_{self.target.name}.relative_risk",
+            self._relative_risk_source,
+            component=self,
+            required_resources=[self.measure],
+        )
+
+    def register_target_modifier(self, builder: Builder) -> None:
+        builder.value.register_value_modifier(
+            self.target_pipeline_name,
+            modifier=self.adjust_target,
+            component=self,
+            required_resources=[self.relative_risk],
+        )
+
+    def register_paf_modifier(self, builder: Builder) -> None:
+        required_columns = get_lookup_columns(
+            [self.lookup_tables["population_attributable_fraction"]]
+        )
+        builder.value.register_value_modifier(
+            self.target_paf_pipeline_name,
+            modifier=self.lookup_tables["population_attributable_fraction"],
+            component=self,
+            required_resources=required_columns,
+        )
+
+    ##################
+    # Helper methods #
+    ##################
+
+    def _get_exposure_class(self, builder: Builder) -> Exposure:
+        exposure_component = builder.components.get_component(self.entity)
+        if not isinstance(exposure_component, Exposure):
+            raise ValueError(
+                f"Exposure model {self.name} requires an Exposure component named {self.entity}"
+            )
+        return exposure_component

vivarium_public_health/exposure/exposure.py

@@ -0,0 +1,254 @@
+from abc import ABC, abstractmethod
+from typing import Any, NamedTuple
+
+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.randomness import RandomnessStream
+from vivarium.framework.resource import Resource
+from vivarium.framework.values import Pipeline
+
+from vivarium_public_health.exposure.distributions import (
+    ContinuousDistribution,
+    DichotomousDistribution,
+    EnsembleDistribution,
+    ExposureDistribution,
+    PolytomousDistribution,
+)
+from vivarium_public_health.risks.data_transformations import get_exposure_post_processor
+from vivarium_public_health.utilities import EntityString, get_lookup_columns
+
+
+class Exposure(Component, ABC):
+    """A base class to store common functionality for for risk-like health factors.
+
+    This class is used to define the determinant of models health factors such as
+    risks and the exposure to these risks, or interventions and the available coverage
+    for these interventions.
+
+    """
+
+    exposure_distributions = {
+        "dichotomous": DichotomousDistribution,
+        "ordered_polytomous": PolytomousDistribution,
+        "unordered_polytomous": PolytomousDistribution,
+        "normal": ContinuousDistribution,
+        "lognormal": ContinuousDistribution,
+        "ensemble": EnsembleDistribution,
+    }
+
+    ##############
+    # Properties #
+    ##############
+
+    @property
+    def name(self) -> str:
+        return self.entity
+
+    @property
+    def configuration_defaults(self) -> dict[str, Any]:
+        return {
+            self.name: {
+                "data_sources": {
+                    f"{self.exposure_type}": f"{self.entity}.{self.exposure_type}",
+                    "ensemble_distribution_weights": f"{self.entity}.exposure_distribution_weights",
+                    "exposure_standard_deviation": f"{self.entity}.exposure_standard_deviation",
+                },
+                "distribution_type": f"{self.entity}.distribution",
+                # rebinned_exposed only used for DichotomousDistribution
+                "rebinned_exposed": [],
+                "category_thresholds": [],
+            }
+        }
+
+    @property
+    def columns_created(self) -> list[str]:
+        columns_to_create = [self.propensity_column_name]
+        if self.create_exposure_column:
+            columns_to_create.append(self.exposure_column_name)
+        return columns_to_create
+
+    @property
+    def initialization_requirements(self) -> list[str | Resource]:
+        return [self.randomness]
+
+    @property
+    @abstractmethod
+    def exposure_type(self) -> str:
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def dichotomous_exposure_category_names(self) -> NamedTuple:
+        """The name of the exposure categories. E.g. "exposed" and "unexposed" or
+        "covered" and "uncovered".
+
+        """
+        raise NotImplementedError
+
+    #####################
+    # Lifecycle methods #
+    #####################
+
+    def __init__(self, entity: str) -> None:
+        """
+
+        Parameters
+        ----------
+        entity
+            the type and name of a entity, specified as "type.name". Type is singular.
+        level_type
+            The type of level for the health factor, e.g., "exposure" or "coverage".
+        """
+        super().__init__()
+        self.entity = EntityString(entity)
+        self.distribution_type = None
+        self.randomness_stream_name = f"initial_{self.entity.name}_propensity"
+        self.propensity_column_name = f"{self.entity.name}_propensity"
+        self.propensity_pipeline_name = f"{self.entity.name}.propensity"
+        self.exposure_pipeline_name = f"{self.entity.name}.{self.exposure_type}"
+        self.exposure_column_name = f"{self.entity.name}_{self.exposure_type}"
+
+    #################
+    # Setup methods #
+    #################
+
+    def build_all_lookup_tables(self, builder: "Builder") -> None:
+        # All lookup tables are built in the exposure distribution
+        pass
+
+    # noinspection PyAttributeOutsideInit
+    def setup(self, builder: Builder) -> None:
+        self.distribution_type = self.get_distribution_type(builder)
+        self.exposure_distribution = self.get_exposure_distribution(builder)
+
+        self.randomness = self.get_randomness_stream(builder)
+        self.propensity = self.get_propensity_pipeline(builder)
+        self.exposure = self.get_exposure_callable(builder)
+        # This will be overwritten in the Risk class if there is a non-loglinear risk effect
+        # on that risk instance
+        self.create_exposure_column = False
+
+    def get_distribution_type(self, builder: Builder) -> str:
+        """Get the distribution type for the risk from the configuration.
+
+        If the configured distribution type is not one of the supported types,
+        it is assumed to be a data source and the data is retrieved using the
+        get_data method.
+
+        Parameters
+        ----------
+        builder
+            The builder object.
+
+        Returns
+        -------
+            The distribution type.
+        """
+        if self.configuration is None:
+            self.configuration = self.get_configuration(builder)
+
+        distribution_type = self.configuration["distribution_type"]
+        if distribution_type not in self.exposure_distributions.keys():
+            # todo deal with incorrect typing
+            distribution_type = self.get_data(builder, distribution_type)
+
+        if self.configuration["rebinned_exposed"]:
+            if distribution_type != "dichotomous" or "polytomous" not in distribution_type:
+                raise ValueError(
+                    f"Unsupported risk distribution type '{distribution_type}' "
+                    f"for {self.name}. Rebinned exposed categories are only "
+                    "supported for dichotomous and polytomous distributions."
+                )
+            distribution_type = "dichotomous"
+        return distribution_type
+
+    def get_exposure_distribution(self, builder: Builder) -> ExposureDistribution:
+        """Creates and sets up the exposure distribution component for the Risk
+        based on its distribution type.
+
+        Parameters
+        ----------
+        builder
+            The builder object.
+
+        Returns
+        -------
+            The exposure distribution.
+
+        Raises
+        ------
+        NotImplementedError
+            If the distribution type is not supported.
+        """
+        try:
+            exposure_distribution = self.exposure_distributions[self.distribution_type](
+                self, self.distribution_type
+            )
+        except KeyError:
+            raise NotImplementedError(
+                f"Distribution type {self.distribution_type} is not supported."
+            )
+
+        exposure_distribution.setup_component(builder)
+        return exposure_distribution
+
+    def get_randomness_stream(self, builder: Builder) -> RandomnessStream:
+        return builder.randomness.get_stream(self.randomness_stream_name, component=self)
+
+    def get_propensity_pipeline(self, builder: Builder) -> Pipeline:
+        return builder.value.register_value_producer(
+            self.propensity_pipeline_name,
+            source=lambda index: (
+                self.population_view.subview([self.propensity_column_name])
+                .get(index)
+                .squeeze(axis=1)
+            ),
+            component=self,
+            required_resources=[self.propensity_column_name],
+        )
+
+    def get_exposure_callable(self, builder: Builder) -> Pipeline:
+        required_columns = get_lookup_columns(
+            self.exposure_distribution.lookup_tables.values()
+        )
+        return builder.value.register_value_producer(
+            self.exposure_pipeline_name,
+            source=self.get_current_exposure,
+            component=self,
+            required_resources=required_columns
+            + [
+                self.propensity,
+                self.exposure_distribution.exposure_parameters,
+            ],
+            preferred_post_processor=get_exposure_post_processor(builder, self.name),
+        )
+
+    ########################
+    # Event-driven methods #
+    ########################
+
+    def on_initialize_simulants(self, pop_data: SimulantData) -> None:
+        propensity = pd.Series(
+            self.randomness.get_draw(pop_data.index), name=self.propensity_column_name
+        )
+        self.population_view.update(propensity)
+        self.update_exposure_column(pop_data.index)
+
+    def on_time_step_prepare(self, event: Event) -> None:
+        self.update_exposure_column(event.index)
+
+    def update_exposure_column(self, index: pd.Index) -> None:
+        if self.create_exposure_column:
+            exposure = pd.Series(self.exposure_type(index), name=self.exposure_column_name)
+            self.population_view.update(exposure)
+
+    ##################################
+    # Pipeline sources and modifiers #
+    ##################################
+
+    def get_current_exposure(self, index: pd.Index) -> pd.Series:
+        propensity = self.propensity(index)
+        return pd.Series(self.exposure_distribution.ppf(propensity), index=index)

vivarium_public_health/population/data_transformations.py

@@ -584,8 +584,8 @@ def rescale_final_age_bin(builder, population_data):
 
 def validate_crude_birth_rate_data(builder, data_year_max):
     population_config = builder.configuration.population.to_dict()
-    untracking_age = population_config.get("untracking_age", None)
-    age_end = population_config.get("age_end", None)
+    untracking_age = population_config.get("untracking_age")
+    age_end = population_config.get("age_end")
     if untracking_age and age_end and age_end != untracking_age:
         raise ValueError(
             "If you specify an exit age, the initial population age end must be the same "

vivarium_public_health/results/__init__.py

@@ -1,6 +1,7 @@
 from .columns import COLUMNS
 from .disability import DisabilityObserver
 from .disease import DiseaseObserver
+from .intervention import CategoricalInterventionObserver
 from .mortality import MortalityObserver
 from .observer import PublicHealthObserver
 from .risk import CategoricalRiskObserver