Albrja/mic 6186/intervention component (#527)

Albrja/mic 6186/intervention component

Refactor and implementation for Intervention component
- *Category*: Implementation/Refactor
- *JIRA issue*: https://jira.ihme.washington.edu/browse/MIC-6186

Changes and notes
-Adds super classes Exposure and Exposure effect
-Adds Intervention and InterventionEffect classes equivalent to Risk and RiskEffect classes
-updates RiskExposureDistribution classes to take the Exposure component instead of the entity string (exposure vs coverage)

Author: albrja

CHANGELOG.rst

@@ -1,3 +1,10 @@
+**4.3.0 - 07/29/25**
+
+  - Refactor to create two new abstract classes, Exposure and ExposureEffect
+  - Add Intervention and InterventionEffect classes
+  - Refactor to pass Exposure component to RiskExposureDistribution classes
+  - Update exposure category names for DichotomousDistribution
+  
 **4.2.6 - 07/25/25**
 
   - Feature: Support new environment creation via 'make build-env'

docs/source/api_reference/exposure/distributions.rst

@@ -0,0 +1 @@
+.. automodule:: vivarium_public_health.exposure.distributions

docs/source/api_reference/exposure/effect.rst

@@ -0,0 +1 @@
+.. automodule:: vivarium_public_health.exposure.effect

docs/source/api_reference/exposure/exposure.rst

@@ -0,0 +1 @@
+.. automodule:: vivarium_public_health.exposure.exposure

docs/source/api_reference/exposure/index.rst

@@ -0,0 +1,11 @@
+=======================
+Exposure Modeling Tools
+=======================
+
+.. automodule:: vivarium_public_health.exposure
+
+.. toctree::
+   :maxdepth: 2
+   :glob:
+
+   *

docs/source/api_reference/risks/distributions.rst

@@ -0,0 +1 @@
+.. automodule:: vivarium_public_health.treatment.intervention

docs/source/api_reference/treatment/intervention.rst

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

src/vivarium_public_health/__init__.py

@@ -0,0 +1,2 @@
+from .effect import ExposureEffect
+from .exposure import Exposure

src/vivarium_public_health/exposure/__init__.py

@@ -1,15 +1,19 @@
 """
 =================================
-Risk Exposure Distribution Models
+Exposure Distribution Models
 =================================
 
 This module contains tools for modeling several different risk
 exposure distributions.
 
 """
 
+from __future__ import annotations
+
+import warnings
 from abc import ABC, abstractmethod
 from collections.abc import Callable
+from typing import TYPE_CHECKING
 
 import numpy as np
 import pandas as pd
@@ -24,36 +28,48 @@
 from vivarium_public_health.risks.data_transformations import pivot_categorical
 from vivarium_public_health.utilities import EntityString, get_lookup_columns
 
+if TYPE_CHECKING:
+    from vivarium_public_health.exposure import Exposure
+
 
 class MissingDataError(Exception):
     pass
 
 
-class RiskExposureDistribution(Component, ABC):
+class ExposureDistribution(Component, ABC):
 
     #####################
     # Lifecycle methods #
     #####################
 
     def __init__(
         self,
-        risk: EntityString,
+        exposure_component: Exposure,
         distribution_type: str,
         exposure_data: int | float | pd.DataFrame | None = None,
     ) -> None:
         super().__init__()
-        self.risk = risk
+        self.exposure_component = exposure_component
         self.distribution_type = distribution_type
+        if (
+            self.distribution_type != "dichotomous"
+            and self.exposure_component.entity.type == "intervention"
+        ):
+            raise NotImplementedError(
+                f"Distribution type {self.distribution_type} is not supported for interventions."
+            )
         self._exposure_data = exposure_data
 
-        self.parameters_pipeline_name = f"{self.risk}.exposure_parameters"
+        self.parameters_pipeline_name = (
+            f"{self.exposure_component.entity}.exposure_parameters"
+        )
 
     #################
     # Setup methods #
     #################
 
     def get_configuration(self, builder: "Builder") -> LayeredConfigTree | None:
-        return builder.configuration[self.risk]
+        return builder.configuration[self.exposure_component.entity]
 
     @abstractmethod
     def build_all_lookup_tables(self, builder: "Builder") -> None:
@@ -62,7 +78,9 @@ def build_all_lookup_tables(self, builder: "Builder") -> None:
     def get_exposure_data(self, builder: Builder) -> 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"])
+        return self.get_data(
+            builder, self.configuration["data_sources"][self.exposure_component.exposure_type]
+        )
 
     # noinspection PyAttributeOutsideInit
     def setup(self, builder: Builder) -> None:
@@ -87,7 +105,7 @@ def ppf(self, quantiles: pd.Series) -> pd.Series:
         raise NotImplementedError
 
 
-class EnsembleDistribution(RiskExposureDistribution):
+class EnsembleDistribution(ExposureDistribution):
     ##############
     # Properties #
     ##############
@@ -106,7 +124,7 @@ def initialization_requirements(self) -> list[str | Resource]:
 
     def __init__(self, risk: EntityString, distribution_type: str = "ensemble") -> None:
         super().__init__(risk, distribution_type)
-        self._propensity = f"ensemble_propensity_{self.risk}"
+        self._propensity = f"ensemble_propensity_{self.exposure_component.entity}"
 
     #################
     # Setup methods #
@@ -129,7 +147,11 @@ def build_all_lookup_tables(self, builder: Builder) -> None:
         distributions = list(raw_weights["parameter"].unique())
 
         raw_weights = pivot_categorical(
-            builder, self.risk, raw_weights, pivot_column="parameter", reset_index=False
+            builder,
+            self.exposure_component.entity,
+            raw_weights,
+            pivot_column="parameter",
+            reset_index=False,
         )
 
         weights, parameters = rd.EnsembleDistribution.get_parameters(
@@ -201,7 +223,7 @@ def ppf(self, quantiles: pd.Series) -> pd.Series:
         return x
 
 
-class ContinuousDistribution(RiskExposureDistribution):
+class ContinuousDistribution(ExposureDistribution):
     #####################
     # Lifecycle methods #
     #####################
@@ -261,12 +283,12 @@ def ppf(self, quantiles: pd.Series) -> pd.Series:
         return x
 
 
-class PolytomousDistribution(RiskExposureDistribution):
+class PolytomousDistribution(ExposureDistribution):
     @property
     def categories(self) -> list[str]:
-        # These need to be sorted so the cumulative sum is in the ocrrect order of categories
+        # These need to be sorted so the cumulative sum is in the correct order of categories
         # and results are therefore reproducible and correct
-        return sorted(self.lookup_tables["exposure"].value_columns)
+        return sorted(self.lookup_tables[self.exposure_component.exposure_type].value_columns)
 
     #################
     # Setup methods #
@@ -277,9 +299,11 @@ def build_all_lookup_tables(self, builder: "Builder") -> None:
         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")
+            exposure_data = pivot_categorical(
+                builder, self.exposure_component.entity, exposure_data, "parameter"
+            )
 
-        self.lookup_tables["exposure"] = self.build_lookup_table(
+        self.lookup_tables[self.exposure_component.exposure_type] = self.build_lookup_table(
             builder, exposure_data, exposure_value_columns
         )
 
@@ -293,9 +317,11 @@ def get_exposure_value_columns(
     def get_exposure_parameter_pipeline(self, builder: Builder) -> Pipeline:
         return builder.value.register_value_producer(
             self.parameters_pipeline_name,
-            source=self.lookup_tables["exposure"],
+            source=self.lookup_tables[self.exposure_component.exposure_type],
             component=self,
-            required_resources=get_lookup_columns([self.lookup_tables["exposure"]]),
+            required_resources=get_lookup_columns(
+                [self.lookup_tables[self.exposure_component.exposure_type]]
+            ),
         )
 
     ##################
@@ -313,12 +339,12 @@ def ppf(self, quantiles: pd.Series) -> pd.Series:
         ).sum(axis=1)
         return pd.Series(
             np.array(self.categories)[category_index],
-            name=self.risk + ".exposure",
+            name=f"{self.exposure_component.entity}.exposure",
             index=quantiles.index,
         )
 
 
-class DichotomousDistribution(RiskExposureDistribution):
+class DichotomousDistribution(ExposureDistribution):
 
     #################
     # Setup methods #
@@ -332,11 +358,15 @@ def build_all_lookup_tables(self, builder: "Builder") -> None:
             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}")
+                raise ValueError(
+                    f"All exposures must be in the range [0, 1] for {self.exposure_component.entity}"
+                )
         elif exposure_data < 0 or exposure_data > 1:
-            raise ValueError(f"Exposure must be in the range [0, 1] for {self.risk}")
+            raise ValueError(
+                f"Exposure must be in the range [0, 1] for {self.exposure_component.entity}"
+            )
 
-        self.lookup_tables["exposure"] = self.build_lookup_table(
+        self.lookup_tables[self.exposure_component.exposure_type] = self.build_lookup_table(
             builder, exposure_data, exposure_value_columns
         )
         self.lookup_tables["paf"] = self.build_lookup_table(builder, 0.0)
@@ -350,20 +380,43 @@ def get_exposure_data(self, builder: Builder) -> int | float | pd.DataFrame:
         # rebin exposure categories
         self.validate_rebin_source(builder, exposure_data)
         rebin_exposed_categories = set(self.configuration["rebinned_exposed"])
+        # Check if risk exposure is exposed vs cat1
+        if (
+            "cat1" in exposure_data["parameter"].unique()
+            and self.exposure_component.entity.type == "risk_factor"
+        ):
+            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,
+            )
+            exposure_data["parameter"] = exposure_data["parameter"].replace(
+                {
+                    "cat1": self.exposure_component.dichotomous_exposure_category_names.exposed,
+                    "cat2": self.exposure_component.dichotomous_exposure_category_names.unexposed,
+                }
+            )
         if rebin_exposed_categories:
-            exposure_data = self._rebin_exposure_data(exposure_data, rebin_exposed_categories)
+            exposure_data = self._rebin_exposure_data(
+                exposure_data,
+                rebin_exposed_categories,
+                self.exposure_component.dichotomous_exposure_category_names.exposed,
+            )
 
-        exposure_data = exposure_data[exposure_data["parameter"] == "cat1"]
+        exposure_data = exposure_data[
+            exposure_data["parameter"]
+            == self.exposure_component.dichotomous_exposure_category_names.exposed
+        ]
         return exposure_data.drop(columns="parameter")
 
     @staticmethod
     def _rebin_exposure_data(
-        exposure_data: pd.DataFrame, rebin_exposed_categories: set
+        exposure_data: pd.DataFrame, rebin_exposed_categories: set, exposed_category_name: str
     ) -> pd.DataFrame:
         exposure_data = exposure_data[
             exposure_data["parameter"].isin(rebin_exposed_categories)
         ]
-        exposure_data["parameter"] = "cat1"
+        exposure_data["parameter"] = exposed_category_name
         exposure_data = (
             exposure_data.groupby(list(exposure_data.columns.difference(["value"])))
             .sum()
@@ -382,7 +435,7 @@ def get_exposure_value_columns(
     def setup(self, builder: Builder) -> None:
         super().setup(builder)
         self.joint_paf = builder.value.register_value_producer(
-            f"{self.risk}.exposure_parameters.paf",
+            f"{self.exposure_component.entity}.exposure_parameters.paf",
             source=lambda index: [self.lookup_tables["paf"](index)],
             component=self,
             preferred_combiner=list_combiner,
@@ -391,10 +444,12 @@ def setup(self, builder: Builder) -> None:
 
     def get_exposure_parameter_pipeline(self, builder: Builder) -> Pipeline:
         return builder.value.register_value_producer(
-            f"{self.risk}.exposure_parameters",
+            f"{self.exposure_component.entity}.exposure_parameters",
             source=self.exposure_parameter_source,
             component=self,
-            required_resources=get_lookup_columns([self.lookup_tables["exposure"]]),
+            required_resources=get_lookup_columns(
+                [self.lookup_tables[self.exposure_component.exposure_type]]
+            ),
         )
 
     ##############
@@ -405,29 +460,31 @@ 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"])
+        rebin_exposed_categories = set(
+            builder.configuration[self.exposure_component.entity]["rebinned_exposed"]
+        )
 
         if (
             rebin_exposed_categories
-            and builder.configuration[self.risk]["category_thresholds"]
+            and builder.configuration[self.exposure_component.entity]["category_thresholds"]
         ):
             raise ValueError(
                 f"Rebinning and category thresholds are mutually exclusive. "
-                f"You provided both for {self.risk.name}."
+                f"You provided both for {self.exposure_component.entity.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"category of {self.exposure_component.entity.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"{self.exposure_component.entity.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."
             )
@@ -437,7 +494,9 @@ def validate_rebin_source(self, builder, data: pd.DataFrame) -> None:
     ##################################
 
     def exposure_parameter_source(self, index: pd.Index) -> pd.Series:
-        base_exposure = self.lookup_tables["exposure"](index).values
+        base_exposure = self.lookup_tables[self.exposure_component.exposure_type](
+            index
+        ).values
         joint_paf = self.joint_paf(index).values
         return pd.Series(base_exposure * (1 - joint_paf), index=index, name="values")
 
@@ -448,8 +507,13 @@ def exposure_parameter_source(self, index: pd.Index) -> pd.Series:
     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",
+            exposed.replace(
+                {
+                    True: self.exposure_component.dichotomous_exposure_category_names.exposed,
+                    False: self.exposure_component.dichotomous_exposure_category_names.unexposed,
+                }
+            ),
+            name=f"{self.exposure_component.entity}.{self.exposure_component.exposure_type}",
             index=quantiles.index,
         )