Fill missing Ensemble Distribution Parameters (#89)

* use singledispatch to simplify formatters

* [COPILOT] Add basic formatting tests

* add some tests for missing data

* make good tests

* cleanup

* remove added tests

* add unsupported tests

* formatting

* add cl

* another change

* format

Author: patricktnast

CHANGELOG.rst

@@ -1,3 +1,8 @@
+**2.2.0 - 10/02/25**
+
+ - Allow user to pass incomplete parameter sets to EnsembleDistribution (filled with zeros)
+ - Backfill some unit tests
+
 **2.1.6 - 08/01/25**
 
  - Use vivarium_dependencies for common setup constraints

src/risk_distributions/formatting.py

@@ -1,3 +1,4 @@
+from functools import singledispatch
 from typing import Any, TypeVar
 
 import numpy as np
@@ -42,22 +43,13 @@ def cast_to_series(mean: Parameter, sd: Parameter) -> tuple[pd.Series, pd.Series
     return mean, sd
 
 
+@singledispatch
 def format_data(data: Parameters, required_columns: list[Any], measure: str) -> pd.DataFrame:
     """Formats parameter data into a dataframe."""
-    if isinstance(data, np.ndarray):
-        data = format_array(data, required_columns, measure)
-    elif isinstance(data, pd.Series):
-        data = format_series(data, required_columns, measure)
-    elif isinstance(data, pd.DataFrame):
-        data = format_data_frame(data, required_columns, measure)
-    elif isinstance(data, (list, tuple)):
-        data = format_list_like(data, required_columns, measure)
-    elif isinstance(data, dict):
-        data = format_dict(data, required_columns, measure)
-
-    return data
+    raise TypeError(f"Unsupported data type {type(data)} for {measure}")
 
 
+@format_data.register
 def format_array(data: np.ndarray, required_columns: list[Any], measure: str) -> pd.DataFrame:
     """Transforms 1d and 2d arrays into dataframes with columns for the
     parameters and (possibly) rows for each parameter variation."""
@@ -111,6 +103,7 @@ def format_array(data: np.ndarray, required_columns: list[Any], measure: str) ->
     return data
 
 
+@format_data.register
 def format_series(data: pd.Series, required_columns: list[Any], measure: str) -> pd.DataFrame:
     """Transforms series data into dataframes with columns for the
     parameters and (possibly) rows for each parameter variation."""
@@ -133,6 +126,7 @@ def format_series(data: pd.Series, required_columns: list[Any], measure: str) ->
     return data
 
 
+@format_data.register
 def format_data_frame(
     data: pd.DataFrame, required_columns: list[Any], measure: str
 ) -> pd.DataFrame:
@@ -154,6 +148,8 @@ def format_data_frame(
     return data
 
 
+@format_data.register(list)
+@format_data.register(tuple)
 def format_list_like(
     data: list | tuple, required_columns: list[Any], measure: str
 ) -> pd.DataFrame:
@@ -163,9 +159,8 @@ def format_list_like(
     return format_array(data, required_columns, measure)
 
 
-def format_dict(
-    data: dict[str, Parameter], required_columns: list[Any], measure: str
-) -> pd.DataFrame:
+@format_data.register
+def format_dict(data: dict, required_columns: list[Any], measure: str) -> pd.DataFrame:
     """Transform dictionaries with scalar or list-like values into dataframes
     with columns for the parameters and (possibly) rows for each parameter
     variation."""

src/risk_distributions/risk_distributions.py

@@ -1,3 +1,4 @@
+import copy
 import warnings
 from collections.abc import Callable
 
@@ -566,7 +567,10 @@ def get_parameters(
         mean: Parameter = None,
         sd: Parameter = None,
     ) -> tuple[pd.DataFrame, dict[str, pd.DataFrame]]:
-        weights = format_data(weights, list(cls._distribution_map.keys()), "weights")
+        expected_columns = list(cls._distribution_map.keys())
+
+        weights = cls.fill_missing_weights(weights, expected_columns)
+        weights = format_data(weights, expected_columns, "weights")
 
         params = {}
         for name, dist in cls._distribution_map.items():
@@ -591,6 +595,27 @@ def get_parameters(
 
         return weights, params
 
+    @staticmethod
+    def fill_missing_weights(weights: Parameters, expected_columns) -> Parameters:
+        weights = copy.deepcopy(weights)
+
+        # Get existing keys/columns/index based on weights type
+        if isinstance(weights, dict):
+            column_names = set(weights.keys())
+        elif isinstance(weights, pd.DataFrame):
+            column_names = set(weights.columns)
+        elif isinstance(weights, pd.Series):
+            column_names = set(weights.index)
+        else:
+            column_names = None  # For list, tuple, np.array, we can't fill missing columns
+
+        # Add missing columns with 0.0 value
+        if column_names and column_names < set(expected_columns):
+            for col in expected_columns:
+                if col not in column_names:
+                    weights[col] = 0.0
+        return weights
+
     def pdf(self, x: pd.Series | np.ndarray | float | int) -> pd.Series | np.ndarray | float:
         single_val = isinstance(x, (float, int))
         values_only = isinstance(x, np.ndarray)

tests/test_ensemble_distribution.py

@@ -9,44 +9,107 @@
 from risk_distributions.risk_distributions import EnsembleDistribution
 
 weights_base = {
-    "betasr": 1 / 12,
-    "exp": 1 / 12,
-    "gamma": 1 / 12,
-    "gumbel": 1 / 12,
-    "invgamma": 1 / 12,
-    "invweibull": 1 / 12,
-    "llogis": 1 / 12,
-    "lnorm": 1 / 12,
-    "mgamma": 1 / 12,
-    "mgumbel": 1 / 12,
-    "norm": 1 / 12,
-    "weibull": 1 / 12,
+    "betasr": 1,
+    "exp": 2,
+    "gamma": 3,
+    "gumbel": 5,
+    "invgamma": 7,
+    "invweibull": 11,
+    "llogis": 13,
+    "lnorm": 17,
+    "mgamma": 19,
+    "mgumbel": 23,
+    "norm": 29,
+    "weibull": 31,
 }
 
-weights_df = pd.DataFrame({k: [v] for k, v in weights_base.items()})
+weights_base_missing = copy.deepcopy(weights_base)
+del weights_base_missing["exp"]
+
+
+def normalize_weights(weights: dict[str, float]) -> dict[str, float]:
+    weights = copy.deepcopy(weights)
+    total = sum(weights.values())
+    for k in weights:
+        weights[k] = weights[k] / total
+    return weights
+
+
+@pytest.fixture
+def expected_weights() -> pd.DataFrame:
+    return pd.DataFrame({k: [v] for k, v in normalize_weights(weights_base).items()})
+
+
+@pytest.fixture
+def expected_weights_missing() -> pd.DataFrame:
+    data = pd.DataFrame({k: [v] for k, v in normalize_weights(weights_base_missing).items()})
+    data["exp"] = 0.0
+    return data
 
 
 @pytest.mark.parametrize(
     "weights",
     [
         weights_base,
+        normalize_weights(weights_base),
         {k: [v] for k, v in weights_base.items()},
         pd.Series(weights_base),
         pd.Series(weights_base).reset_index(drop=True),
-        weights_df,
+        pd.DataFrame({k: [v] for k, v in weights_base.items()}),
         list(weights_base.values()),
         tuple(weights_base.values()),
         np.array(list(weights_base.values())),  # Column Vector
         np.array([list(weights_base.values())]),  # Row Vector
         np.array([list(weights_base.values())]).T,
     ],
 )
-def test_weight_formats(weights: Parameters) -> None:
+def test_weight_formats(weights: Parameters, expected_weights: pd.DataFrame) -> None:
     weights_original = copy.deepcopy(weights)
     dist = EnsembleDistribution(
         weights,
         mean=1,
         sd=1,
     )
     assert_equal(weights_original, weights)
-    pd.testing.assert_frame_equal(dist.weights, pd.DataFrame(weights_df))
+    pd.testing.assert_frame_equal(dist.weights, expected_weights)
+
+
+@pytest.mark.parametrize(
+    "weights",
+    [
+        weights_base_missing,
+        normalize_weights(weights_base_missing),
+        {k: [v] for k, v in weights_base_missing.items()},
+        pd.Series(weights_base_missing),
+        pd.DataFrame({k: [v] for k, v in weights_base_missing.items()}),
+    ],
+)
+def test_missing_weights(weights: Parameters, expected_weights_missing: pd.DataFrame) -> None:
+    weights_original = copy.deepcopy(weights)
+    dist = EnsembleDistribution(
+        weights,
+        mean=1,
+        sd=1,
+    )
+    assert_equal(weights_original, weights)
+    pd.testing.assert_frame_equal(dist.weights, expected_weights_missing)
+
+
+@pytest.mark.parametrize(
+    "weights",
+    [
+        pd.Series(weights_base_missing).reset_index(drop=True),
+        list(weights_base_missing.values()),
+        tuple(weights_base_missing.values()),
+        np.array(list(weights_base_missing.values())),  # Column Vector
+        np.array([list(weights_base_missing.values())]),  # Row Vector
+        np.array([list(weights_base_missing.values())]).T,
+    ],
+)
+def test_missing_weights_invalid(weights: Parameters) -> None:
+    with pytest.raises(ValueError):
+        EnsembleDistribution(
+            weights,
+            mean=1,
+            sd=1,
+        )

tests/test_formatting.py

@@ -121,6 +121,13 @@ def test_format_data_frame(data_columns, required_columns, match):
         format_data_frame(data, required_columns, measure="test")
 
 
+@pytest.mark.parametrize("data", ["string", {1, 2, 3}, None])
+def test_format_data_unsupported_types(data):
+    """Test format_data with unsupported data types."""
+    with pytest.raises(TypeError, match="Unsupported data type"):
+        format_data(data, ["param1"], "test")
+
+
 @pytest.mark.parametrize(
     "data, required_columns, expected",
     [