Prune failed Matching trials instead of scoring the training mean

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, 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.

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

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
+    },
 )
 ```
 

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")

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

docs/use_cases/index.md

@@ -23,27 +23,27 @@ 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 = {
     1: "less_than_hs",
     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

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(
@@ -325,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
@@ -639,11 +652,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 +680,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]
@@ -709,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}"

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