Albright/mic-5889/use-birth-exposure-key (#513)

Albright/mic-5889/use-birth-exposure-key

Use birth exposure key for LBWSG components
- *Category*: Feature
- *JIRA issue*: https://jira.ihme.washington.edu/browse/MIC-5889

Changes and notes
-Use birth exposure key from artifact for LBWSG 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

CHANGELOG.rst

@@ -1,3 +1,7 @@
+**4.0.0 - 03/17/25**
+
+ - Feature: Use birth exposure artifact key for LBWSG components
+
 **3.2.2 - 2/21/25**
 
  - Bugfix: Use relative risk data to get age intervals for LBWSGRiskEffect

src/vivarium_public_health/risks/implementations/low_birth_weight_and_short_gestation.py

@@ -202,6 +202,9 @@ def get_exposure_column_name(axis: str) -> str:
     @property
     def configuration_defaults(self) -> dict[str, Any]:
         configuration_defaults = super().configuration_defaults
+        configuration_defaults[self.name]["data_sources"][
+            "exposure"
+        ] = f"{self.risk}.birth_exposure"
         configuration_defaults[self.name]["distribution_type"] = "lbwsg"
         return configuration_defaults
 

tests/risks/test_low_birth_weight_and_short_gestation.py

@@ -41,9 +41,7 @@ def test_parsing_lbwsg_descriptions(description, expected_weight_values, expecte
     assert age_interval.right == expected_age_values[1]
 
 
-def test_lbwsg_risk_effect_rr_pipeline(
-    base_config, base_plugins, mocker, mock_rr_interpolators
-):
+def test_lbwsg_risk_effect_rr_pipeline(base_config, base_plugins, mock_rr_interpolators):
 
     risk = LBWSGRisk()
     lbwsg_effect = LBWSGRiskEffect("cause.test_cause.cause_specific_mortality_rate")
@@ -63,7 +61,7 @@ def test_lbwsg_risk_effect_rr_pipeline(
 
     # Add data dict to add to artifact
     data = {
-        f"{risk.name}.exposure": exposure,
+        f"{risk.name}.birth_exposure": exposure,
         f"{risk.name}.relative_risk": rr_data,
         f"{risk.name}.population_attributable_fraction": 0,
         f"{risk.name}.categories": categories,
@@ -122,6 +120,59 @@ def map_age_groups(value):
                 assert (actual_rr == 1.0).all()
 
 
+def test_use_birth_exposure(base_config, base_plugins, mock_rr_interpolators):
+    risk = LBWSGRisk()
+    lbwsg_effect = LBWSGRiskEffect("cause.test_cause.cause_specific_mortality_rate")
+
+    # Add mock data to artifact
+    categories = {
+        "cat81": "Neonatal preterm and LBWSG (estimation years) - [28, 30) wks, [2500, 3000) g",
+        "cat82": "Neonatal preterm and LBWSG (estimation years) - [28, 30) wks, [3000, 3500) g",
+    }
+    # Create exposure with matching demograph index as age_bins
+    age_bins = make_age_bins()
+    ages = age_bins.drop(columns="age_group_name")
+    # Have to match age bins and rr data to make age intervals
+    rr_data = make_categorical_data(ages)
+    # Format birth exposure data
+    exposure = pd.DataFrame(
+        {
+            "sex": ["Male", "Female", "Male", "Female"],
+            "year_start": [2021, 2021, 2021, 2021],
+            "year_end": [2022, 2022, 2022, 2022],
+            "parameter": ["cat81", "cat81", "cat82", "cat82"],
+            "value": [0.75, 0.75, 0.25, 0.25],
+        }
+    )
+
+    # Add data dict to add to artifact
+    data = {
+        f"{risk.name}.birth_exposure": exposure,
+        f"{risk.name}.relative_risk": rr_data,
+        f"{risk.name}.population_attributable_fraction": 0,
+        f"{risk.name}.categories": categories,
+        f"{risk.name}.relative_risk_interpolator": mock_rr_interpolators,
+    }
+
+    # Only have neontal age groups
+    age_start = 0.0
+    age_end = 1.0
+    base_config.update(
+        {
+            "population": {
+                "initialization_age_start": age_start,
+                "initialization_age_max": age_end,
+            }
+        }
+    )
+    sim = _setup_risk_effect_simulation(base_config, base_plugins, risk, lbwsg_effect, data)
+    pop = sim.get_population()
+
+    # Assert LBWSG birth exposure columns were created
+    assert "birth_weight_exposure" in pop.columns
+    assert "gestational_age_exposure" in pop.columns
+
+
 def make_categorical_data(data: pd.DataFrame) -> pd.DataFrame:
     # Takes age gropus and adds sex, years, categories, and values
     dfs = []