Albrja/mic-5451/Epic/update required resrouces (#519)

Albrja/mic-5451/Epic/update required resrouces

Update required resources for all components
- *Category*: Feature
- *JIRA issue*: https://jira.ihme.washington.edu/browse/MIC-5451

Changes and notes
-Update required resources for all components

### Testing
<!--
Details on how code was verified, any unit tests local for the
repo, regression testing, etc. At a minimum, this should include an
integration test for a framework change. Consider: plots, images,
(small) csv file.
-->

Author: albrja

.github/CODEOWNERS

@@ -1,2 +1,2 @@
 # default owners
-* @albrja @collijk @hussain-jafari @patricktnast @rmudambi @stevebachmeier
+* @albrja @hussain-jafari @patricktnast @rmudambi @stevebachmeier

CHANGELOG.rst

@@ -1,3 +1,7 @@
+**4.0.1 - 03/21/25**
+
+ - Feature: Update required resources for all components
+
 **4.0.0 - 03/17/25**
 
  - Feature: Use birth exposure artifact key for LBWSG components

src/vivarium_public_health/disease/model.py

@@ -114,7 +114,8 @@ def setup(self, builder: Builder) -> None:
         builder.value.register_value_modifier(
             "cause_specific_mortality_rate",
             self.adjust_cause_specific_mortality_rate,
-            requires_columns=["age", "sex"],
+            component=self,
+            required_resources=["age", "sex"],
         )
 
     def on_initialize_simulants(self, pop_data: SimulantData) -> None:

src/vivarium_public_health/disease/special_disease.py

@@ -16,6 +16,7 @@
 from vivarium import Component
 from vivarium.framework.event import Event
 from vivarium.framework.population import SimulantData
+from vivarium.framework.resource import Resource
 from vivarium.framework.values import list_combiner, union_post_processor
 
 from vivarium_public_health.disease.transition import TransitionString
@@ -123,12 +124,8 @@ def columns_required(self) -> list[str] | None:
         return ["alive"]
 
     @property
-    def initialization_requirements(self) -> dict[str, list[str]]:
-        return {
-            "requires_columns": [],
-            "requires_values": [f"{self.risk.name}.exposure"],
-            "requires_streams": [],
-        }
+    def initialization_requirements(self) -> list[str | Resource]:
+        return [self.exposure_pipeline]
 
     @property
     def state_names(self):
@@ -171,17 +168,21 @@ def setup(self, builder):
         self.disability_weight = builder.value.register_value_producer(
             f"{self.cause.name}.disability_weight",
             source=self.compute_disability_weight,
-            requires_columns=get_lookup_columns(
+            component=self,
+            required_resources=get_lookup_columns(
                 [self.lookup_tables["raw_disability_weight"]]
             ),
         )
         builder.value.register_value_modifier(
-            "all_causes.disability_weight", modifier=self.disability_weight
+            "all_causes.disability_weight",
+            modifier=self.disability_weight,
+            component=self,
         )
         builder.value.register_value_modifier(
             "cause_specific_mortality_rate",
             self.adjust_cause_specific_mortality_rate,
-            requires_columns=get_lookup_columns(
+            component=self,
+            required_resources=get_lookup_columns(
                 [self.lookup_tables["cause_specific_mortality_rate"]]
             ),
         )
@@ -190,29 +191,32 @@ def setup(self, builder):
         self.excess_mortality_rate = builder.value.register_value_producer(
             self.excess_mortality_rate_pipeline_name,
             source=self.compute_excess_mortality_rate,
-            requires_columns=get_lookup_columns(
+            component=self,
+            required_resources=get_lookup_columns(
                 [self.lookup_tables["excess_mortality_rate"]]
-            ),
-            requires_values=[self.excess_mortality_rate_paf_pipeline_name],
+            )
+            + [self.joint_paf],
         )
         self.joint_paf = builder.value.register_value_producer(
             self.excess_mortality_rate_paf_pipeline_name,
             source=lambda idx: [self.lookup_tables["population_attributable_fraction"](idx)],
+            component=self,
             preferred_combiner=list_combiner,
             preferred_post_processor=union_post_processor,
         )
         builder.value.register_value_modifier(
             "mortality_rate",
             modifier=self.adjust_mortality_rate,
-            requires_values=[self.excess_mortality_rate_pipeline_name],
+            component=self,
+            required_resources=[self.excess_mortality_rate],
         )
 
         distribution = builder.data.load(f"{self.risk}.distribution")
-        exposure_pipeline = builder.value.get_value(f"{self.risk.name}.exposure")
+        self.exposure_pipeline = builder.value.get_value(f"{self.risk.name}.exposure")
         threshold = builder.configuration[self.name].threshold
 
         self.filter_by_exposure = self.get_exposure_filter(
-            distribution, exposure_pipeline, threshold
+            distribution, self.exposure_pipeline, threshold
         )
 
     #################

src/vivarium_public_health/disease/state.py

@@ -15,6 +15,7 @@
 from vivarium.framework.engine import Builder
 from vivarium.framework.population import PopulationView, SimulantData
 from vivarium.framework.randomness import RandomnessStream
+from vivarium.framework.resource import Resource
 from vivarium.framework.state_machine import State, Transient, Transition, Trigger
 from vivarium.framework.values import Pipeline, list_combiner, union_post_processor
 from vivarium.types import DataInput, LookupTableData
@@ -57,12 +58,8 @@ def columns_required(self) -> list[str] | None:
         return [self.model, "alive"]
 
     @property
-    def initialization_requirements(self) -> dict[str, list[str]]:
-        return {
-            "requires_columns": [self.model],
-            "requires_values": [],
-            "requires_streams": [],
-        }
+    def initialization_requirements(self) -> list[str | Resource]:
+        return [self.model]
 
     #####################
     # Lifecycle methods #
@@ -327,6 +324,10 @@ class DiseaseState(BaseDiseaseState):
     # Properties #
     ##############
 
+    @property
+    def initialization_requirements(self) -> list[str | Resource]:
+        return super().initialization_requirements + [self.randomness_prevalence]
+
     @property
     def configuration_defaults(self) -> dict[str, Any]:
         configuration_defaults = super().configuration_defaults
@@ -465,19 +466,22 @@ def setup(self, builder: Builder) -> None:
         self.disability_weight = self.get_disability_weight_pipeline(builder)
 
         builder.value.register_value_modifier(
-            "all_causes.disability_weight", modifier=self.disability_weight
+            "all_causes.disability_weight",
+            modifier=self.disability_weight,
+            component=self,
         )
 
         self.has_excess_mortality = is_non_zero(
             self.lookup_tables["excess_mortality_rate"].data
         )
-        self.excess_mortality_rate = self.get_excess_mortality_rate_pipeline(builder)
         self.joint_paf = self.get_joint_paf(builder)
+        self.excess_mortality_rate = self.get_excess_mortality_rate_pipeline(builder)
 
         builder.value.register_value_modifier(
             "mortality_rate",
             modifier=self.adjust_mortality_rate,
-            requires_values=[self.excess_mortality_rate_pipeline_name],
+            component=self,
+            required_resources=[self.excess_mortality_rate],
         )
 
         self.randomness_prevalence = self.get_randomness_prevalence(builder)
@@ -531,7 +535,8 @@ def get_dwell_time_pipeline(self, builder: Builder) -> Pipeline:
         return builder.value.register_value_producer(
             f"{self.state_id}.dwell_time",
             source=self.lookup_tables["dwell_time"],
-            requires_columns=required_columns,
+            component=self,
+            required_resources=required_columns,
         )
 
     def get_disability_weight_source(self, disability_weight: DataInput | None) -> DataInput:
@@ -556,7 +561,8 @@ def get_disability_weight_pipeline(self, builder: Builder) -> Pipeline:
         return builder.value.register_value_producer(
             f"{self.state_id}.disability_weight",
             source=self.compute_disability_weight,
-            requires_columns=lookup_columns + ["alive", self.model],
+            component=self,
+            required_resources=lookup_columns + ["alive", self.model],
         )
 
     def get_excess_mortality_rate_source(
@@ -583,21 +589,24 @@ def get_excess_mortality_rate_pipeline(self, builder: Builder) -> Pipeline:
         return builder.value.register_rate_producer(
             self.excess_mortality_rate_pipeline_name,
             source=self.compute_excess_mortality_rate,
-            requires_columns=lookup_columns + ["alive", self.model],
-            requires_values=[self.excess_mortality_rate_paf_pipeline_name],
+            component=self,
+            required_resources=lookup_columns + ["alive", self.model, self.joint_paf],
         )
 
     def get_joint_paf(self, builder: Builder) -> Pipeline:
         paf = builder.lookup.build_table(0)
         return builder.value.register_value_producer(
             self.excess_mortality_rate_paf_pipeline_name,
             source=lambda idx: [paf(idx)],
+            component=self,
             preferred_combiner=list_combiner,
             preferred_post_processor=union_post_processor,
         )
 
     def get_randomness_prevalence(self, builder: Builder) -> RandomnessStream:
-        return builder.randomness.get_stream(f"{self.state_id}_prevalent_cases")
+        return builder.randomness.get_stream(
+            f"{self.state_id}_prevalent_cases", component=self
+        )
 
     ##################
     # Public methods #

src/vivarium_public_health/disease/transition.py

@@ -129,19 +129,20 @@ def __init__(
     # noinspection PyAttributeOutsideInit
     def setup(self, builder: Builder) -> None:
         lookup_columns = get_lookup_columns([self.lookup_tables["transition_rate"]])
-        self.transition_rate = builder.value.register_rate_producer(
-            self.transition_rate_pipeline_name,
-            source=self.compute_transition_rate,
-            requires_columns=lookup_columns + ["alive"],
-            requires_values=[f"{self.transition_rate_pipeline_name}.paf"],
-        )
         paf = builder.lookup.build_table(0)
         self.joint_paf = builder.value.register_value_producer(
             f"{self.transition_rate_pipeline_name}.paf",
             source=lambda index: [paf(index)],
+            component=self,
             preferred_combiner=list_combiner,
             preferred_post_processor=union_post_processor,
         )
+        self.transition_rate = builder.value.register_rate_producer(
+            self.transition_rate_pipeline_name,
+            source=self.compute_transition_rate,
+            component=self,
+            required_resources=lookup_columns + ["alive", self.joint_paf],
+        )
 
     #################
     # Setup methods #

src/vivarium_public_health/mslt/delay.py

@@ -16,6 +16,7 @@
 from vivarium.framework.engine import Builder
 from vivarium.framework.event import Event
 from vivarium.framework.population import SimulantData
+from vivarium.framework.resource import Resource
 
 
 class DelayedRisk(Component):
@@ -111,12 +112,8 @@ def columns_required(self) -> list[str] | None:
         return ["age", "sex", "population"]
 
     @property
-    def initialization_requirements(self) -> dict[str, list[str]]:
-        return {
-            "requires_columns": ["age", "sex", "population"],
-            "requires_values": [],
-            "requires_streams": [],
-        }
+    def initialization_requirements(self) -> list[str | Resource]:
+        return ["age", "sex", "population"]
 
     #####################
     # Lifecycle methods #
@@ -168,9 +165,11 @@ def setup(self, builder: Builder) -> None:
         )
         inc_name = "{}.incidence".format(self.risk)
         inc_int_name = "{}_intervention.incidence".format(self.risk)
-        self.incidence = builder.value.register_rate_producer(inc_name, source=inc_data)
+        self.incidence = builder.value.register_rate_producer(
+            inc_name, source=inc_data, component=self
+        )
         self.int_incidence = builder.value.register_rate_producer(
-            inc_int_name, source=inc_data
+            inc_int_name, source=inc_data, component=self
         )
 
         # Load the remission rates for the BAU and intervention scenarios.
@@ -183,9 +182,11 @@ def setup(self, builder: Builder) -> None:
         )
         rem_name = "{}.remission".format(self.risk)
         rem_int_name = "{}_intervention.remission".format(self.risk)
-        self.remission = builder.value.register_rate_producer(rem_name, source=rem_data)
+        self.remission = builder.value.register_rate_producer(
+            rem_name, source=rem_data, component=self
+        )
         self.int_remission = builder.value.register_rate_producer(
-            rem_int_name, source=rem_data
+            rem_int_name, source=rem_data, component=self
         )
 
         # We apply separate mortality rates to the different exposure bins.
@@ -255,6 +256,7 @@ def setup(self, builder: Builder) -> None:
             source=builder.lookup.build_table(
                 mortality_data, key_columns=["sex"], parameter_columns=["age", "year"]
             ),
+            component=self,
         )
 
     #################
@@ -312,7 +314,7 @@ def register_modifier(self, builder: Builder, disease: str) -> None:
         for template in rate_templates:
             rate_name = template.format(disease)
             modifier = lambda ix, rate: self.incidence_adjustment(disease, ix, rate)
-            builder.value.register_value_modifier(rate_name, modifier)
+            builder.value.register_value_modifier(rate_name, modifier, component=self)
 
     ########################
     # Event-driven methods #