Refactor Module Imports (#58)

* gitignore

* change to common stratify index

* make some dfs use from_product instead

* more products

* align indexes but not for pop structure

* add unique strat columns to fixtures

* cleanup

* lint

* move title to measures.py (#56)

* move title to measures.py

* lint

* add type

* fix test

* change function import strategy

* lint

* lint

* rename with "test"

* add module names

* remove accidental copy

* revert interface change

* fix typo

Author: patricktnast

src/vivarium_testing_utils/automated_validation/comparison.py

@@ -5,19 +5,12 @@
 
 from vivarium_testing_utils.automated_validation.constants import DRAW_INDEX, SEED_INDEX
 from vivarium_testing_utils.automated_validation.data_loader import DataSource
-from vivarium_testing_utils.automated_validation.data_transformation.calculations import (
-    filter_data,
-    get_singular_indices,
-    marginalize,
-)
+from vivarium_testing_utils.automated_validation.data_transformation import calculations
 from vivarium_testing_utils.automated_validation.data_transformation.measures import (
     Measure,
     RatioMeasure,
 )
-from vivarium_testing_utils.automated_validation.visualization.dataframe_utils import (
-    format_draws_sample,
-    format_metadata,
-)
+from vivarium_testing_utils.automated_validation.visualization import dataframe_utils
 
 
 class Comparison(ABC):
@@ -100,14 +93,14 @@ def __init__(
         self.test_source = test_source
         self.test_scenarios: dict[str, str] = test_scenarios if test_scenarios else {}
         self.test_datasets = {
-            key: filter_data(dataset, self.test_scenarios, drop_singles=False)
+            key: calculations.filter_data(dataset, self.test_scenarios, drop_singles=False)
             for key, dataset in test_datasets.items()
         }
         self.reference_source = reference_source
         self.reference_scenarios: dict[str, str] = (
             reference_scenarios if reference_scenarios else {}
         )
-        self.reference_data = filter_data(
+        self.reference_data = calculations.filter_data(
             reference_data, self.reference_scenarios, drop_singles=False
         )
 
@@ -131,7 +124,7 @@ def metadata(self) -> pd.DataFrame:
         measure_key = self.measure.measure_key
         test_info = self._get_metadata_from_datasets("test")
         reference_info = self._get_metadata_from_datasets("reference")
-        return format_metadata(measure_key, test_info, reference_info)
+        return dataframe_utils.format_metadata(measure_key, test_info, reference_info)
 
     def get_diff(
         self,
@@ -232,7 +225,7 @@ def _get_metadata_from_datasets(
             num_draws = dataframe.index.get_level_values(DRAW_INDEX).nunique()
             data_info["num_draws"] = f"{num_draws:,}"
             draw_values = list(dataframe.index.get_level_values(DRAW_INDEX).unique())
-            data_info[DRAW_INDEX + "s"] = format_draws_sample(draw_values)
+            data_info[DRAW_INDEX + "s"] = dataframe_utils.format_draws_sample(draw_values)
 
         # Seeds information
         if SEED_INDEX in dataframe.index.names:
@@ -266,13 +259,17 @@ def _align_datasets(self) -> tuple[pd.DataFrame, pd.DataFrame]:
         # If the test data has any index levels that are not in the reference data, marginalize
         # over those index levels.
         test_datasets = {
-            key: marginalize(self.test_datasets[key], test_indexes_to_marginalize)
+            key: calculations.marginalize(
+                self.test_datasets[key], test_indexes_to_marginalize
+            )
             for key in self.test_datasets
         }
 
         # Drop any singular index levels from the reference data if they are not in the test data.
         # If any ref-only index level is not singular, raise an error.
-        redundant_ref_indexes = set(get_singular_indices(self.reference_data).keys())
+        redundant_ref_indexes = set(
+            calculations.get_singular_indices(self.reference_data).keys()
+        )
         if not reference_indexes_to_drop.issubset(redundant_ref_indexes):
             # TODO: MIC-6075
             diff = reference_indexes_to_drop - redundant_ref_indexes

src/vivarium_testing_utils/automated_validation/constants.py

@@ -1,4 +1,22 @@
+from __future__ import annotations
+
+from enum import Enum
+
 DRAW_PREFIX = "draw_"
 
 DRAW_INDEX = "input_draw"
 SEED_INDEX = "random_seed"
+
+
+class DataSource(Enum):
+    SIM = "sim"
+    GBD = "gbd"
+    ARTIFACT = "artifact"
+    CUSTOM = "custom"
+
+    @classmethod
+    def from_str(cls, source: str) -> DataSource:
+        try:
+            return cls(source)
+        except ValueError:
+            raise ValueError(f"Source {source} not recognized. Must be one of {DataSource}")

src/vivarium_testing_utils/automated_validation/data_loader.py

@@ -1,36 +1,20 @@
 from __future__ import annotations
 
-from enum import Enum
 from pathlib import Path
 
 import pandas as pd
 import yaml
 from vivarium import Artifact
 
-from vivarium_testing_utils.automated_validation.data_transformation.calculations import (
-    clean_artifact_data,
-    marginalize,
+from vivarium_testing_utils.automated_validation.constants import DataSource
+from vivarium_testing_utils.automated_validation.data_transformation import (
+    calculations,
+    utils,
 )
 from vivarium_testing_utils.automated_validation.data_transformation.data_schema import (
     SimOutputData,
     SingleNumericColumn,
 )
-from vivarium_testing_utils.automated_validation.data_transformation.utils import check_io
-
-
-class DataSource(Enum):
-    SIM = "sim"
-    GBD = "gbd"
-    ARTIFACT = "artifact"
-    CUSTOM = "custom"
-
-    @classmethod
-    def from_str(cls, source: str) -> DataSource:
-        try:
-            return cls(source)
-        except ValueError:
-            raise ValueError(f"Source {source} not recognized. Must be one of {DataSource}")
-
 
 NONSTANDARD_ARTIFACT_KEYS = {"population.age_bins"}
 
@@ -129,7 +113,7 @@ def _add_to_cache(self, dataset_key: str, source: DataSource, data: pd.DataFrame
             raise ValueError(f"Dataset {dataset_key} already exists in the cache.")
         self._raw_datasets[source].update({dataset_key: data.copy()})
 
-    @check_io(out=SimOutputData)
+    @utils.check_io(out=SimOutputData)
     def _load_from_sim(self, dataset_key: str) -> pd.DataFrame:
         """Load the data from the simulation output directory and set the non-value columns as indices."""
         sim_data = pd.read_parquet(self._results_dir / f"{dataset_key}.parquet")
@@ -168,12 +152,12 @@ def _load_nonstandard_artifact(self, dataset_key: str) -> pd.DataFrame:
         self._artifact.clear_cache()
         return data
 
-    @check_io(out=SingleNumericColumn)
+    @utils.check_io(out=SingleNumericColumn)
     def _load_from_artifact(self, dataset_key: str) -> pd.DataFrame:
         """Load data directly from artifact, assuming correctly formatted data."""
         data: pd.DataFrame = self._artifact.load(dataset_key)
         self._artifact.clear_cache()
-        return clean_artifact_data(dataset_key, data)
+        return calculations.clean_artifact_data(dataset_key, data)
 
     def _load_from_gbd(self, dataset_key: str) -> pd.DataFrame:
         raise NotImplementedError
@@ -186,7 +170,7 @@ def _load_from_gbd(self, dataset_key: str) -> pd.DataFrame:
 
 def _convert_to_total_person_time(data: pd.DataFrame) -> pd.DataFrame:
     old_index_names = data.index.names
-    data = marginalize(data, ["entity_type", "entity", "sub_entity"])
+    data = calculations.marginalize(data, ["entity_type", "entity", "sub_entity"])
     data["entity_type"] = "none"
     data["entity"] = "total"
     data["sub_entity"] = "total"

src/vivarium_testing_utils/automated_validation/data_transformation/age_groups.py

@@ -12,10 +12,10 @@
 AgeTuple = tuple[str, int | float, int | float]
 AgeRange = tuple[int | float, int | float]
 
+from vivarium_testing_utils.automated_validation.data_transformation import utils
 from vivarium_testing_utils.automated_validation.data_transformation.data_schema import (
     SingleNumericColumn,
 )
-from vivarium_testing_utils.automated_validation.data_transformation.utils import check_io
 
 
 class AgeGroup:
@@ -396,7 +396,7 @@ def can_coerce_to(self, other: AgeSchema) -> bool:
         return True
 
 
-def format_dataframe(target_schema: AgeSchema, df: pd.DataFrame) -> pd.DataFrame:
+def _format_dataframe(target_schema: AgeSchema, df: pd.DataFrame) -> pd.DataFrame:
     """
     Format a DataFrame to match the current schema.
 
@@ -451,7 +451,7 @@ def format_dataframe(target_schema: AgeSchema, df: pd.DataFrame) -> pd.DataFrame
         return data
 
 
-@check_io(df=SingleNumericColumn, out=SingleNumericColumn)
+@utils.check_io(df=SingleNumericColumn, out=SingleNumericColumn)
 def rebin_count_dataframe(
     target_schema: AgeSchema,
     df: pd.DataFrame,
@@ -535,3 +535,17 @@ def _get_transform_matrix(source_schema: AgeSchema, target_schema: AgeSchema) ->
             if fraction > 0:
                 transform_matrix.loc[target_group.name, source_group.name] = fraction
     return transform_matrix
+
+
+def format_dataframe_from_age_bin_df(
+    data: pd.DataFrame, age_bin_df: pd.DataFrame
+) -> pd.DataFrame:
+    """Try to merge the age groups with the data. If it fails, just return the data."""
+    context_age_schema = AgeSchema.from_dataframe(age_bin_df)
+    try:
+        return _format_dataframe(context_age_schema, data)
+    except ValueError:
+        logger.info(
+            "Could not resolve age groups. The DataFrame likely has no age data. Returning dataframe as-is."
+        )
+        return data

src/vivarium_testing_utils/automated_validation/data_transformation/calculations.py

@@ -10,18 +10,11 @@
     DRAW_INDEX,
     DRAW_PREFIX,
 )
-from vivarium_testing_utils.automated_validation.data_transformation.age_groups import (
-    AgeSchema,
-    format_dataframe,
-)
+from vivarium_testing_utils.automated_validation.data_transformation import utils
 from vivarium_testing_utils.automated_validation.data_transformation.data_schema import (
     DrawData,
     SingleNumericColumn,
 )
-from vivarium_testing_utils.automated_validation.data_transformation.utils import (
-    check_io,
-    series_to_dataframe,
-)
 
 
 def filter_data(
@@ -53,7 +46,7 @@ def filter_data(
     return data
 
 
-@check_io(
+@utils.check_io(
     numerator_data=SingleNumericColumn,
     denominator_data=SingleNumericColumn,
     out=SingleNumericColumn,
@@ -109,10 +102,10 @@ def linear_combination(
     data: pd.DataFrame, coeff_a: float, col_a: str, coeff_b: float, col_b: str
 ) -> pd.DataFrame:
     """Return a series that is the linear combination of two columns in a DataFrame."""
-    return series_to_dataframe((data[col_a] * coeff_a) + (data[col_b] * coeff_b))
+    return utils.series_to_dataframe((data[col_a] * coeff_a) + (data[col_b] * coeff_b))
 
 
-@check_io(out=SingleNumericColumn)
+@utils.check_io(out=SingleNumericColumn)
 def clean_artifact_data(
     dataset_key: str,
     data: pd.DataFrame,
@@ -125,7 +118,7 @@ def clean_artifact_data(
     return data
 
 
-@check_io(data=DrawData, out=SingleNumericColumn)
+@utils.check_io(data=DrawData, out=SingleNumericColumn)
 def _clean_artifact_draws(
     data: pd.DataFrame,
 ) -> pd.DataFrame:
@@ -144,18 +137,6 @@ def _clean_artifact_draws(
     return data
 
 
-def resolve_age_groups(data: pd.DataFrame, age_groups: pd.DataFrame) -> pd.DataFrame:
-    """Try to merge the age groups with the data. If it fails, just return the data."""
-    context_age_schema = AgeSchema.from_dataframe(age_groups)
-    try:
-        return format_dataframe(context_age_schema, data)
-    except ValueError:
-        logger.info(
-            "Could not resolve age groups. The DataFrame likely has no age data. Returning dataframe as-is."
-        )
-        return data
-
-
 def get_singular_indices(data: pd.DataFrame) -> dict[str, Any]:
     """Get index levels and their values that are singular (i.e. have only one unique value)."""
     singular_metadata: dict[str, Any] = {}

src/vivarium_testing_utils/automated_validation/data_transformation/formatting.py

@@ -1,11 +1,7 @@
 import pandas as pd
 
 from vivarium_testing_utils.automated_validation.constants import DRAW_INDEX, SEED_INDEX
-from vivarium_testing_utils.automated_validation.data_transformation.calculations import (
-    filter_data,
-    marginalize,
-    stratify,
-)
+from vivarium_testing_utils.automated_validation.data_transformation import calculations
 
 
 class SimDataFormatter:
@@ -29,11 +25,11 @@ def __init__(self, measure: str, entity: str, filter_value: str) -> None:
 
     def format_dataset(self, dataset: pd.DataFrame) -> pd.DataFrame:
         """Clean up unused columns, and filter for the state."""
-        dataset = marginalize(dataset, self.unused_columns)
+        dataset = calculations.marginalize(dataset, self.unused_columns)
         if self.filter_value == "total":
-            dataset = marginalize(dataset, [*self.filters])
+            dataset = calculations.marginalize(dataset, [*self.filters])
         else:
-            dataset = filter_data(dataset, self.filters)
+            dataset = calculations.filter_data(dataset, self.filters)
         return dataset
 
 
@@ -82,7 +78,7 @@ def format_dataset(self, dataset: pd.DataFrame) -> pd.DataFrame:
         levels_to_stratify = [
             level for level in between_scenario_levels if level in dataset.index.names
         ]
-        return stratify(
+        return calculations.stratify(
             data=dataset,
             stratification_cols=levels_to_stratify,
         )
@@ -124,7 +120,7 @@ def __init__(self, entity: str, sum_all: bool = False) -> None:
         self.unused_columns = ["measure", "entity_type", "entity"]
 
     def format_dataset(self, dataset: pd.DataFrame) -> pd.DataFrame:
-        dataset = marginalize(dataset, self.unused_columns)
+        dataset = calculations.marginalize(dataset, self.unused_columns)
         if self.sum_all:
             # Get the levels to group by (all except 'sub_entity')
             group_levels = [