Fix calibration multi-population input contract

README.md

@@ -31,7 +31,7 @@ pip install rtichoke
 To use `rtichoke`, you'll usually need two main inputs:
 
 * `probs`: A dictionary containing model-predicted probabilities.
-* `reals`: A dictionary containing the observed outcomes.
+* `reals`: Observed outcomes, provided either as one array or as a dictionary keyed by population.
 
 Here's a quick example of creating a ROC curve for a single model:
 
@@ -54,6 +54,30 @@ fig = rk.create_roc_curve(
 fig.show()
 ```
 
+### Compare populations
+
+When predictions and outcomes are both dictionaries with the same keys, rtichoke pairs them population-by-population. The populations do **not** need to have the same sample size.
+
+```python
+probs = {
+    "Train": np.array([0.10, 0.90, 0.20, 0.80, 0.30, 0.70]),
+    "Test": np.array([0.15, 0.85, 0.25, 0.75]),
+}
+reals = {
+    "Train": np.array([0, 1, 0, 1, 0, 1]),
+    "Test": np.array([0, 1, 0, 0]),
+}
+
+fig = rk.create_calibration_curve(
+    probs=probs,
+    reals=reals,
+)
+
+fig.show()
+```
+
+Here, `Train` contains six observations and `Test` contains four. Each probability vector only needs to match the outcome vector for its own population.
+
 ## Key Features
 
 * **Simple API**: Create complex visualizations with a small amount of code.

src/rtichoke/calibration/calibration.py

@@ -470,17 +470,39 @@ def _make_deciles_dat_binary(
     n_bins: int = 10,
 ) -> pl.DataFrame:
     if isinstance(reals, dict):
-        reference_groups_keys = list(reals.keys())
-        y_list = [
-            np.asarray(reals[str(reference_group)]).ravel()
-            for reference_group in reference_groups_keys
-        ]
-        lengths = np.array([len(y) for y in y_list], dtype=np.int64)
-        offsets = np.concatenate([np.array([0], dtype=np.int64), np.cumsum(lengths)])
-        n_total = int(offsets[-1])
-
         frames: list[pl.DataFrame] = []
-        for model, p_all in probs.items():
+
+        if probs.keys() == reals.keys():
+            for population in reals:
+                p = np.asarray(probs[population]).ravel()
+                y = np.asarray(reals[population]).ravel()
+                if p.shape[0] != y.shape[0]:
+                    raise ValueError(
+                        f"Length mismatch for population '{population}': "
+                        f"probs has length {p.shape[0]} but reals has length {y.shape[0]}."
+                    )
+                frames.append(
+                    pl.DataFrame(
+                        {
+                            "reference_group": population,
+                            "model": population,
+                            "prob": p.astype(float, copy=False),
+                            "real": y.astype(float, copy=False),
+                        }
+                    )
+                )
+        elif len(probs) == 1:
+            reference_groups_keys = list(reals.keys())
+            y_list = [
+                np.asarray(reals[str(reference_group)]).ravel()
+                for reference_group in reference_groups_keys
+            ]
+            lengths = np.array([len(y) for y in y_list], dtype=np.int64)
+            offsets = np.concatenate(
+                [np.array([0], dtype=np.int64), np.cumsum(lengths)]
+            )
+            n_total = int(offsets[-1])
+            model, p_all = next(iter(probs.items()))
             p_all = np.asarray(p_all).ravel()
             if p_all.shape[0] != n_total:
                 raise ValueError(
@@ -491,7 +513,6 @@ def _make_deciles_dat_binary(
             for i, pop in enumerate(reference_groups_keys):
                 start = int(offsets[i])
                 end = int(offsets[i + 1])
-
                 frames.append(
                     pl.DataFrame(
                         {
@@ -502,6 +523,10 @@ def _make_deciles_dat_binary(
                         }
                     )
                 )
+        else:
+            raise ValueError(
+                "When probs and reals are dictionaries, their population keys must match."
+            )
 
         df = pl.concat(frames, how="vertical")
 

tests/test_calibration.py

@@ -26,3 +26,20 @@ def test_create_calibration_curve_smooth_single_point():
     # Check histogram data
     histogram = fig.data[2]
     assert histogram.type == "bar"
+
+
+def test_create_calibration_curve_multiple_populations_unequal_sizes():
+    probs = {
+        "Train": np.array([0.1, 0.9, 0.2, 0.8, 0.3, 0.7]),
+        "Test": np.array([0.2, 0.8, 0.3, 0.7]),
+    }
+    reals = {
+        "Train": np.array([0, 1, 0, 1, 0, 1]),
+        "Test": np.array([0, 1, 0, 0]),
+    }
+
+    for calibration_type in ("discrete", "smooth"):
+        fig = create_calibration_curve(
+            probs, reals, calibration_type=calibration_type
+        )
+        assert {trace.name for trace in fig.data if trace.name} >= {"Train", "Test"}

tests/test_calibration_times.py

@@ -22,3 +22,34 @@ def test_create_calibration_curve_times():
     assert fig is not None
     assert len(fig.data) > 0
     assert len(fig.layout.sliders) > 0
+
+
+def test_create_calibration_curve_times_unequal_size_populations():
+    probs = {
+        "Train": np.array([0.1, 0.9, 0.2, 0.8, 0.3, 0.7]),
+        "Test": np.array([0.2, 0.8, 0.3, 0.7]),
+    }
+    reals = {
+        "Train": np.array([0, 1, 0, 1, 0, 1]),
+        "Test": np.array([0, 1, 0, 0]),
+    }
+    times = {
+        "Train": np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]),
+        "Test": np.array([1.0, 2.0, 3.0, 4.0]),
+    }
+    heuristics_sets = [
+        {
+            "censoring_heuristic": "excluded",
+            "competing_heuristic": "adjusted_as_negative",
+        }
+    ]
+
+    fig = create_calibration_curve_times(
+        probs,
+        reals,
+        times,
+        fixed_time_horizons=[3.0, 6.0],
+        heuristics_sets=heuristics_sets,
+    )
+
+    assert {trace.name for trace in fig.data if trace.name} >= {"Train", "Test"}