From e7360e185a2eb145a1fda6ea1cbb5406b15d1250 Mon Sep 17 00:00:00 2001 From: vahid-ahmadi Date: Wed, 16 Sep 2026 13:04:34 +0100 Subject: [PATCH 1/2] Prune failed Matching trials instead of scoring the training mean Both tuning handlers caught every exception and substituted the training mean, with no log and no counter, and that score went straight into the Optuna objective. A mean-predictor is not a neutral score - on a low-signal target it can beat a genuine matching fit on quantile loss - so a parameter set under which matching always failed could be selected as best and reported as the winning method. Both now log the exception and raise TrialPruned. The predict path keeps its NaN fill, which is the right behaviour there, but now reports the total number of unmatched records rather than leaving silent NaN blocks. Fixes #210 --- changelog.d/210.fixed.md | 1 + microimpute/models/matching.py | 42 ++++++++++++++++++++++++++-------- 2 files changed, 33 insertions(+), 10 deletions(-) create mode 100644 changelog.d/210.fixed.md diff --git a/changelog.d/210.fixed.md b/changelog.d/210.fixed.md new file mode 100644 index 00000000..279eeac4 --- /dev/null +++ b/changelog.d/210.fixed.md @@ -0,0 +1 @@ +Matching hyperparameter tuning now prunes a trial when matching fails, instead of silently scoring it as if it had predicted the training mean. The predict path reports how many records could not be matched. diff --git a/microimpute/models/matching.py b/microimpute/models/matching.py index 07952d56..5b6af9a5 100644 --- a/microimpute/models/matching.py +++ b/microimpute/models/matching.py @@ -241,6 +241,18 @@ def _predict_chunked( combined_results = pd.concat(all_results) combined_results = combined_results.loc[X_test_copy.index] + # A failed chunk leaves NaN blocks in the output. Report the total + # so a caller knows what share of the result is missing without + # having to check for it themselves. + n_failed = int(combined_results.isna().any(axis=1).sum()) + if n_failed: + self.logger.warning( + f"{n_failed} of {len(combined_results)} records " + f"({n_failed / len(combined_results):.1%}) could not be " + "matched and are NaN in the result." + ) + self.n_failed_records = n_failed + return self._process_matching_results( combined_results, X_test_copy, quantiles, return_probs ) @@ -639,11 +651,18 @@ def objective(trial: optuna.Trial) -> float: ) y_pred_chunks.append(fused0[var].values) y_val_chunks.append(chunk_y_val.values) - except Exception: - # If chunk fails, use mean of training data as prediction - mean_val = X_train_fold[var].mean() - y_pred_chunks.append(np.full(len(chunk_data), mean_val)) - y_val_chunks.append(chunk_y_val.values) + except Exception as e: + # Substituting the training mean here would + # score this trial as a mean-predictor, which + # can beat a genuine matching fit on a + # low-signal target. Prune instead, so a + # parameter set that cannot match is never + # selected as best. + self.logger.warning( + f"Matching failed for '{var}' on fold " + f"{fold_idx} chunk {i}: {e}. Pruning trial." + ) + raise optuna.TrialPruned() from e # Combine chunk results y_pred = np.concatenate(y_pred_chunks) @@ -660,11 +679,14 @@ def objective(trial: optuna.Trial) -> float: ) y_pred = fused0[var].values y_val_combined = y_val.values - except Exception: - # If matching fails, use mean of training data as prediction - mean_val = X_train_fold[var].mean() - y_pred = np.full(len(X_val_var), mean_val) - y_val_combined = y_val.values + except Exception as e: + # See above: score the trial on matching, or not at + # all. + self.logger.warning( + f"Matching failed for '{var}' on fold " + f"{fold_idx}: {e}. Pruning trial." + ) + raise optuna.TrialPruned() from e # Use appropriate metric based on variable type metric = variable_metrics[var] From 33501b6547945ecb37c4e8e273193f2988491cf2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mar=C3=ADa=20Juaristi?= <127882282+juaristi22@users.noreply.github.com> Date: Thu, 17 Sep 2026 11:34:00 +0100 Subject: [PATCH 2/2] Fix issues from review: track Matching prediction failures --- changelog.d/210.fixed.md | 2 +- .../cross-validation.md | 12 +- docs/imputation-benchmarking/preprocessing.md | 8 +- .../imputation-benchmarking/visualizations.md | 8 +- docs/models/imputer/implement-new-model.md | 4 +- docs/use_cases/index.md | 12 +- microimpute/models/matching.py | 30 ++-- tests/test_models/test_matching_failures.py | 164 ++++++++++++++++++ 8 files changed, 202 insertions(+), 38 deletions(-) create mode 100644 tests/test_models/test_matching_failures.py diff --git a/changelog.d/210.fixed.md b/changelog.d/210.fixed.md index 279eeac4..f465c4b7 100644 --- a/changelog.d/210.fixed.md +++ b/changelog.d/210.fixed.md @@ -1 +1 @@ -Matching hyperparameter tuning now prunes a trial when matching fails, instead of silently scoring it as if it had predicted the training mean. The predict path reports how many records could not be matched. +Matching hyperparameter tuning now prunes a trial when matching fails, instead of silently scoring it as if it had predicted the training mean, and reports when no trial succeeds. Predictions report how many records could not be matched and reset the failure count on every returned result, including small unchunked predictions. diff --git a/docs/imputation-benchmarking/cross-validation.md b/docs/imputation-benchmarking/cross-validation.md index 8da745ff..c1b6f93e 100644 --- a/docs/imputation-benchmarking/cross-validation.md +++ b/docs/imputation-benchmarking/cross-validation.md @@ -41,23 +41,23 @@ Returns a dictionary containing separate results for each metric type: ```python { "quantile_loss": { - "results": pd.DataFrame, # rows: ["train", "test"], cols: quantiles (mean across folds) + "results": pd.DataFrame, # rows: ["train", "test"], cols: quantiles (mean across folds) "results_std": pd.DataFrame, # rows: ["train", "test"], cols: quantiles (std across folds) "mean_train": float, "mean_test": float, "std_train": float, "std_test": float, - "variables": List[str] # numerical variables evaluated + "variables": List[str], # numerical variables evaluated }, "log_loss": { - "results": pd.DataFrame, # rows: ["train", "test"], cols: quantiles + "results": pd.DataFrame, # rows: ["train", "test"], cols: quantiles "results_std": pd.DataFrame, # rows: ["train", "test"], cols: quantiles (std across folds) "mean_train": float, "mean_test": float, "std_train": float, "std_test": float, - "variables": List[str] # categorical variables evaluated - } + "variables": List[str], # categorical variables evaluated + }, } ``` @@ -77,7 +77,7 @@ results = cross_validate_model( data=diabetes_df, predictors=["age", "sex", "bmi", "bp"], imputed_variables=["s1", "s4"], - n_splits=5 + n_splits=5, ) # Check performance for numerical variables diff --git a/docs/imputation-benchmarking/preprocessing.md b/docs/imputation-benchmarking/preprocessing.md index 736d73ed..05693943 100644 --- a/docs/imputation-benchmarking/preprocessing.md +++ b/docs/imputation-benchmarking/preprocessing.md @@ -117,10 +117,10 @@ result = autoimpute( predictors=["age", "education"], imputed_variables=["income", "wealth"], preprocessing={ - "income": "log", # Log transform (positive values only) - "wealth": "asinh", # Asinh transform (handles zeros/negatives) - "age": "normalize" # Z-score normalization - } + "income": "log", # Log transform (positive values only) + "wealth": "asinh", # Asinh transform (handles zeros/negatives) + "age": "normalize", # Z-score normalization + }, ) ``` diff --git a/docs/imputation-benchmarking/visualizations.md b/docs/imputation-benchmarking/visualizations.md index 21dde021..fe8f87b8 100644 --- a/docs/imputation-benchmarking/visualizations.md +++ b/docs/imputation-benchmarking/visualizations.md @@ -83,11 +83,7 @@ comparison_viz = method_comparison_results( ) # Generate plot -fig = comparison_viz.plot( - title="Method comparison", - show_mean=True, - plot_type="bar" -) +fig = comparison_viz.plot(title="Method comparison", show_mean=True, plot_type="bar") fig.show() # Get summary statistics @@ -165,7 +161,7 @@ perf_viz = model_performance_results( results=cv_results, model_name="QRF", method_name="Cross-validation", - metric="quantile_loss" + metric="quantile_loss", ) fig = perf_viz.plot(title="QRF performance") diff --git a/docs/models/imputer/implement-new-model.md b/docs/models/imputer/implement-new-model.md index edf21254..794d6193 100644 --- a/docs/models/imputer/implement-new-model.md +++ b/docs/models/imputer/implement-new-model.md @@ -73,9 +73,7 @@ class NewModelResults(ImputerResults): except Exception as e: self.logger.error(f"Error during Model prediction: {str(e)}") - raise RuntimeError( - f"Failed to predict with Model: {str(e)}" - ) from e + raise RuntimeError(f"Failed to predict with Model: {str(e)}") from e ``` ## Implementing the main model class diff --git a/docs/use_cases/index.md b/docs/use_cases/index.md index 32c7af96..2c9ff07c 100644 --- a/docs/use_cases/index.md +++ b/docs/use_cases/index.md @@ -23,7 +23,7 @@ Before imputation, make sure both datasets have compatible variables. Identify c ```python # Identify common variables -common_variables = ['age', 'income', 'education', 'marital_status', 'region'] +common_variables = ["age", "income", "education", "marital_status", "region"] # Ensure variable formats match (example: education coding) education_mapping = { @@ -31,19 +31,19 @@ education_mapping = { 2: "high_school", 3: "some_college", 4: "bachelor", - 5: "graduate" + 5: "graduate", } # Apply standardization to both datasets for dataset in [scf_data, cps_data]: - dataset['education'] = dataset['education'].map(education_mapping) + dataset["education"] = dataset["education"].map(education_mapping) # Convert income to same units (thousands) - if 'income' in dataset.columns: - dataset['income'] = dataset['income'] / 1000 + if "income" in dataset.columns: + dataset["income"] = dataset["income"] / 1000 # Identify target variable in donor dataset -target_variable = ['networth'] +target_variable = ["networth"] ``` ## Performing imputation diff --git a/microimpute/models/matching.py b/microimpute/models/matching.py index 5b6af9a5..640178e0 100644 --- a/microimpute/models/matching.py +++ b/microimpute/models/matching.py @@ -74,6 +74,7 @@ def __init__( self.categorical_targets = categorical_targets or {} self.boolean_targets = boolean_targets or {} self.dummy_processor = dummy_processor + self.n_failed_records = 0 @validate_call(config=VALIDATE_CONFIG) def _predict( @@ -241,18 +242,6 @@ def _predict_chunked( combined_results = pd.concat(all_results) combined_results = combined_results.loc[X_test_copy.index] - # A failed chunk leaves NaN blocks in the output. Report the total - # so a caller knows what share of the result is missing without - # having to check for it themselves. - n_failed = int(combined_results.isna().any(axis=1).sum()) - if n_failed: - self.logger.warning( - f"{n_failed} of {len(combined_results)} records " - f"({n_failed / len(combined_results):.1%}) could not be " - "matched and are NaN in the result." - ) - self.n_failed_records = n_failed - return self._process_matching_results( combined_results, X_test_copy, quantiles, return_probs ) @@ -337,6 +326,18 @@ def _process_matching_results( ) raise RuntimeError("Failed to process matching results") from convert_error + # Both single-call and chunked predictions replace the previous count. + # Only missing target values represent unmatched output records. + self.n_failed_records = int( + fused0[self.imputed_variables].isna().any(axis=1).sum() + ) + if self.n_failed_records: + self.logger.warning( + f"{self.n_failed_records} of {len(fused0)} records " + f"({self.n_failed_records / len(fused0):.1%}) could not be " + "matched and are NaN in the result." + ) + # Create output dictionary with results imputations: Dict[float, pd.DataFrame] = {} prob_results = {} if return_probs else None @@ -731,6 +732,11 @@ def objective(trial: optuna.Trial) -> float: study.optimize(objective, n_trials=n_trials) + if not any( + trial.state == optuna.trial.TrialState.COMPLETE for trial in study.trials + ): + raise ValueError("No matching hyperparameter trial succeeded") + best_value = study.best_value self.logger.info( f"Matching - Lowest average normalized quantile loss ({n_cv_folds}-fold CV): {best_value}" diff --git a/tests/test_models/test_matching_failures.py b/tests/test_models/test_matching_failures.py new file mode 100644 index 00000000..1a212ad9 --- /dev/null +++ b/tests/test_models/test_matching_failures.py @@ -0,0 +1,164 @@ +"""Failure handling with a deterministic custom backend and real Optuna studies.""" + +import importlib.util +import sys +import types +from pathlib import Path + +import numpy as np +import optuna +import pandas as pd +import pytest + + +@pytest.fixture +def matching_class(monkeypatch): + """Load Matching without installing or globally replacing its optional R adapter.""" + try: + from microimpute.models.matching import Matching + except ModuleNotFoundError as error: + if not error.name.startswith("rpy2"): + raise + import microimpute.models + + adapter = types.ModuleType("microimpute.utils.statmatch_hotdeck") + + def unavailable_adapter(*args, **kwargs): + raise AssertionError("These tests must use their custom matching backend") + + adapter.nnd_hotdeck_using_rpy2 = unavailable_adapter + path = Path(microimpute.models.__file__).with_name("matching.py") + spec = importlib.util.spec_from_file_location("_matching_failure_tests", path) + module = importlib.util.module_from_spec(spec) + with monkeypatch.context() as scoped: + scoped.setitem(sys.modules, adapter.__name__, adapter) + scoped.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + return module.Matching + return Matching + + +def donor_data(size=30): + x = np.arange(size, dtype=float) + return pd.DataFrame({"x": x, "y": 2 * x + 1}) + + +def exact_backend(receiver, donor, matching_variables, z_variables, **kwargs): + """Predict the known linear relation, making the successful-trial score exact.""" + result = receiver.copy() + for variable in z_variables: + result[variable] = 2 * receiver["x"].to_numpy() + 1 + return result, result.copy() + + +@pytest.fixture +def studies(monkeypatch): + """Queue one failing candidate followed by a successful candidate deterministically.""" + created = [] + create_study = optuna.create_study + + def tracked_study(**kwargs): + study = create_study(**kwargs) + study.enqueue_trial({"dist_fun": "Manhattan", "k": 1}) + study.enqueue_trial({"dist_fun": "Euclidean", "k": 1}) + created.append(study) + return study + + monkeypatch.setattr(optuna, "create_study", tracked_study) + return created + + +@pytest.mark.parametrize("size", [30, 3003]) +def test_failed_tuning_candidate_cannot_win(matching_class, studies, size): + """Prune failures, including a second failed chunk after the first succeeds.""" + successful_manhattan_chunks = [] + + def backend(**kwargs): + if kwargs["dist_fun"] == "Manhattan": + if size == 30 or len(kwargs["receiver"]) == 1: + raise RuntimeError("candidate cannot match") + successful_manhattan_chunks.append(len(kwargs["receiver"])) + return exact_backend(**kwargs) + + data = donor_data(size) + fitted, params = matching_class(backend).fit( + data, ["x"], ["y"], tune_hyperparameters=True + ) + trials = studies[0].trials + assert trials[0].state == optuna.trial.TrialState.PRUNED + assert trials[1].state == optuna.trial.TrialState.COMPLETE + assert trials[1].value == pytest.approx(0.0) + assert params["dist_fun"] != "Manhattan" + if size > 30: + assert successful_manhattan_chunks and set(successful_manhattan_chunks) == { + 1000 + } + prediction = fitted.predict(data[["x"]].iloc[:3], quantiles=[0.5])[0.5] + np.testing.assert_array_equal(prediction.y, data.y.iloc[:3]) + + +@pytest.mark.parametrize("size", [30, 3003]) +def test_all_failed_trials_raise_without_a_model(matching_class, studies, size): + """An all-pruned study must never fall back to a fitted mean predictor.""" + + def backend(**kwargs): + raise RuntimeError("no valid donor match") + + with pytest.raises(ValueError, match="No matching hyperparameter trial succeeded"): + matching_class(backend).fit( + donor_data(size), ["x"], ["y"], tune_hyperparameters=True + ) + assert studies[0].trials + assert all( + trial.state == optuna.trial.TrialState.PRUNED for trial in studies[0].trials + ) + + +def test_failure_count_available_before_and_after_small_prediction(matching_class): + """A new fitted model and its successful unchunked result report zero failures.""" + fitted = matching_class(exact_backend).fit(donor_data(), ["x"], ["y"]) + assert fitted.n_failed_records == 0 + prediction = fitted.predict(donor_data()[["x"]].iloc[:3], quantiles=[0.5])[0.5] + assert fitted.n_failed_records == 0 + assert not prediction.isna().any().any() + + +@pytest.mark.parametrize("next_size", [3, 2001]) +def test_partial_prediction_preserves_rows_and_resets_count(matching_class, next_size): + """A later successful prediction replaces the previous failure count.""" + + def backend(**kwargs): + if len(kwargs["receiver"]) == 1: + raise RuntimeError("last chunk cannot match") + return exact_backend(**kwargs) + + fitted = matching_class(backend).fit(donor_data(), ["x"], ["y"]) + receiver = pd.DataFrame( + {"x": np.arange(2001, dtype=float)}, index=np.arange(10000, 12001) + ) + prediction = fitted.predict(receiver, quantiles=[0.5])[0.5] + assert prediction.index.equals(receiver.index) + np.testing.assert_array_equal(prediction.y.iloc[:-1], 2 * receiver.x.iloc[:-1] + 1) + assert pd.isna(prediction.y.iloc[-1]) + assert fitted.n_failed_records == 1 + + fitted.matching_hotdeck = exact_backend + prediction = fitted.predict(receiver.iloc[:next_size], quantiles=[0.5])[0.5] + assert not prediction.isna().any().any() + assert fitted.n_failed_records == 0 + + +def test_small_prediction_reports_missing_targets_only(matching_class, caplog): + """Count a partially missing single-call result, not unrelated backend columns.""" + + def backend(**kwargs): + result, _ = exact_backend(**kwargs) + result["unused"] = np.nan + result.iloc[-1, result.columns.get_loc("y")] = np.nan + return result, result.copy() + + fitted = matching_class(backend).fit(donor_data(), ["x"], ["y"]) + prediction = fitted.predict(donor_data()[["x"]].iloc[:3], quantiles=[0.5])[0.5] + assert prediction.y.isna().sum() == 1 + assert fitted.n_failed_records == 1 + assert "1 of 3 records (33.3%) could not be matched" in caplog.text