diff --git a/README.md b/README.md index 2f5ad00d..f72bdfd0 100644 --- a/README.md +++ b/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. diff --git a/src/rtichoke/calibration/calibration.py b/src/rtichoke/calibration/calibration.py index 567c05c4..2efbf57d 100644 --- a/src/rtichoke/calibration/calibration.py +++ b/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") diff --git a/tests/test_calibration.py b/tests/test_calibration.py index 4e79687f..f17a2b91 100644 --- a/tests/test_calibration.py +++ b/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"} diff --git a/tests/test_calibration_times.py b/tests/test_calibration_times.py index b6570c95..8e73c63a 100644 --- a/tests/test_calibration_times.py +++ b/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"}