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

vivarium_public_health/risks/base_risk.py

@@ -7,20 +7,26 @@ This module contains tools for modeling categorical and continuous risk
 exposure.
 
 """
+
 from typing import Any, Dict, List
 
 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.values import Pipeline
 
-from vivarium_public_health.risks.data_transformations import (
-    get_exposure_post_processor,
+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.risks.distributions import SimulationDistribution
-from vivarium_public_health.utilities import EntityString
+from vivarium_public_health.utilities import EntityString, get_lookup_columns
 
 
 class Risk(Component):
@@ -81,25 +87,42 @@ class Risk(Component):
 
     """
 
-    CONFIGURATION_DEFAULTS = {
-        "risk": {
-            "exposure": "data",
-            "rebinned_exposed": [],
-            "category_thresholds": [],
-        }
+    exposure_distributions = {
+        "dichotomous": DichotomousDistribution,
+        "ordered_polytomous": PolytomousDistribution,
+        "unordered_polytomous": PolytomousDistribution,
+        "normal": ContinuousDistribution,
+        "lognormal": ContinuousDistribution,
+        "ensemble": EnsembleDistribution,
     }
 
     ##############
     # Properties #
     ##############
 
+    @property
+    def name(self) -> str:
+        return self.risk
+
     @property
     def configuration_defaults(self) -> Dict[str, Any]:
-        return {self.risk.name: self.CONFIGURATION_DEFAULTS["risk"]}
+        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": [],
+            }
+        }
 
     @property
     def columns_created(self) -> List[str]:
-        return [self.propensity_column_name]
+        return [self.propensity_column_name, self.exposure_column_name]
 
     @property
     def initialization_requirements(self) -> Dict[str, List[str]]:
@@ -122,30 +145,98 @@ class Risk(Component):
         """
         super().__init__()
         self.risk = EntityString(risk)
-        self.exposure_distribution = self.get_exposure_distribution()
-        self._sub_components = [self.exposure_distribution]
+        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:
+        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)
 
-    ##########################
-    # Initialization methods #
-    ##########################
+    def get_distribution_type(self, builder: Builder) -> str:
+        """
+        Get the distribution type for the risk from the configuration.
 
-    def get_exposure_distribution(self) -> SimulationDistribution:
-        return SimulationDistribution(self.risk)
+        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.
 
-    #################
-    # Setup methods #
-    #################
+        Parameters
+        ----------
+        builder : Builder
+            the builder object
+
+        Returns
+        -------
+        str
+            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 : Builder
+            the builder object
+
+        Returns
+        -------
+        RiskExposureDistribution
+            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)
@@ -162,12 +253,18 @@ class Risk(Component):
         )
 
     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,
-            requires_columns=["age", "sex"],
-            requires_values=[self.propensity_pipeline_name],
-            preferred_post_processor=get_exposure_post_processor(builder, self.risk),
+            requires_columns=required_columns,
+            requires_values=[
+                self.propensity_pipeline_name,
+                self.exposure_distribution.parameters_pipeline_name,
+            ],
+            preferred_post_processor=get_exposure_post_processor(builder, self.name),
         )
 
     ########################
@@ -175,11 +272,19 @@ class Risk(Component):
     ########################
 
     def on_initialize_simulants(self, pop_data: SimulantData) -> None:
-        self.population_view.update(
-            pd.Series(
-                self.randomness.get_draw(pop_data.index), name=self.propensity_column_name
-            )
+        propensity_values = self.randomness.get_draw(pop_data.index)
+        df = pd.DataFrame(
+            {
+                self.propensity_column_name: self.randomness.get_draw(pop_data.index),
+                self.exposure_column_name: self.exposure_distribution.ppf(propensity_values),
+            }
         )
+        self.population_view.update(df)
+
+    def on_time_step_prepare(self, event: Event) -> None:
+        exposure_values = self.exposure(event.index)
+        exposure_col = pd.Series(exposure_values, name=self.exposure_column_name)
+        self.population_view.update(exposure_col)
 
     ##################################
     # Pipeline sources and modifiers #