vivarium-public-health 4.3.2__py3-none-any.whl → 4.3.4__py3-none-any.whl

vivarium_public_health/exposure/effect.py

@@ -1,356 +0,0 @@
-"""
-###################
-# 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

@@ -1,254 +0,0 @@
-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)