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

vivarium_public_health/results/intervention.py

@@ -0,0 +1,139 @@
+"""
+======================
+Intervention Observers
+======================
+
+This module contains tools for observing risk exposure during the simulation.
+
+"""
+
+import pandas as pd
+from vivarium.framework.engine import Builder
+
+from vivarium_public_health.results.columns import COLUMNS
+from vivarium_public_health.results.observer import PublicHealthObserver
+from vivarium_public_health.utilities import to_years
+
+
+class CategoricalInterventionObserver(PublicHealthObserver):
+    """
+    A class for observering interventions. This class has the same implementation as
+    the 'CategoricalRiskObserver' class.
+
+    """
+
+    @property
+    def columns_required(self) -> list[str] | None:
+        """The columns required by this observer."""
+        return ["alive"]
+
+    #####################
+    # Lifecycle methods #
+    #####################
+
+    def __init__(self, intervention: str) -> None:
+        """Constructor for this observer.
+
+        Parameters
+        ----------
+        intervention
+            The name of the intervention being observed
+        """
+        super().__init__()
+        self.intervention = intervention
+        self.coverage_pipeline_name = f"{self.intervention}.coverage"
+
+    #################
+    # Setup methods #
+    #################
+
+    def setup(self, builder: Builder) -> None:
+        """Set up the observer."""
+        self.step_size = builder.time.step_size()
+        self.categories = builder.data.load(f"intervention.{self.intervention}.categories")
+
+    def get_configuration_name(self) -> str:
+        return self.intervention
+
+    def register_observations(self, builder: Builder) -> None:
+        """Register a stratification and observation.
+
+        Notes
+        -----
+        While it's typical for all stratification registrations to be encapsulated
+        in a single class (i.e. the
+        :class:ResultsStratifier <vivarium_public_health.results.stratification.ResultsStratifier),
+        this observer registers an additional one. While it could be registered
+        in the ``ResultsStratifier`` as well, it is specific to this observer and
+        so it is registered here while we have easy access to the required categories
+        and value names.
+        """
+        builder.results.register_stratification(
+            f"{self.intervention}",
+            list(self.categories.keys()),
+            requires_values=[self.coverage_pipeline_name],
+        )
+        self.register_adding_observation(
+            builder=builder,
+            name=f"person_time_{self.intervention}",
+            pop_filter=f'alive == "alive" and tracked==True',
+            when="time_step__prepare",
+            requires_columns=["alive"],
+            requires_values=[self.coverage_pipeline_name],
+            additional_stratifications=self.configuration.include + [self.intervention],
+            excluded_stratifications=self.configuration.exclude,
+            aggregator=self.aggregate_intervention_category_person_time,
+        )
+
+    ###############
+    # Aggregators #
+    ###############
+
+    def aggregate_intervention_category_person_time(self, x: pd.DataFrame) -> float:
+        """Aggregate the person time for this time step."""
+        return len(x) * to_years(self.step_size())
+
+    ##############################
+    # Results formatting methods #
+    ##############################
+
+    def format(self, measure: str, results: pd.DataFrame) -> pd.DataFrame:
+        """Rename the appropriate column to 'sub_entity'.
+
+        The primary thing this method does is rename the risk column
+        to 'sub_entity'. We do this here instead of the 'get_sub_entity_column'
+        method simply because we do not want the risk column at all. If we keep
+        it here and then return it as the sub-entity column later, the final
+        results would have both.
+
+        Parameters
+        ----------
+        measure
+            The measure.
+        results
+            The results to format.
+
+        Returns
+        -------
+            The formatted results.
+        """
+        results = results.reset_index()
+        results.rename(columns={self.intervention: COLUMNS.SUB_ENTITY}, inplace=True)
+        return results
+
+    def get_measure_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        """Get the 'measure' column values."""
+        return pd.Series("person_time", index=results.index)
+
+    def get_entity_type_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        """Get the 'entity_type' column values."""
+        return pd.Series("rei", index=results.index)
+
+    def get_entity_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        """Get the 'entity' column values."""
+        return pd.Series(self.intervention, index=results.index)
+
+    def get_sub_entity_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        """Get the 'sub_entity' column values."""
+        # The sub-entity col was created in the 'format' method
+        return results[COLUMNS.SUB_ENTITY]

vivarium_public_health/risks/base_risk.py

@@ -7,30 +7,14 @@ This module contains tools for modeling categorical and continuous risk
 exposure.
 
 """
+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.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
+from vivarium_public_health.exposure import Exposure
 
 
-class Risk(Component):
+class Risk(Exposure):
     """A model for a risk factor defined by either a continuous or a categorical value.
 
     For example,
@@ -89,217 +73,34 @@ class Risk(Component):
 
     """
 
-    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
+    def exposure_type(self) -> str:
+        """The measure of the risk exposure."""
+        return "exposure"
 
     @property
-    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": [],
-            }
-        }
+    def dichotomous_exposure_category_names(self) -> NamedTuple:
+        """The name of the exposed category for this intervention."""
 
-    @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
+        class __Categories(NamedTuple):
+            exposed: str = "exposed"
+            unexposed: str = "unexposed"
 
-    @property
-    def initialization_requirements(self) -> list[str | Resource]:
-        return [self.randomness]
+        categories = __Categories()
+        return categories
 
     #####################
     # 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:
-        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
+        super().setup(builder)
+        # We want to set this to True if 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.risk.name}_on_")
+                if component.startswith(f"non_log_linear_risk_effect.{self.entity.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)