vivarium-public-health 4.3.1__py3-none-any.whl → 4.3.3__py3-none-any.whl

vivarium_public_health/__init__.py

@@ -46,13 +46,7 @@ from vivarium_public_health.risks import (
     Risk,
     RiskEffect,
 )
-from vivarium_public_health.treatment import (
-    AbsoluteShift,
-    Intervention,
-    InterventionEffect,
-    LinearScaleUp,
-    TherapeuticInertia,
-)
+from vivarium_public_health.treatment import AbsoluteShift, LinearScaleUp, TherapeuticInertia
 
 __all__ = [
     __author__,

vivarium_public_health/_version.py

@@ -1 +1 @@
-__version__ = "4.3.1"
+__version__ = "4.3.3"

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")
-    age_end = population_config.get("age_end")
+    untracking_age = population_config.get("untracking_age", None)
+    age_end = population_config.get("age_end", None)
     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,7 +1,6 @@
 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

vivarium_public_health/risks/base_risk.py

@@ -7,15 +7,30 @@ This module contains tools for modeling categorical and continuous risk
 exposure.
 
 """
-import warnings
-from typing import NamedTuple
 
+from typing import Any
+
+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 import Exposure
+from vivarium_public_health.risks.data_transformations import get_exposure_post_processor
+from vivarium_public_health.risks.distributions import (
+    ContinuousDistribution,
+    DichotomousDistribution,
+    EnsembleDistribution,
+    PolytomousDistribution,
+    RiskExposureDistribution,
+)
+from vivarium_public_health.utilities import EntityString, get_lookup_columns
 
 
-class Risk(Exposure):
+class Risk(Component):
     """A model for a risk factor defined by either a continuous or a categorical value.
 
     For example,
@@ -74,52 +89,217 @@ class Risk(Exposure):
 
     """
 
-    @property
-    def risk(self) -> str:
-        warnings.warn(
-            "The 'risk' attribute is deprecated. Use 'entity' instead.",
-            DeprecationWarning,
-            stacklevel=2,
-        )
-        return self.entity
-
-    @risk.setter
-    def risk(self, value: str) -> None:
-        warnings.warn(
-            "The 'risk' attribute is deprecated. Use 'entity' instead.",
-            DeprecationWarning,
-            stacklevel=2,
-        )
-        self.entity = value
+    exposure_distributions = {
+        "dichotomous": DichotomousDistribution,
+        "ordered_polytomous": PolytomousDistribution,
+        "unordered_polytomous": PolytomousDistribution,
+        "normal": ContinuousDistribution,
+        "lognormal": ContinuousDistribution,
+        "ensemble": EnsembleDistribution,
+    }
+
+    ##############
+    # Properties #
+    ##############
 
     @property
-    def exposure_type(self) -> str:
-        """The measure of the risk exposure."""
-        return "exposure"
+    def name(self) -> str:
+        return self.risk
 
     @property
-    def dichotomous_exposure_category_names(self) -> NamedTuple:
-        """The name of the exposed category for this intervention."""
+    def configuration_defaults(self) -> dict[str, Any]:
+        return {
+            self.name: {
+                "data_sources": {
+                    "exposure": f"{self.risk}.exposure",
+                    "ensemble_distribution_weights": f"{self.risk}.exposure_distribution_weights",
+                    "exposure_standard_deviation": f"{self.risk}.exposure_standard_deviation",
+                },
+                "distribution_type": f"{self.risk}.distribution",
+                # rebinned_exposed only used for DichotomousDistribution
+                "rebinned_exposed": [],
+                "category_thresholds": [],
+            }
+        }
 
-        class __Categories(NamedTuple):
-            exposed: str = "exposed"
-            unexposed: str = "unexposed"
+    @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
 
-        categories = __Categories()
-        return categories
+    @property
+    def initialization_requirements(self) -> list[str | Resource]:
+        return [self.randomness]
 
     #####################
     # Lifecycle methods #
     #####################
 
+    def __init__(self, risk: str):
+        """
+
+        Parameters
+        ----------
+        risk
+            the type and name of a risk, specified as "type.name". Type is singular.
+        """
+        super().__init__()
+        self.risk = EntityString(risk)
+        self.distribution_type = None
+
+        self.randomness_stream_name = f"initial_{self.risk.name}_propensity"
+        self.propensity_column_name = f"{self.risk.name}_propensity"
+        self.propensity_pipeline_name = f"{self.risk.name}.propensity"
+        self.exposure_pipeline_name = f"{self.risk.name}.exposure"
+        self.exposure_column_name = f"{self.risk.name}_exposure"
+
+    #################
+    # 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:
-        super().setup(builder)
-        # We want to set this to True if there is a non-loglinear risk effect
+        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_pipeline(builder)
+
+        # We want to set this to True iff there is a non-loglinear risk effect
         # on this risk instance
         self.create_exposure_column = bool(
             [
                 component
                 for component in builder.components.list_components()
-                if component.startswith(f"non_log_linear_risk_effect.{self.entity.name}_on_")
+                if component.startswith(f"non_log_linear_risk_effect.{self.risk.name}_on_")
             ]
         )
+
+    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) -> RiskExposureDistribution:
+        """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.risk, 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_pipeline(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(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/risks/data_transformations.py

@@ -1,32 +1,254 @@
 """
-==============================
-Data Transformations (Legacy)
-==============================
+=========================
+Risk Data Transformations
+=========================
 
-.. deprecated:: 4.3.0
-   This module is deprecated. Use :mod:`vivarium_public_health.exposure.data_transformations` instead.
+This module contains tools for handling raw risk exposure and relative
+risk data and performing any necessary data transformations.
 
-Backward compatibility module for risk data_transformations.
+"""
 
-This module provides backward compatibility for imports that expect risk
-distribution classes to be in vivarium_public_health.risks.data_transformations.
+import numpy as np
+import pandas as pd
+from vivarium.framework.engine import Builder
 
-The actual distribution classes have been moved to:
-vivarium_public_health.exposure.data_transformations
+from vivarium_public_health.utilities import EntityString, TargetString
 
-This module will be deprecated in a future version.
-"""
+#############
+# Utilities #
+#############
+
+
+def pivot_categorical(
+    builder: Builder,
+    risk: EntityString,
+    data: pd.DataFrame,
+    pivot_column: str = "parameter",
+    reset_index: bool = True,
+) -> pd.DataFrame:
+    """Pivots data that is long on categories to be wide."""
+    # todo remove dependency on artifact manager having exactly one value column
+    value_column = builder.data.value_columns()(f"{risk}.exposure")[0]
+    index_cols = [
+        column for column in data.columns if column not in [value_column, pivot_column]
+    ]
+    data = data.pivot_table(index=index_cols, columns=pivot_column, values=value_column)
+    if reset_index:
+        data = data.reset_index()
+    data.columns.name = None
+
+    return data
+
+
+##########################
+# Exposure data handlers #
+##########################
+
+
+def get_exposure_post_processor(builder, risk: str):
+    thresholds = builder.configuration[risk]["category_thresholds"]
+
+    if thresholds:
+        thresholds = [-np.inf] + thresholds + [np.inf]
+        categories = [f"cat{i}" for i in range(1, len(thresholds))]
+
+        def post_processor(exposure, _):
+            return pd.Series(
+                pd.cut(exposure, thresholds, labels=categories), index=exposure.index
+            ).astype(str)
+
+    else:
+        post_processor = None
+
+    return post_processor
+
+
+def load_exposure_data(builder: Builder, risk: EntityString) -> pd.DataFrame:
+    risk_component = builder.components.get_component(risk)
+    return risk_component.get_data(
+        builder, builder.configuration[risk_component.name]["data_sources"]["exposure"]
+    )
+
+
+###############################
+# Relative risk data handlers #
+###############################
+
+
+def rebin_relative_risk_data(
+    builder, risk: EntityString, relative_risk_data: pd.DataFrame
+) -> pd.DataFrame:
+    """Rebin relative risk data if necessary.
+
+    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 risk in builder.configuration.to_dict():
+        return relative_risk_data
+
+    rebin_exposed_categories = set(builder.configuration[risk]["rebinned_exposed"])
+
+    if rebin_exposed_categories:
+        # todo make sure this works
+        exposure_data = load_exposure_data(builder, risk)
+        relative_risk_data = _rebin_relative_risk_data(
+            relative_risk_data, exposure_data, rebin_exposed_categories
+        )
+
+    return relative_risk_data
+
+
+def _rebin_relative_risk_data(
+    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"])
+
+
+##############
+# Validators #
+##############
+
+
+def validate_distribution_data_source(builder: Builder, risk: EntityString) -> None:
+    """Checks that the exposure distribution specification is valid."""
+    exposure_type = builder.configuration[risk]["data_sources"]["exposure"]
+    rebin = builder.configuration[risk]["rebinned_exposed"]
+    category_thresholds = builder.configuration[risk]["category_thresholds"]
+
+    if risk.type == "alternative_risk_factor":
+        if exposure_type != "data" or rebin:
+            raise ValueError(
+                "Parameterized risk components are not available for alternative risks."
+            )
+
+        if not category_thresholds:
+            raise ValueError("Must specify category thresholds to use alternative risks.")
+
+    elif risk.type not in ["risk_factor", "coverage_gap"]:
+        raise ValueError(f"Unknown risk type {risk.type} for risk {risk.name}")
+
+
+def validate_relative_risk_data_source(builder, risk: EntityString, target: TargetString):
+    from vivarium_public_health.risks import RiskEffect
+
+    source_key = RiskEffect.get_name(risk, target)
+    source_config = builder.configuration[source_key]
+
+    provided_keys = set(
+        k
+        for k, v in source_config["distribution_args"].to_dict().items()
+        if isinstance(v, (int, float))
+    )
+
+    source_map = {
+        "data": set(),
+        "relative risk value": {"relative_risk"},
+        "normal distribution": {"mean", "se"},
+        "log distribution": {"log_mean", "log_se", "tau_squared"},
+    }
+
+    if provided_keys not in source_map.values():
+        raise ValueError(
+            f"The acceptable parameter options for specifying relative risk are: "
+            f"{source_map.values()}. You provided {provided_keys} for {source_key}."
+        )
+
+    source_type = [k for k, v in source_map.items() if provided_keys == v][0]
+
+    if source_type == "relative risk value":
+        if not 1 <= source_type <= 100:
+            raise ValueError(
+                "If specifying a single value for relative risk, it should be in the range [1, 100]. "
+                f"You provided {source_type} for {source_key}."
+            )
+    elif source_type == "normal distribution":
+        if source_config["mean"] <= 0 or source_config["se"] <= 0:
+            raise ValueError(
+                f"To specify parameters for a normal distribution for a risk effect, you must provide"
+                f"both mean and se above 0. This is not the case for {source_key}."
+            )
+    elif source_type == "log distribution":
+        if source_config["log_mean"] <= 0 or source_config["log_se"] <= 0:
+            raise ValueError(
+                f"To specify parameters for a log distribution for a risk effect, you must provide"
+                f"both log_mean and log_se above 0. This is not the case for {source_key}."
+            )
+        if source_config["tau_squared"] < 0:
+            raise ValueError(
+                f"To specify parameters for a log distribution for a risk effect, you must provide"
+                f"tau_squared >= 0. This is not the case for {source_key}."
+            )
+    else:
+        pass
+
+    return source_type
+
+
+def validate_relative_risk_rebin_source(
+    builder, risk: EntityString, target: TargetString, data: pd.DataFrame
+):
+    if data.index.size == 0:
+        raise ValueError(
+            f"Subsetting {risk} relative risk data to {target.name} {target.measure} "
+            "returned an empty DataFrame. Check your artifact."
+        )
+    if risk in builder.configuration.to_dict():
+        validate_rebin_source(builder, risk, data)
+
+
+def validate_rebin_source(builder, risk: EntityString, data: pd.DataFrame) -> None:
+
+    if not isinstance(data, pd.DataFrame):
+        return
+
+    rebin_exposed_categories = set(builder.configuration[risk]["rebinned_exposed"])
+
+    if rebin_exposed_categories and builder.configuration[risk]["category_thresholds"]:
+        raise ValueError(
+            f"Rebinning and category thresholds are mutually exclusive. "
+            f"You provided both for {risk.name}."
+        )
 
-import warnings
+    if rebin_exposed_categories and "polytomous" not in builder.data.load(
+        f"{risk}.distribution"
+    ):
+        raise ValueError(
+            f"Rebinning is only supported for polytomous risks. You provided "
+            f"rebinning exposed categoriesfor {risk.name}, which is of "
+            f"type {builder.data.load(f'{risk}.distribution')}."
+        )
 
-# Issue a deprecation warning when this module is imported
-warnings.warn(
-    "Importing from 'vivarium_public_health.risks.data_transformations' is deprecated. "
-    "Please import from 'vivarium_public_health.exposure.data_transformations' instead.",
-    DeprecationWarning,
-    stacklevel=2,
-)
+    invalid_cats = rebin_exposed_categories.difference(set(data.parameter))
+    if invalid_cats:
+        raise ValueError(
+            f"The following provided categories for the rebinned exposed "
+            f"category of {risk.name} are not found in the exposure data: "
+            f"{invalid_cats}."
+        )
 
-# Import all the classes from the new location
-from vivarium_public_health.exposure.data_transformations import *
-from vivarium_public_health.exposure.data_transformations import _rebin_relative_risk_data
+    if rebin_exposed_categories == set(data.parameter):
+        raise ValueError(
+            f"The provided categories for the rebinned exposed category of "
+            f"{risk.name} comprise all categories for the exposure data. "
+            f"At least one category must be left out of the provided categories "
+            f"to be rebinned into the unexposed category."
+        )