vivarium-public-health 2.3.3__py3-none-any.whl → 3.0.1__py3-none-any.whl

vivarium_public_health/risks/distributions.py

@@ -7,52 +7,86 @@ This module contains tools for modeling several different risk
 exposure distributions.
 
 """
-from typing import Dict, List
+
+from abc import ABC, abstractmethod
+from typing import Callable, Dict, List, Optional, Union
 
 import numpy as np
 import pandas as pd
-from risk_distributions import EnsembleDistribution, LogNormal, Normal
+import risk_distributions as rd
+from layered_config_tree import LayeredConfigTree
 from vivarium import Component
 from vivarium.framework.engine import Builder
 from vivarium.framework.population import SimulantData
 from vivarium.framework.values import Pipeline, list_combiner, union_post_processor
 
-from vivarium_public_health.risks.data_transformations import get_distribution_data
-from vivarium_public_health.utilities import EntityString
+from vivarium_public_health.risks.data_transformations import pivot_categorical
+from vivarium_public_health.utilities import EntityString, get_lookup_columns
 
 
 class MissingDataError(Exception):
     pass
 
 
-# FIXME: This is a hack.  It's wrapping up an adaptor pattern in another
-# adaptor pattern, which is gross, but would require some more difficult
-# refactoring which is thoroughly out of scope right now. -J.C. 8/25/19
-class SimulationDistribution(Component):
-    """Wrapper around a variety of distribution implementations."""
+class RiskExposureDistribution(Component, ABC):
 
     #####################
     # Lifecycle methods #
     #####################
 
-    def __init__(self, risk: str):
+    def __init__(
+        self,
+        risk: EntityString,
+        distribution_type: str,
+        exposure_data: Optional[Union[int, float, pd.DataFrame]] = None,
+    ) -> None:
         super().__init__()
-        self.risk = EntityString(risk)
+        self.risk = risk
+        self.distribution_type = distribution_type
+        self._exposure_data = exposure_data
+
+        self.parameters_pipeline_name = f"{self.risk}.exposure_parameters"
 
+    #################
+    # Setup methods #
+    #################
+
+    def get_configuration(self, builder: "Builder") -> Optional[LayeredConfigTree]:
+        return builder.configuration[self.risk]
+
+    @abstractmethod
+    def build_all_lookup_tables(self, builder: "Builder") -> None:
+        raise NotImplementedError
+
+    def get_exposure_data(self, builder: Builder) -> Union[int, float, pd.DataFrame]:
+        if self._exposure_data is not None:
+            return self._exposure_data
+        return self.get_data(builder, self.configuration["data_sources"]["exposure"])
+
+    # noinspection PyAttributeOutsideInit
     def setup(self, builder: Builder) -> None:
-        distribution_data = get_distribution_data(builder, self.risk)
-        self.implementation = get_distribution(self.risk, **distribution_data)
-        self.implementation.setup_component(builder)
+        self.exposure_parameters = self.get_exposure_parameter_pipeline(builder)
+        if self.exposure_parameters.name != self.parameters_pipeline_name:
+            raise ValueError(
+                "Expected exposure parameters pipeline to be named "
+                f"{self.parameters_pipeline_name}, "
+                f"but found {self.exposure_parameters.name}."
+            )
+
+    @abstractmethod
+    def get_exposure_parameter_pipeline(self, builder: Builder) -> Pipeline:
+        raise NotImplementedError
 
     ##################
     # Public methods #
     ##################
 
-    def ppf(self, q):
-        return self.implementation.ppf(q)
+    @abstractmethod
+    def ppf(self, quantiles: pd.Series) -> pd.Series:
+        raise NotImplementedError
 
 
-class EnsembleSimulation(Component):
+class EnsembleDistribution(RiskExposureDistribution):
     ##############
     # Properties #
     ##############
@@ -73,38 +107,71 @@ class EnsembleSimulation(Component):
     # Lifecycle methods #
     #####################
 
-    def __init__(self, risk, weights, mean, sd):
-        super().__init__()
-        self.risk = EntityString(risk)
-        self._weights, self._parameters = self.get_parameters(weights, mean, sd)
+    def __init__(self, risk: EntityString, distribution_type: str = "ensemble") -> None:
+        super().__init__(risk, distribution_type)
         self._propensity = f"ensemble_propensity_{self.risk}"
 
-    def setup(self, builder: Builder) -> None:
-        self.weights = builder.lookup.build_table(
-            self._weights, key_columns=["sex"], parameter_columns=["age", "year"]
+    #################
+    # Setup methods #
+    #################
+
+    def build_all_lookup_tables(self, builder: Builder) -> None:
+        exposure_data = self.get_exposure_data(builder)
+        standard_deviation = self.get_data(
+            builder,
+            self.configuration["data_sources"]["exposure_standard_deviation"],
+        )
+        weights_source = self.configuration["data_sources"]["ensemble_distribution_weights"]
+        raw_weights = self.get_data(builder, weights_source)
+
+        glnorm_mask = raw_weights["parameter"] == "glnorm"
+        if np.any(raw_weights.loc[glnorm_mask, self.get_value_columns(weights_source)]):
+            raise NotImplementedError("glnorm distribution is not supported")
+        raw_weights = raw_weights[~glnorm_mask]
+
+        distributions = list(raw_weights["parameter"].unique())
+
+        raw_weights = pivot_categorical(
+            builder, self.risk, raw_weights, pivot_column="parameter", reset_index=False
         )
+
+        weights, parameters = rd.EnsembleDistribution.get_parameters(
+            raw_weights,
+            mean=get_risk_distribution_parameter(self.get_value_columns, exposure_data),
+            sd=get_risk_distribution_parameter(self.get_value_columns, standard_deviation),
+        )
+
+        distribution_weights_table = self.build_lookup_table(
+            builder, weights.reset_index(), distributions
+        )
+        self.lookup_tables["ensemble_distribution_weights"] = distribution_weights_table
+        key_columns = distribution_weights_table.key_columns
+        parameter_columns = distribution_weights_table.parameter_columns
+
         self.parameters = {
-            k: builder.lookup.build_table(
-                v, key_columns=["sex"], parameter_columns=["age", "year"]
+            parameter: builder.lookup.build_table(
+                data.reset_index(),
+                key_columns=key_columns,
+                parameter_columns=parameter_columns,
             )
-            for k, v in self._parameters.items()
+            for parameter, data in parameters.items()
         }
 
+    def setup(self, builder: Builder) -> None:
+        super().setup(builder)
         self.randomness = builder.randomness.get_stream(self._propensity)
 
-    ##########################
-    # Initialization methods #
-    ##########################
-
-    def get_parameters(self, weights, mean, sd):
-        index_cols = ["sex", "age_start", "age_end", "year_start", "year_end"]
-        weights = weights.set_index(index_cols)
-        mean = mean.set_index(index_cols)["value"]
-        sd = sd.set_index(index_cols)["value"]
-        weights, parameters = EnsembleDistribution.get_parameters(weights, mean=mean, sd=sd)
-        return weights.reset_index(), {
-            name: p.reset_index() for name, p in parameters.items()
-        }
+    def get_exposure_parameter_pipeline(self, builder: Builder) -> Pipeline:
+        # This pipeline is not needed for ensemble distributions, so just
+        # register a dummy pipeline
+        def raise_not_implemented():
+            raise NotImplementedError(
+                "EnsembleDistribution does not use exposure parameters."
+            )
+
+        return builder.value.register_value_producer(
+            self.parameters_pipeline_name, lambda *_: raise_not_implemented()
+        )
 
     ########################
     # Event-driven methods #
@@ -120,149 +187,256 @@ class EnsembleSimulation(Component):
     # Public methods #
     ##################
 
-    def ppf(self, q):
-        if not q.empty:
-            q = clip(q)
-            weights = self.weights(q.index)
+    def ppf(self, quantiles: pd.Series) -> pd.Series:
+        if not quantiles.empty:
+            quantiles = clip(quantiles)
+            weights = self.lookup_tables["ensemble_distribution_weights"](quantiles.index)
             parameters = {
-                name: parameter(q.index) for name, parameter in self.parameters.items()
+                name: param(quantiles.index) for name, param in self.parameters.items()
             }
-            ensemble_propensity = self.population_view.get(q.index).iloc[:, 0]
-            x = EnsembleDistribution(weights, parameters).ppf(q, ensemble_propensity)
+            ensemble_propensity = self.population_view.get(quantiles.index).iloc[:, 0]
+            x = rd.EnsembleDistribution(weights, parameters).ppf(
+                quantiles, ensemble_propensity
+            )
             x[x.isnull()] = 0
         else:
             x = pd.Series([])
         return x
 
 
-class ContinuousDistribution(Component):
+class ContinuousDistribution(RiskExposureDistribution):
     #####################
     # Lifecycle methods #
     #####################
 
-    def __init__(self, risk, mean, sd, distribution=None):
-        super().__init__()
-        self.risk = EntityString(risk)
-        self._distribution = distribution
-        self._parameters = self.get_parameters(mean, sd)
+    def __init__(self, risk: EntityString, distribution_type: str) -> None:
+        super().__init__(risk, distribution_type)
+        self.standard_deviation = None
+        try:
+            self._distribution = {
+                "normal": rd.Normal,
+                "lognormal": rd.LogNormal,
+            }[distribution_type]
+        except KeyError:
+            raise NotImplementedError(
+                f"Distribution type {distribution_type} is not supported for "
+                f"risk {risk.name}."
+            )
 
-    def setup(self, builder: Builder) -> None:
-        self.parameters = builder.lookup.build_table(
-            self._parameters, key_columns=["sex"], parameter_columns=["age", "year"]
+    #################
+    # Setup methods #
+    #################
+
+    def build_all_lookup_tables(self, builder: "Builder") -> None:
+        exposure_data = self.get_exposure_data(builder)
+        standard_deviation = self.get_data(
+            builder, self.configuration["data_sources"]["exposure_standard_deviation"]
+        )
+        parameters = self._distribution.get_parameters(
+            mean=get_risk_distribution_parameter(self.get_value_columns, exposure_data),
+            sd=get_risk_distribution_parameter(self.get_value_columns, standard_deviation),
         )
 
-    ##########################
-    # Initialization methods #
-    ##########################
+        self.lookup_tables["parameters"] = self.build_lookup_table(
+            builder, parameters.reset_index(), list(parameters.columns)
+        )
 
-    def get_parameters(self, mean, sd):
-        index = ["sex", "age_start", "age_end", "year_start", "year_end"]
-        mean = mean.set_index(index)["value"]
-        sd = sd.set_index(index)["value"]
-        return self._distribution.get_parameters(mean=mean, sd=sd).reset_index()
+    def get_exposure_parameter_pipeline(self, builder: Builder) -> Pipeline:
+        return builder.value.register_value_producer(
+            self.parameters_pipeline_name,
+            source=self.lookup_tables["parameters"],
+            requires_columns=get_lookup_columns([self.lookup_tables["parameters"]]),
+        )
 
     ##################
     # Public methods #
     ##################
 
-    def ppf(self, q):
-        if not q.empty:
-            q = clip(q)
-            x = self._distribution(parameters=self.parameters(q.index)).ppf(q)
+    def ppf(self, quantiles: pd.Series) -> pd.Series:
+        if not quantiles.empty:
+            quantiles = clip(quantiles)
+            parameters = self.exposure_parameters(quantiles.index)
+            x = self._distribution(parameters=parameters).ppf(quantiles)
             x[x.isnull()] = 0
         else:
             x = pd.Series([])
         return x
 
 
-class PolytomousDistribution(Component):
-    #####################
-    # Lifecycle methods #
-    #####################
-
-    def __init__(self, risk: str, exposure_data: pd.DataFrame):
-        super().__init__()
-        self.risk = EntityString(risk)
-        self._exposure_data = exposure_data
-        self.exposure_parameters_pipeline_name = f"{self.risk}.exposure_parameters"
-
-    # noinspection PyAttributeOutsideInit
-    def setup(self, builder: Builder) -> None:
-        self.categories = self.get_categories()
-        self.exposure = self.get_exposure_parameters(builder)
+class PolytomousDistribution(RiskExposureDistribution):
+    @property
+    def categories(self) -> List[str]:
+        # These need to be sorted so the cumulative sum is in the ocrrect order of categories
+        # and results are therefore reproducible and correct
+        return sorted(self.lookup_tables["exposure"].value_columns)
 
     #################
     # Setup methods #
     #################
 
-    def get_categories(self) -> List[str]:
-        return sorted(
-            [column for column in self._exposure_data if "cat" in column],
-            key=lambda column: int(column[3:]),
+    def build_all_lookup_tables(self, builder: "Builder") -> None:
+        exposure_data = self.get_exposure_data(builder)
+        exposure_value_columns = self.get_exposure_value_columns(exposure_data)
+
+        if isinstance(exposure_data, pd.DataFrame):
+            exposure_data = pivot_categorical(builder, self.risk, exposure_data, "parameter")
+
+        self.lookup_tables["exposure"] = self.build_lookup_table(
+            builder, exposure_data, exposure_value_columns
         )
 
-    def get_exposure_parameters(self, builder: Builder) -> Pipeline:
+    def get_exposure_value_columns(
+        self, exposure_data: Union[int, float, pd.DataFrame]
+    ) -> Optional[List[str]]:
+        if isinstance(exposure_data, pd.DataFrame):
+            return list(exposure_data["parameter"].unique())
+        return None
+
+    def get_exposure_parameter_pipeline(self, builder: Builder) -> Pipeline:
         return builder.value.register_value_producer(
-            self.exposure_parameters_pipeline_name,
-            source=builder.lookup.build_table(
-                self._exposure_data,
-                key_columns=["sex"],
-                parameter_columns=["age", "year"],
-            ),
+            self.parameters_pipeline_name,
+            source=self.lookup_tables["exposure"],
+            requires_columns=get_lookup_columns([self.lookup_tables["exposure"]]),
         )
 
     ##################
     # Public methods #
     ##################
 
-    def ppf(self, x: pd.Series) -> pd.Series:
-        exposure = self.exposure(x.index)
+    def ppf(self, quantiles: pd.Series) -> pd.Series:
+        exposure = self.exposure_parameters(quantiles.index)
         sorted_exposures = exposure[self.categories]
         if not np.allclose(1, np.sum(sorted_exposures, axis=1)):
             raise MissingDataError("All exposure data returned as 0.")
         exposure_sum = sorted_exposures.cumsum(axis="columns")
         category_index = pd.concat(
-            [exposure_sum[c] < x for c in exposure_sum.columns], axis=1
+            [exposure_sum[c] < quantiles for c in exposure_sum.columns], axis=1
         ).sum(axis=1)
         return pd.Series(
             np.array(self.categories)[category_index],
             name=self.risk + ".exposure",
-            index=x.index,
+            index=quantiles.index,
         )
 
 
-class DichotomousDistribution(Component):
-    #####################
-    # Lifecycle methods #
-    #####################
+class DichotomousDistribution(RiskExposureDistribution):
 
-    def __init__(self, risk: str, exposure_data: pd.DataFrame):
-        super().__init__()
-        self.risk = risk
-        self._exposure_data = exposure_data.drop(columns="cat2")
+    #################
+    # Setup methods #
+    #################
 
-    # noinspection PyAttributeOutsideInit
-    def setup(self, builder: Builder) -> None:
-        self._base_exposure = builder.lookup.build_table(
-            self._exposure_data, key_columns=["sex"], parameter_columns=["age", "year"]
+    def build_all_lookup_tables(self, builder: "Builder") -> None:
+        exposure_data = self.get_exposure_data(builder)
+        exposure_value_columns = self.get_exposure_value_columns(exposure_data)
+
+        if isinstance(exposure_data, pd.DataFrame):
+            any_negatives = (exposure_data[exposure_value_columns] < 0).any().any()
+            any_over_one = (exposure_data[exposure_value_columns] > 1).any().any()
+            if any_negatives or any_over_one:
+                raise ValueError(f"All exposures must be in the range [0, 1] for {self.risk}")
+        elif exposure_data < 0 or exposure_data > 1:
+            raise ValueError(f"Exposure must be in the range [0, 1] for {self.risk}")
+
+        self.lookup_tables["exposure"] = self.build_lookup_table(
+            builder, exposure_data, exposure_value_columns
         )
-        self.exposure_proportion = builder.value.register_value_producer(
-            f"{self.risk}.exposure_parameters", source=self.exposure
+        self.lookup_tables["paf"] = self.build_lookup_table(builder, 0.0)
+
+    def get_exposure_data(self, builder: Builder) -> Union[int, float, pd.DataFrame]:
+        exposure_data = super().get_exposure_data(builder)
+
+        if isinstance(exposure_data, (int, float)):
+            return exposure_data
+
+        # rebin exposure categories
+        self.validate_rebin_source(builder, exposure_data)
+        rebin_exposed_categories = set(self.configuration["rebinned_exposed"])
+        if rebin_exposed_categories:
+            exposure_data = self._rebin_exposure_data(exposure_data, rebin_exposed_categories)
+
+        exposure_data = exposure_data[exposure_data["parameter"] == "cat1"]
+        return exposure_data.drop(columns="parameter")
+
+    @staticmethod
+    def _rebin_exposure_data(
+        exposure_data: pd.DataFrame, rebin_exposed_categories: set
+    ) -> pd.DataFrame:
+        exposure_data = exposure_data[
+            exposure_data["parameter"].isin(rebin_exposed_categories)
+        ]
+        exposure_data["parameter"] = "cat1"
+        exposure_data = (
+            exposure_data.groupby(list(exposure_data.columns.difference(["value"])))
+            .sum()
+            .reset_index()
         )
-        base_paf = builder.lookup.build_table(0)
+        return exposure_data
+
+    def get_exposure_value_columns(
+        self, exposure_data: Union[int, float, pd.DataFrame]
+    ) -> Optional[List[str]]:
+        if isinstance(exposure_data, pd.DataFrame):
+            return self.get_value_columns(exposure_data)
+        return None
+
+    # noinspection PyAttributeOutsideInit
+    def setup(self, builder: Builder) -> None:
+        super().setup(builder)
         self.joint_paf = builder.value.register_value_producer(
             f"{self.risk}.exposure_parameters.paf",
-            source=lambda index: [base_paf(index)],
+            source=lambda index: [self.lookup_tables["paf"](index)],
             preferred_combiner=list_combiner,
             preferred_post_processor=union_post_processor,
         )
 
+    def get_exposure_parameter_pipeline(self, builder: Builder) -> Pipeline:
+        return builder.value.register_value_producer(
+            f"{self.risk}.exposure_parameters",
+            source=self.exposure_parameter_source,
+            requires_columns=get_lookup_columns([self.lookup_tables["exposure"]]),
+        )
+
+    ##############
+    # Validators #
+    ##############
+
+    def validate_rebin_source(self, builder, data: pd.DataFrame) -> None:
+        if not isinstance(data, pd.DataFrame):
+            return
+
+        rebin_exposed_categories = set(builder.configuration[self.risk]["rebinned_exposed"])
+
+        if (
+            rebin_exposed_categories
+            and builder.configuration[self.risk]["category_thresholds"]
+        ):
+            raise ValueError(
+                f"Rebinning and category thresholds are mutually exclusive. "
+                f"You provided both for {self.risk.name}."
+            )
+
+        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 {self.risk.name} are not found in the exposure data: "
+                f"{invalid_cats}."
+            )
+
+        if rebin_exposed_categories == set(data.parameter):
+            raise ValueError(
+                f"The provided categories for the rebinned exposed category of "
+                f"{self.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."
+            )
+
     ##################################
     # Pipeline sources and modifiers #
     ##################################
 
-    def exposure(self, index: pd.Index) -> pd.Series:
-        base_exposure = self._base_exposure(index).values
+    def exposure_parameter_source(self, index: pd.Index) -> pd.Series:
+        base_exposure = self.lookup_tables["exposure"](index).values
         joint_paf = self.joint_paf(index).values
         return pd.Series(base_exposure * (1 - joint_paf), index=index, name="values")
 
@@ -270,42 +444,17 @@ class DichotomousDistribution(Component):
     # Public methods #
     ##################
 
-    def ppf(self, x: pd.Series) -> pd.Series:
-        exposed = x < self.exposure_proportion(x.index)
+    def ppf(self, quantiles: pd.Series) -> pd.Series:
+        exposed = quantiles < self.exposure_parameters(quantiles.index)
         return pd.Series(
             exposed.replace({True: "cat1", False: "cat2"}),
             name=self.risk + ".exposure",
-            index=x.index,
-        )
-
-
-def get_distribution(risk, distribution_type, exposure, exposure_standard_deviation, weights):
-    if distribution_type == "dichotomous":
-        distribution = DichotomousDistribution(risk, exposure)
-    elif "polytomous" in distribution_type:
-        distribution = PolytomousDistribution(risk, exposure)
-    elif distribution_type == "normal":
-        distribution = ContinuousDistribution(
-            risk, mean=exposure, sd=exposure_standard_deviation, distribution=Normal
-        )
-    elif distribution_type == "lognormal":
-        distribution = ContinuousDistribution(
-            risk, mean=exposure, sd=exposure_standard_deviation, distribution=LogNormal
-        )
-    elif distribution_type == "ensemble":
-        distribution = EnsembleSimulation(
-            risk,
-            weights,
-            mean=exposure,
-            sd=exposure_standard_deviation,
+            index=quantiles.index,
         )
-    else:
-        raise NotImplementedError(f"Unhandled distribution type {distribution_type}")
-    return distribution
 
 
 def clip(q):
-    """Adjust the percentile boundary casses.
+    """Adjust the percentile boundary cases.
 
     The  risk distributions package uses the 99.9th and 0.001st percentiles
     of a log-normal distribution as the bounds of the distribution support.
@@ -319,3 +468,24 @@ def clip(q):
     q[q > Q_UPPER_BOUND] = Q_UPPER_BOUND
     q[q < Q_LOWER_BOUND] = Q_LOWER_BOUND
     return q
+
+
+def get_risk_distribution_parameter(
+    value_columns_getter: Callable[[Union[pd.DataFrame]], List[str]],
+    data: Union[float, pd.DataFrame],
+) -> Union[float, pd.Series]:
+    if isinstance(data, pd.DataFrame):
+        value_columns = value_columns_getter(data)
+        if len(value_columns) > 1:
+            raise ValueError(
+                "Expected a single value column for risk data, but found "
+                f"{len(value_columns)}: {value_columns}."
+            )
+        # don't return parameter col in continuous and ensemble distribution
+        # means to match standard deviation index
+        if "parameter" in data.columns and set(data["parameter"]) == {"continuous"}:
+            data = data.drop("parameter", axis=1)
+        index = [col for col in data.columns if col not in value_columns]
+        data = data.set_index(index)[value_columns].squeeze(axis=1)
+
+    return data