Albrja/mic 6508/load gbd data (#84)

Albrja/mic 6508/load gbd data

Add feature to lead gbd data from DataLoader
- *Category*: Feature
- *JIRA issue*: https://jira.ihme.washington.edu/browse/MIC-6508

Changes and notes
-pin pandas below 2.0 to resolve package dependencies with Vivarium Inputs
-use load_standard_data to load GBD data in DataLoader

Author: albrja

setup.py

@@ -45,7 +45,7 @@
 
     install_requirements = [
         "vivarium_dependencies[numpy,pyyaml,scipy,click,tables,loguru,networkx]",
-        "pandas>2.0.0",
+        "pandas",
         "vivarium_build_utils>=2.0.1,<3.0.0",
         "pyarrow",
         "seaborn",
@@ -57,8 +57,8 @@
 
     validation_requirements = [
         "vivarium>=3.4.0",
-        "vivarium-inputs",
-        "pandera",
+        "vivarium-inputs>=7.1.0, <8.0.0",
+        "pandera<0.23.0",
     ]
 
     interactive_requirements = [

src/vivarium_testing_utils/automated_validation/constants.py

@@ -20,3 +20,6 @@ def from_str(cls, source: str) -> DataSource:
             return cls(source)
         except ValueError:
             raise ValueError(f"Source {source} not recognized. Must be one of {DataSource}")
+
+
+LOCATION_ARTIFACT_KEY = "population.location"

src/vivarium_testing_utils/automated_validation/data_loader.py

@@ -6,8 +6,13 @@
 import pandas as pd
 import yaml
 from vivarium import Artifact
+from vivarium_inputs.interface import load_standard_data
 
-from vivarium_testing_utils.automated_validation.constants import DRAW_PREFIX, DataSource
+from vivarium_testing_utils.automated_validation.constants import (
+    DRAW_PREFIX,
+    LOCATION_ARTIFACT_KEY,
+    DataSource,
+)
 from vivarium_testing_utils.automated_validation.data_transformation import (
     calculations,
     utils,
@@ -23,7 +28,7 @@ def __init__(self, sim_output_dir: Path, cache_size_mb: int = 1000):
         self._cache_size_mb = cache_size_mb
 
         self._results_dir = self._sim_output_dir / "results"
-        self._raw_data_cache: dict[DataSource, dict[str, pd.DataFrame]] = {
+        self._raw_data_cache: dict[DataSource, dict[str, pd.DataFrame | str]] = {
             data_source: {} for data_source in DataSource
         }
         self._loader_mapping = {
@@ -39,6 +44,8 @@ def __init__(self, sim_output_dir: Path, cache_size_mb: int = 1000):
             self._add_to_cache(
                 data_key="person_time_total", data=person_time_total, source=DataSource.SIM
             )
+        # TODO: MIC-6533 - Update when all locations are in one artifact in the future.
+        self.location = self.get_data(LOCATION_ARTIFACT_KEY, DataSource.ARTIFACT)
 
     def _create_person_time_total_dataset(self) -> pd.DataFrame | None:
         """
@@ -82,7 +89,8 @@ def get_artifact_keys(self) -> list[str]:
     def get_data(self, data_key: str, source: DataSource) -> Any:
         """Return the data from the cache if it exists, otherwise load it from the source."""
         try:
-            return self._raw_data_cache[source][data_key].copy()
+            data = self._raw_data_cache[source][data_key]
+            return data.copy() if isinstance(data, pd.DataFrame) else data
         except KeyError:
             if source == DataSource.CUSTOM:
                 raise ValueError(
@@ -100,11 +108,14 @@ def _load_from_source(self, data_key: str, source: DataSource) -> Any:
         """Load the data from the given source via the loader mapping."""
         return self._loader_mapping[source](data_key)
 
-    def _add_to_cache(self, data_key: str, source: DataSource, data: pd.DataFrame) -> None:
+    def _add_to_cache(
+        self, data_key: str, source: DataSource, data: pd.DataFrame | str
+    ) -> None:
         """Update the raw_data_cache with the given data."""
         if data_key in self._raw_data_cache.get(source, {}):
             raise ValueError(f"Data for {data_key} already exist in the cache.")
-        self._raw_data_cache[source].update({data_key: data.copy()})
+        cache_data = data.copy() if isinstance(data, pd.DataFrame) else data
+        self._raw_data_cache[source].update({data_key: cache_data})
 
     @utils.check_io(out=SimOutputData)
     def _load_from_sim(self, data_key: str) -> pd.DataFrame:
@@ -148,11 +159,18 @@ def _load_from_artifact(self, data_key: str) -> Any:
             and not data.columns.empty
             and data.columns.str.startswith(DRAW_PREFIX).all()
         ):
-            data = calculations.clean_artifact_draws(data)
+            data = calculations.clean_draw_columns(data)
         return data
 
-    def _load_from_gbd(self, data_key: str) -> pd.DataFrame:
-        raise NotImplementedError
+    def _load_from_gbd(self, data_key: str) -> Any:
+        data = load_standard_data(data_key, self.location)
+        if (
+            isinstance(data, pd.DataFrame)
+            and not data.columns.empty
+            and data.columns.str.startswith(DRAW_PREFIX).all()
+        ):
+            data = calculations.clean_draw_columns(data)
+        return data
 
     def _get_raw_data_from_source(
         self, measure_keys: dict[str, str], source: DataSource

src/vivarium_testing_utils/automated_validation/data_transformation/age_groups.py

@@ -507,9 +507,7 @@ def rebin_count_dataframe(
         result_matrix_for_col.columns.name = AGE_GROUP_COLUMN
 
         # Stack the new age group columns into the index
-        stacked_series_for_col = result_matrix_for_col.stack(
-            level=AGE_GROUP_COLUMN, future_stack=True
-        )
+        stacked_series_for_col = result_matrix_for_col.stack(level=AGE_GROUP_COLUMN)
         stacked_series_for_col.name = val_col
 
         all_results_series.append(stacked_series_for_col)

src/vivarium_testing_utils/automated_validation/data_transformation/calculations.py

@@ -124,7 +124,7 @@ def linear_combination(
 
 
 @utils.check_io(data=DrawData, out=SingleNumericColumn)
-def clean_artifact_draws(
+def clean_draw_columns(
     data: pd.DataFrame,
 ) -> pd.DataFrame:
     """Clean the artifact data by dropping unnecessary columns and renaming the value column."""

src/vivarium_testing_utils/automated_validation/data_transformation/data_schema.py

@@ -1,5 +1,5 @@
 import pandas as pd
-import pandera.pandas as pa
+import pandera as pa
 from pandera.typing import Index
 
 

src/vivarium_testing_utils/automated_validation/data_transformation/utils.py

@@ -3,7 +3,7 @@
 from typing import Any, Callable, TypeVar
 
 import pandas as pd
-import pandera.pandas as pa
+import pandera as pa
 
 F = TypeVar("F", bound=Callable[..., Any])
 

tests/automated_validation/conftest.py

@@ -1,3 +1,4 @@
+import shutil
 from pathlib import Path
 from unittest import mock
 
@@ -7,7 +8,11 @@
 from pytest import TempPathFactory
 from vivarium.framework.artifact import Artifact
 
-from vivarium_testing_utils.automated_validation.constants import DRAW_INDEX, SEED_INDEX
+from vivarium_testing_utils.automated_validation.constants import (
+    DRAW_INDEX,
+    LOCATION_ARTIFACT_KEY,
+    SEED_INDEX,
+)
 from vivarium_testing_utils.automated_validation.data_loader import (
     _convert_to_total_person_time,
 )
@@ -529,7 +534,7 @@ def _make_artifact_prevalence() -> pd.DataFrame:
 
 
 @pytest.fixture(scope="session")
-def _artifact_keys_mapper() -> dict[str, pd.DataFrame | dict[str, str]]:
+def _artifact_keys_mapper() -> dict[str, str | pd.DataFrame | dict[str, str]]:
     _raw_artifact_disease_incidence = _create_raw_artifact_disease_incidence()
     _raw_artifact_risk_exposure = _create_raw_artifact_risk_exposure()
     _sample_age_group_df = _create_sample_age_group_df()
@@ -543,6 +548,7 @@ def _artifact_keys_mapper() -> dict[str, pd.DataFrame | dict[str, str]]:
         "risk_factor.risky_risk.categories": _risk_categories,
         "population.structure": _population_structure,
         "cause.disease.prevalence": _artifact_prevalence,
+        LOCATION_ARTIFACT_KEY: "Ethiopia",
     }
 
 
@@ -603,3 +609,11 @@ def reference_weights() -> pd.DataFrame:
             names=["year", "sex", "age"],
         ),
     )
+
+
+def is_on_slurm() -> bool:
+    """Returns True if the current environment is a SLURM cluster."""
+    return not shutil.which("sbatch") is not None
+
+
+NO_GBD_ACCESS = is_on_slurm()

tests/automated_validation/data_transformation/test_calculations.py

@@ -244,19 +244,19 @@ def test_aggregate_sum_preserves_string_order() -> None:
         (
             ["sex"],
             [
-                2.83,
                 6.92,
+                2.83,
             ],  # Male: (20*2 + 100*3)/(20+100) ≈ 2.83, Female: (2*5 + 50*7)/(2+50) ≈ 6.92
-            pd.Index(["Male", "Female"], name="sex"),
+            pd.Index(["Female", "Male"], name="sex"),
         ),
         # Test aggregating by color
         (
             ["color"],
             [
-                2.27,
                 4.33,
+                2.27,
             ],  # Red: (20*2 + 2*5)/(20+2) ≈ 2.27, Blue: (100*3 + 50*7)/(100+50) ≈ 4.33
-            pd.Index(["Red", "Blue"], name="color"),
+            pd.Index(["Blue", "Red"], name="color"),
         ),
         # Test no aggregation - keeping all index levels
         (

tests/automated_validation/data_transformation/test_data_schema.py

@@ -1,5 +1,5 @@
 import pandas as pd
-import pandera.pandas as pa
+import pandera as pa
 import pytest
 from pandera.errors import SchemaError