VTU Mypy (#34)

* type hint for calculations

* in flux

* in progress

* do approximately half of typing

* removoe more series types

* lint

* data_loader

* type measures.py

* fix strats

* type test_interface

* type test_formatting

* type test_data_schema

* refactor patch

* lint

* remove fuzzy checker ignores

* "delete unused types.py file"

* merge changes from refactor and fix

* address comments

* remve ref ro 'fuzzy'

* move error up

* remove trailing commas

* change to collection

* empty commit

* lint

* remove unused import

Author: patricktnast

pyproject.toml

@@ -30,7 +30,10 @@ exclude = [
     "build",
     "setup.py",
     "docs/source/conf.py",
+    "src/vivarium_testing_utils/automated_validation/comparison.py", # Only stubbed out currently
+    "src/vivarium_testing_utils/automated_validation/plot_utils.py", # Only stubbed out currently
 ]
+plugins = ["pandera.mypy"]
 
 # handle mypy errors when 3rd party packages are not typed.
 [[tool.mypy.overrides]]

setup.py

@@ -47,9 +47,9 @@
         "vivarium_dependencies[pandas,numpy,pyyaml,scipy,click,tables,loguru,networkx]",
         "vivarium_build_utils>=2.0.1,<3.0.0",
         "pyarrow",
+        "vivarium",
         # Type stubs
         "types-setuptools",
-        "vivarium",
     ]
 
     setup_requires = ["setuptools_scm"]

src/vivarium_testing_utils/automated_validation/comparison.py

@@ -1,11 +1,8 @@
 from abc import ABC, abstractmethod
+from typing import Collection
 
 import pandas as pd
 
-from vivarium_testing_utils.automated_validation.data_transformation.data_schema import (
-    RatioData,
-    SingleNumericColumn,
-)
 from vivarium_testing_utils.automated_validation.data_transformation.measures import (
     Measure,
     RatioMeasure,
@@ -24,36 +21,36 @@ class Comparison(ABC):
     stratifications: list[str]
 
     @abstractmethod
-    def verify(self, stratifications: list[str]):
+    def verify(self, stratifications: Collection[str] = ()):
         pass
 
     @abstractmethod
-    def summarize(self, stratifications: list[str]):
+    def summarize(self, stratifications: Collection[str] = ()):
         pass
 
     @abstractmethod
-    def heads(self, stratifications: list[str]):
+    def heads(self, stratifications: Collection[str] = ()):
         pass
 
 
-class FuzzyComparison:
+class FuzzyComparison(Comparison):
     def __init__(
         self,
         measure: RatioMeasure,
         test_data: pd.DataFrame,
         reference_data: pd.DataFrame,
-        stratifications: list[str] = [],
+        stratifications: Collection[str] = (),
     ):
         self.measure = measure
         self.test_data = test_data
         self.reference_data = reference_data
         self.stratifications = stratifications
 
-    def verify(self, stratifications: list[str]):
+    def verify(self, stratifications: Collection[str] = ()):
         raise NotImplementedError
 
-    def summarize(self, stratifications: list[str]):
+    def summarize(self, stratifications: Collection[str] = ()):
         raise NotImplementedError
 
-    def heads(self, stratifications: list[str]):
+    def heads(self, stratifications: Collection[str] = ()):
         raise NotImplementedError

src/vivarium_testing_utils/automated_validation/data_loader.py

@@ -4,7 +4,6 @@
 from pathlib import Path
 
 import pandas as pd
-import pandera as pa
 import yaml
 from vivarium import Artifact
 
@@ -33,16 +32,18 @@ def from_str(cls, source: str) -> DataSource:
 
 
 class DataLoader:
-    def __init__(self, sim_output_dir: str, cache_size_mb: int = 1000):
-        self._sim_output_dir = Path(sim_output_dir)
-        self._results_dir = self._sim_output_dir / "results"
+    def __init__(self, sim_output_dir: Path, cache_size_mb: int = 1000):
+        self._sim_output_dir = sim_output_dir
         self._cache_size_mb = cache_size_mb
-        self._raw_datasets = {data_source: {} for data_source in DataSource}
+
+        self._results_dir = self._sim_output_dir / "results"
+        self._raw_datasets: dict[DataSource, dict[str, pd.DataFrame]] = {
+            data_source: {} for data_source in DataSource
+        }
         self._loader_mapping = {
             DataSource.SIM: self._load_from_sim,
             DataSource.GBD: self._load_from_gbd,
             DataSource.ARTIFACT: self._load_from_artifact,
-            DataSource.CUSTOM: self._raise_custom_data_error,
         }
         self._artifact = self._load_artifact(self._sim_output_dir)
 
@@ -59,14 +60,19 @@ def get_dataset(self, dataset_key: str, source: DataSource) -> pd.DataFrame:
         try:
             return self._raw_datasets[source][dataset_key].copy()
         except KeyError:
+            if source == DataSource.CUSTOM:
+                raise ValueError(
+                    f"No custom dataset found for {dataset_key}."
+                    "Please upload a dataset using ValidationContext.upload_custom_data."
+                )
             dataset = self._load_from_source(dataset_key, source)
             self._add_to_cache(dataset_key, source, dataset)
             return dataset
 
-    def upload_custom_data(self, dataset_key: str, data: pd.DataFrame | pd.Series) -> None:
+    def upload_custom_data(self, dataset_key: str, data: pd.DataFrame) -> None:
         self._add_to_cache(dataset_key, DataSource.CUSTOM, data)
 
-    def _load_from_source(self, dataset_key: str, source: DataSource) -> None:
+    def _load_from_source(self, dataset_key: str, source: DataSource) -> pd.DataFrame:
         """Load the data from the given source via the loader mapping."""
         return self._loader_mapping[source](dataset_key)
 
@@ -102,8 +108,8 @@ def _load_from_sim(self, dataset_key: str) -> pd.DataFrame:
         return multi_index_df
 
     @staticmethod
-    def _load_artifact(results_dir: str) -> Artifact:
-        model_spec_path = Path(results_dir) / "model_specification.yaml"
+    def _load_artifact(results_dir: Path) -> Artifact:
+        model_spec_path = results_dir / "model_specification.yaml"
         artifact_path = yaml.safe_load(model_spec_path.open("r"))["configuration"][
             "input_data"
         ]["artifact_path"]
@@ -117,9 +123,3 @@ def _load_from_artifact(self, dataset_key: str) -> pd.DataFrame:
 
     def _load_from_gbd(self, dataset_key: str) -> pd.DataFrame:
         raise NotImplementedError
-
-    def _raise_custom_data_error(self, dataset_key: str) -> pd.DataFrame:
-        raise ValueError(
-            f"No custom dataset found for {dataset_key}."
-            "Please upload a dataset using ValidationContext.upload_custom_data."
-        )

src/vivarium_testing_utils/automated_validation/data_transformation/calculations.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import TypeVar
 
 import pandas as pd
@@ -12,22 +14,20 @@
     series_to_dataframe,
 )
 
-DataSet = TypeVar("DataSet", pd.DataFrame, pd.Series)
-
 DRAW_PREFIX = "draw_"
 
 
-def align_indexes(datasets: list[DataSet]) -> list[DataSet]:
+def align_indexes(datasets: list[pd.DataFrame]) -> list[pd.DataFrame]:
     """Put each dataframe on a common index by choosing the intersection of index columns
     and marginalizing over the rest."""
     # Get the common index columns
-    common_index = set.intersection(*(set(data.index.names) for data in datasets))
+    common_index = list(set.intersection(*(set(data.index.names) for data in datasets)))
 
     # Marginalize over the rest
     return [marginalize(data, common_index) for data in datasets]
 
 
-def filter_data(data: DataSet, filter_cols: dict[str, list]) -> DataSet:
+def filter_data(data: pd.DataFrame, filter_cols: dict[str, list[str]]) -> pd.DataFrame:
     """Filter a DataFrame by the given index columns and values.
 
     The filter_cols argument
@@ -57,21 +57,21 @@ def ratio(data: pd.DataFrame, numerator: str, denominator: str) -> pd.DataFrame:
     return series_to_dataframe(data[numerator] / data[denominator])
 
 
-def aggregate_sum(data: DataSet, groupby_cols: list[str]) -> DataSet:
+def aggregate_sum(data: pd.DataFrame, groupby_cols: list[str]) -> pd.DataFrame:
     """Aggregate the dataframe over the specified index columns by summing."""
     if not groupby_cols:
         return data
     return data.groupby(groupby_cols).sum()
 
 
-def stratify(data: DataSet, stratification_cols: list[str]) -> DataSet:
+def stratify(data: pd.DataFrame, stratification_cols: list[str]) -> pd.DataFrame:
     """Stratify the data by the index columns, summing over everything else. Syntactic sugar for aggregate."""
     return aggregate_sum(data, stratification_cols)
 
 
-def marginalize(data: DataSet, marginalize_cols: list[str]) -> DataSet:
+def marginalize(data: pd.DataFrame, marginalize_cols: list[str]) -> pd.DataFrame:
     """Sum over marginalize columns, keeping the rest. Syntactic sugar for aggregate."""
-    return aggregate_sum(data, data.index.names.difference(marginalize_cols))
+    return aggregate_sum(data, [x for x in data.index.names if x not in marginalize_cols])
 
 
 def linear_combination(

src/vivarium_testing_utils/automated_validation/data_transformation/formatting.py

@@ -1,14 +1,9 @@
-from abc import ABC, abstractmethod
-
 import pandas as pd
 
 from vivarium_testing_utils.automated_validation.data_transformation.calculations import (
     filter_data,
     marginalize,
 )
-from vivarium_testing_utils.automated_validation.data_transformation.data_schema import (
-    SimOutputData,
-)
 
 
 class SimDataFormatter:
@@ -30,7 +25,7 @@ def __init__(self, type: str, cause: str, filter_value: str) -> None:
         self.filter_value = filter_value
         self.new_value_column_name = f"{self.filter_value}_{self.type}"
 
-    def format_dataset(self, dataset: SimOutputData) -> SimOutputData:
+    def format_dataset(self, dataset: pd.DataFrame) -> pd.DataFrame:
         """Clean up redundant columns, filter for the state, and rename the value column."""
         for column, value in self.redundant_columns.items():
             dataset = _drop_redundant_index(
@@ -56,13 +51,13 @@ def __init__(self, cause: str, start_state: str, end_state: str) -> None:
 class PersonTime(SimDataFormatter):
     """Formatter for simulation data that contains person time."""
 
-    def __init__(self, cause: str, state=None) -> None:
+    def __init__(self, cause: str, state: str | None = None) -> None:
         super().__init__("person_time", cause, state or "total")
 
 
 def _drop_redundant_index(
     data: pd.DataFrame, idx_column_name: str, idx_column_value: str
-) -> None:
+) -> pd.DataFrame:
     """Validate that a DataFrame column is singular-valued, then drop it from the index."""
     # TODO: Make sure we handle this case appropriately when we
     # want to automatically add many comparisons

src/vivarium_testing_utils/automated_validation/data_transformation/measures.py

@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+from typing import Any
 
 import pandas as pd
 import pandera as pa
@@ -25,21 +26,32 @@ class Measure(ABC):
     """A Measure contains key information and methods to take raw data from a DataSource
     and process it into an epidemiological measure suitable for use in a Comparison."""
 
-    sim_datasets: dict[str, str]
-    artifact_datasets: dict[str, str]
+    measure_key: str
+
+    @property
+    @abstractmethod
+    def sim_datasets(self) -> dict[str, str]:
+        """Return a dictionary of required datasets for this measure."""
+        pass
+
+    @property
+    @abstractmethod
+    def artifact_datasets(self) -> dict[str, str]:
+        """Return a dictionary of required datasets for this measure."""
+        pass
 
     @abstractmethod
-    def get_measure_data_from_artifact(self, *args, **kwargs) -> pd.DataFrame:
+    def get_measure_data_from_artifact(self, *args: Any, **kwargs: Any) -> pd.DataFrame:
         """Process artifact data into a format suitable for calculations."""
         pass
 
     @abstractmethod
-    def get_measure_data_from_sim(self, *args, **kwargs) -> pd.DataFrame:
+    def get_measure_data_from_sim(self, *args: Any, **kwargs: Any) -> pd.DataFrame:
         """Process raw simulation data into a format suitable for calculations."""
         pass
 
     @check_io(out=SingleNumericColumn)
-    def get_measure_data(self, source: DataSource, *args, **kwargs) -> pd.DataFrame:
+    def get_measure_data(self, source: DataSource, *args: Any, **kwargs: Any) -> pd.DataFrame:
         """Process data from the specified source into a format suitable for calculations."""
         if source == DataSource.SIM:
             return self.get_measure_data_from_sim(*args, **kwargs)
@@ -80,15 +92,6 @@ def artifact_datasets(self) -> dict[str, str]:
             "artifact_data": self.measure_key,
         }
 
-    @abstractmethod
-    def get_ratio_data_from_sim(
-        self,
-        numerator_data: pd.DataFrame,
-        denominator_data: pd.DataFrame,
-    ) -> pd.DataFrame:
-        """Process raw simulation data into a format suitable for calculations."""
-        pass
-
     @check_io(artifact_data=SingleNumericColumn, out=SingleNumericColumn)
     def get_measure_data_from_artifact(self, artifact_data: pd.DataFrame) -> pd.DataFrame:
         return artifact_data
@@ -103,7 +106,7 @@ def get_measure_data_from_ratio(self, ratio_data: pd.DataFrame) -> pd.DataFrame:
         )
 
     @check_io(out=SingleNumericColumn)
-    def get_measure_data_from_sim(self, *args, **kwargs) -> pd.DataFrame:
+    def get_measure_data_from_sim(self, *args: Any, **kwargs: Any) -> pd.DataFrame:
         """Process raw simulation data into a format suitable for calculations."""
         return self.get_measure_data_from_ratio(self.get_ratio_data_from_sim(*args, **kwargs))
 

src/vivarium_testing_utils/automated_validation/data_transformation/utils.py

@@ -1,10 +1,14 @@
 from __future__ import annotations
 
+from typing import Any, Callable, TypeVar
+
 import pandas as pd
 import pandera as pa
 
+F = TypeVar("F", bound=Callable[..., Any])
+
 
-def check_io(**model_dict):
+def check_io(**model_dict: type) -> Callable[[F], F]:
     """
     A wrapper for pa.check_io that automatically converts SchemaModels to schemas.