Albrja/mic 6510/gbd to databundle (#86)

Albrja/mic 6510/gbd to databundle

Pipe GBD feature to DataBundle
- *Category*: Feature
- *JIRA issue*: https://jira.ihme.washington.edu/browse/MIC-6510

Changes and notes
-allows GBD to be a source for DataBundle
-same implementation as artifact

Author: albrja

src/vivarium_testing_utils/automated_validation/bundle.py

@@ -47,8 +47,8 @@ def __init__(
     def dataset_names(self) -> dict[str, str]:
         """Return a dictionary of required datasets for the specified source."""
         if self.source == DataSource.SIM:
-            return self.measure.sim_datasets
-        elif self.source == DataSource.ARTIFACT:
+            return self.measure.sim_output_datasets
+        elif self.source in ([DataSource.ARTIFACT, DataSource.GBD]):
             return self.measure.sim_input_datasets
         else:
             raise ValueError(f"Unsupported data source: {self.source}")
@@ -107,11 +107,9 @@ def _get_formatted_datasets(
             datasets = self.measure.get_ratio_datasets_from_sim(
                 **raw_datasets,
             )
-        elif self.source == DataSource.ARTIFACT:
+        elif self.source in [DataSource.ARTIFACT, DataSource.GBD]:
             data = self.measure.get_measure_data_from_sim_inputs(**raw_datasets)
             datasets = {"data": data}
-        elif self.source == DataSource.GBD:
-            raise NotImplementedError
         elif self.source == DataSource.CUSTOM:
             raise NotImplementedError
         else:
@@ -132,7 +130,7 @@ def _get_aggregated_weights(
         self, data_loader: DataLoader, age_group_data: pd.DataFrame
     ) -> pd.DataFrame | None:
         """Fetches and aggregates weights if required by the measure."""
-        if self.source != DataSource.ARTIFACT:
+        if self.source not in [DataSource.ARTIFACT, DataSource.GBD]:
             return None
 
         raw_weights = data_loader._get_raw_data_from_source(
@@ -147,10 +145,8 @@ def get_measure_data(
         """Get the measure data, optionally aggregated over specified stratifications."""
         if self.source == DataSource.SIM:
             return self._aggregate_scenario_stratifications(self.datasets, stratifications)
-        elif self.source == DataSource.ARTIFACT:
-            return self._aggregate_artifact_stratifications(stratifications)
-        elif self.source == DataSource.GBD:
-            raise NotImplementedError
+        elif self.source in [DataSource.ARTIFACT, DataSource.GBD]:
+            return self._aggregate_sim_input_stratifications(stratifications)
         elif self.source == DataSource.CUSTOM:
             raise NotImplementedError
         else:
@@ -167,7 +163,7 @@ def _aggregate_scenario_stratifications(
         }
         return self.measure.get_measure_data_from_ratio(**datasets)
 
-    def _aggregate_artifact_stratifications(
+    def _aggregate_sim_input_stratifications(
         self, stratifications: Collection[str] | Literal["all"]
     ) -> pd.DataFrame:
         """Aggregate the artifact data over specified stratifications. Stratifactions will be retained

src/vivarium_testing_utils/automated_validation/data_transformation/measures.py

@@ -58,7 +58,7 @@ def __str__(self) -> str:
 
     @property
     @abstractmethod
-    def sim_datasets(self) -> dict[str, str]:
+    def sim_output_datasets(self) -> dict[str, str]:
         """Return a dictionary of required datasets for this measure."""
         pass
 
@@ -101,7 +101,7 @@ def __init__(
         self.denominator = denominator
 
     @property
-    def sim_datasets(self) -> dict[str, str]:
+    def sim_output_datasets(self) -> dict[str, str]:
         """Return a dictionary of required datasets for this measure."""
         return {
             "numerator_data": self.numerator.raw_dataset_name,

tests/automated_validation/data_transformation/test_measures.py

@@ -42,7 +42,7 @@ def test_incidence(
     measure = Incidence(cause)
     assert measure.measure_key == f"cause.{cause}.incidence_rate"
     assert measure.title == "Disease Incidence Rate"
-    assert measure.sim_datasets == {
+    assert measure.sim_output_datasets == {
         "numerator_data": f"transition_count_{cause}",
         "denominator_data": f"person_time_{cause}",
     }
@@ -72,7 +72,7 @@ def test_prevalence(person_time_data: pd.DataFrame) -> None:
     measure = Prevalence(cause)
     assert measure.measure_key == f"cause.{cause}.prevalence"
     assert measure.title == "Disease Prevalence"
-    assert measure.sim_datasets == {
+    assert measure.sim_output_datasets == {
         "numerator_data": f"person_time_{cause}",
         "denominator_data": f"person_time_{cause}",
     }
@@ -160,7 +160,7 @@ def test_si_remission(
     measure = SIRemission(cause)
     assert measure.measure_key == f"cause.{cause}.remission_rate"
     assert measure.title == "Disease Remission Rate"
-    assert measure.sim_datasets == {
+    assert measure.sim_output_datasets == {
         "numerator_data": f"transition_count_{cause}",
         "denominator_data": f"person_time_{cause}",
     }
@@ -190,7 +190,7 @@ def test_all_cause_mortality_rate(
     measure = CauseSpecificMortalityRate("all_causes")
     assert measure.measure_key == "cause.all_causes.cause_specific_mortality_rate"
     assert measure.title == "All Causes Cause Specific Mortality Rate"
-    assert measure.sim_datasets == {
+    assert measure.sim_output_datasets == {
         "numerator_data": "deaths",
         "denominator_data": "person_time_total",
     }
@@ -226,7 +226,7 @@ def test_cause_specific_mortality_rate(
     measure = CauseSpecificMortalityRate(cause)
     assert measure.measure_key == f"cause.{cause}.cause_specific_mortality_rate"
     assert measure.title == "Disease Cause Specific Mortality Rate"
-    assert measure.sim_datasets == {
+    assert measure.sim_output_datasets == {
         "numerator_data": f"deaths",
         "denominator_data": "person_time_total",
     }
@@ -261,7 +261,7 @@ def test_excess_mortality_rate(
     measure = ExcessMortalityRate(cause)
     assert measure.measure_key == f"cause.{cause}.excess_mortality_rate"
     assert measure.title == "Disease Excess Mortality Rate"
-    assert measure.sim_datasets == {
+    assert measure.sim_output_datasets == {
         "numerator_data": f"deaths",
         "denominator_data": f"person_time_{cause}",
     }
@@ -296,7 +296,7 @@ def test_risk_exposure(risk_state_person_time_data: pd.DataFrame) -> None:
     measure = RiskExposure(risk_factor)
     assert measure.measure_key == f"risk_factor.{risk_factor}.exposure"
     assert measure.title == "Child Stunting Exposure"
-    assert measure.sim_datasets == {
+    assert measure.sim_output_datasets == {
         "numerator_data": f"person_time_{risk_factor}",
         "denominator_data": f"person_time_{risk_factor}",
     }
@@ -366,7 +366,7 @@ def test_population_structure(person_time_data: pd.DataFrame) -> None:
 
     assert measure.measure_key == "population.structure"
     assert measure.title == "Population Structure"
-    assert measure.sim_datasets == {
+    assert measure.sim_output_datasets == {
         "numerator_data": "person_time_total",
         "denominator_data": "person_time_total",
     }
@@ -465,7 +465,7 @@ def test_categorical_relative_risk(
     assert measure.title == "Effect of Risky Risk on Disease Excess Mortality Rate"
     assert measure.affected_entity == affected_entity
     assert measure.affected_measure_name == "excess_mortality_rate"
-    assert measure.sim_datasets == {
+    assert measure.sim_output_datasets == {
         "numerator_data": "deaths",
         "denominator_data": f"person_time_{affected_entity}",
     }

tests/automated_validation/test_data_bundle.py

@@ -6,9 +6,11 @@
 import pandas as pd
 import pytest
 from pytest_mock import MockFixture
+from vivarium_inputs import interface
 
+from tests.automated_validation.conftest import NO_GBD_ACCESS
 from vivarium_testing_utils.automated_validation.bundle import RatioMeasureDataBundle
-from vivarium_testing_utils.automated_validation.constants import DataSource
+from vivarium_testing_utils.automated_validation.constants import DRAW_INDEX, DataSource
 from vivarium_testing_utils.automated_validation.data_loader import DataLoader
 from vivarium_testing_utils.automated_validation.data_transformation import age_groups
 from vivarium_testing_utils.automated_validation.data_transformation.measures import (
@@ -35,7 +37,7 @@ def test_data_bundle_init(
     )
 
     if data_source == DataSource.SIM:
-        expected_keys = set(measure.sim_datasets.keys())
+        expected_keys = set(measure.sim_output_datasets.keys())
     else:
         expected_keys = set(measure.sim_input_datasets.keys())
     assert set(bundle.dataset_names) == expected_keys
@@ -100,12 +102,10 @@ def test_get_metadata(
     assert metadata["size"] == "4 rows × 1 columns"
 
 
-@pytest.mark.parametrize("source", [DataSource.GBD, DataSource.CUSTOM])
-def test_dataset_names_value_error(
+def test_custom_data_source_dataset_names_value_error(
     mocker: MockFixture,
     mock_ratio_measure: RatioMeasure,
     sample_age_group_df: pd.DataFrame,
-    source: DataSource,
 ) -> None:
     """Test _get_formatted_datasets raises NotImplementedError for GBD source."""
     mock_data_loader = mocker.MagicMock(spec=DataLoader)
@@ -114,7 +114,7 @@ def test_dataset_names_value_error(
     with pytest.raises(ValueError):
         RatioMeasureDataBundle(
             measure=mock_ratio_measure,
-            source=source,
+            source=DataSource.CUSTOM,
             data_loader=mock_data_loader,
             age_group_df=sample_age_group_df,
         )
@@ -186,7 +186,7 @@ def test_aggregate_reference_stratifications(
         data_loader=mocker.MagicMock(spec=DataLoader),
         age_group_df=sample_age_group_df,
     )
-    aggregated = bundle._aggregate_artifact_stratifications(stratifications)
+    aggregated = bundle._aggregate_sim_input_stratifications(stratifications)
 
     if stratifications == "all":
         aggregated.equals(reference_data)
@@ -208,3 +208,53 @@ def test_aggregate_reference_stratifications(
             ),
         )
         pd.testing.assert_frame_equal(aggregated, expected)
+
+
+@pytest.mark.slow
+def test_data_bundle_gbd_source(sim_result_dir: Path) -> None:
+    """Test that GBD data source is handled correctly in RatioMeasureDataBundle."""
+    if NO_GBD_ACCESS:
+        pytest.skip("GBD access not available for this test.")
+
+    age_bins = interface.get_age_bins()
+    age_bins.index.rename({"age_group_name": age_groups.AGE_GROUP_COLUMN}, inplace=True)
+
+    incidence = Incidence("diarrheal_diseases")
+    bundle = RatioMeasureDataBundle(
+        measure=incidence,
+        source=DataSource.GBD,
+        data_loader=DataLoader(sim_result_dir),
+        age_group_df=age_bins,
+    )
+
+    assert set(bundle.dataset_names) == {"data"}
+    # Validate datasets and weights schema
+    dataset_index_names = {
+        "sex",
+        age_groups.AGE_GROUP_COLUMN,
+        "year_start",
+        "year_end",
+        DRAW_INDEX,
+    }
+    assert set(bundle.datasets["data"].index.names) == dataset_index_names
+    assert set(bundle.datasets["data"].columns) == {"value"}
+    assert bundle.weights is not None
+    assert set(bundle.weights.index.names) == dataset_index_names.union({"location"})
+    assert set(bundle.weights.columns) == {"value"}
+
+    # Validate data aggregation
+    stratify_1 = bundle.get_measure_data("all")
+    pd.testing.assert_frame_equal(stratify_1, bundle.datasets["data"])
+    stratify_2 = bundle.get_measure_data(["sex", age_groups.AGE_GROUP_COLUMN])
+    assert set(stratify_2.index.names) == {"sex", age_groups.AGE_GROUP_COLUMN, DRAW_INDEX}
+
+    metadata = bundle.get_metadata()
+    assert metadata["source"] == "gbd"
+    assert metadata["index_columns"] == "sex, year_start, year_end, input_draw, age_group"
+    assert set(metadata.keys()) == {
+        "source",
+        "index_columns",
+        "size",
+        "num_draws",
+        "input_draws",
+    }

tests/automated_validation/test_data_loader.py

@@ -210,7 +210,7 @@ def test___get_raw_data_from_source(
     data_loader = DataLoader(sim_result_dir)
     measure = Incidence("disease")
     test_raw_data = data_loader._get_raw_data_from_source(
-        measure.sim_datasets, DataSource.SIM
+        measure.sim_output_datasets, DataSource.SIM
     )
     ref_raw_data = data_loader._get_raw_data_from_source(
         measure.sim_input_datasets, DataSource.ARTIFACT