From 79ac62b4ad134557b2988527d2473b5c3379fb57 Mon Sep 17 00:00:00 2001
From: vahid-ahmadi <va.vahidahmadi@gmail.com>
Date: Wed, 16 Sep 2026 13:04:34 +0100
Subject: [PATCH 1/5] 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 5a4f8d6e8355b6332f6737bfd1be15c69afd59ef 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/5] 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

From 6a94c6dea18b213fb251e79856209e1203137dc9 Mon Sep 17 00:00:00 2001
From: vahid-ahmadi <va.vahidahmadi@gmail.com>
Date: Mon, 21 Sep 2026 10:36:33 +0100
Subject: [PATCH 3/5] Reset the failure counter, report the cause, expose it on
 the result

Per review, three follow-ups.

n_failed_records kept the previous successful call's value when a
prediction raised before reaching _process_matching_results, so a caller
reading it after an exception got a stale number. It is now reset on
entry to _predict, and documented - including that Matching runs
single-threaded, so the attribute is safe in practice but would race
under concurrent calls on one fitted object.

The count is also mirrored onto result.attrs['n_failed_records'], so a
caller does not have to reach into the fitted model for it.

An all-pruned study reported only that nothing succeeded. It now carries
the most recent underlying failure, which is what a user needs when
matching fails structurally - a bad dtype or a missing R package - and
the cause was otherwise reachable only through __cause__.

Also drops the docs reformatting, which belongs to #218.
---
 .../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                | 36 +++++++++++++++++--
 6 files changed, 59 insertions(+), 21 deletions(-)

diff --git a/docs/imputation-benchmarking/cross-validation.md b/docs/imputation-benchmarking/cross-validation.md
index c1b6f93e..8da745ff 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 05693943..736d73ed 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 fe8f87b8..21dde021 100644
--- a/docs/imputation-benchmarking/visualizations.md
+++ b/docs/imputation-benchmarking/visualizations.md
@@ -83,7 +83,11 @@ 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
@@ -161,7 +165,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 794d6193..edf21254 100644
--- a/docs/models/imputer/implement-new-model.md
+++ b/docs/models/imputer/implement-new-model.md
@@ -73,7 +73,9 @@ 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 2c9ff07c..32c7af96 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 640178e0..a7a79a82 100644
--- a/microimpute/models/matching.py
+++ b/microimpute/models/matching.py
@@ -94,11 +94,24 @@ def _predict(
             Dictionary mapping quantiles to imputed values.
             If return_probs=True, includes 'probabilities' key with one-hot encodings.
 
+        Side effects:
+            Sets ``self.n_failed_records`` to the number of recipient records
+            that could not be matched and are NaN in the result, and mirrors it
+            on ``result.attrs["n_failed_records"]``. It is reset to 0 on entry,
+            so it always describes the most recent call. Matching runs
+            single-threaded (``autoimpute`` forces ``n_jobs=1`` when a Matching
+            model is present), so concurrent calls on one fitted object would
+            race on it.
+
         Raises:
             ValueError: If model is not properly set up or
                 input data is invalid.
             RuntimeError: If matching or prediction fails.
         """
+        # Reset before any work: a prediction that raises must not leave the
+        # previous successful call's count readable as if it described this one.
+        self.n_failed_records = 0
+
         try:
             self.logger.info(f"Performing matching for {len(X_test)} recipient records")
 
@@ -376,6 +389,12 @@ def _process_matching_results(
                 if return_probs and prob_results:
                     imputations["probabilities"] = prob_results
 
+                # Mirror the unmatched count onto each frame, so a caller can
+                # see it without reaching into the fitted model.
+                for frame in imputations.values():
+                    if isinstance(frame, pd.DataFrame):
+                        frame.attrs["n_failed_records"] = self.n_failed_records
+
                 return imputations
             else:
                 # If no quantiles specified, use a default one
@@ -585,7 +604,13 @@ def _tune_hyperparameters(
             f"Tuning Matching hyperparameters with {n_cv_folds}-fold CV and {n_trials} trials"
         )
 
+        # Keeps the most recent underlying failure so an all-pruned study can
+        # report why, rather than only that nothing succeeded.
+        last_trial_error: Optional[BaseException] = None
+
         def objective(trial: optuna.Trial) -> float:
+            nonlocal last_trial_error
+
             params = {
                 "dist_fun": trial.suggest_categorical(
                     "dist_fun",
@@ -663,7 +688,8 @@ def objective(trial: optuna.Trial) -> float:
                                     f"Matching failed for '{var}' on fold "
                                     f"{fold_idx} chunk {i}: {e}. Pruning trial."
                                 )
-                                raise optuna.TrialPruned() from e
+                                last_trial_error = e
+                            raise optuna.TrialPruned() from e
 
                         # Combine chunk results
                         y_pred = np.concatenate(y_pred_chunks)
@@ -687,6 +713,7 @@ def objective(trial: optuna.Trial) -> float:
                                 f"Matching failed for '{var}' on fold "
                                 f"{fold_idx}: {e}. Pruning trial."
                             )
+                            last_trial_error = e
                             raise optuna.TrialPruned() from e
 
                     # Use appropriate metric based on variable type
@@ -735,7 +762,12 @@ def objective(trial: optuna.Trial) -> float:
         if not any(
             trial.state == optuna.trial.TrialState.COMPLETE for trial in study.trials
         ):
-            raise ValueError("No matching hyperparameter trial succeeded")
+            raise ValueError(
+                "No matching hyperparameter trial succeeded. Last error: "
+                f"{last_trial_error}"
+                if last_trial_error
+                else "No matching hyperparameter trial succeeded"
+            )
 
         best_value = study.best_value
         self.logger.info(

From 09ba11b5ab836ffe6f967df87527db97c7af9011 Mon Sep 17 00:00:00 2001
From: vahid-ahmadi <va.vahidahmadi@gmail.com>
Date: Mon, 21 Sep 2026 10:38:28 +0100
Subject: [PATCH 4/5] Fix indentation that moved a raise out of its except
 block

The previous commit de-indented 'raise optuna.TrialPruned() from e' out
of 'except Exception as e', so 'e' was unbound and the trial failed with
UnboundLocalError instead of pruning - which the matching failure tests
caught.
---
 microimpute/models/matching.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/microimpute/models/matching.py b/microimpute/models/matching.py
index a7a79a82..3bbd3894 100644
--- a/microimpute/models/matching.py
+++ b/microimpute/models/matching.py
@@ -689,7 +689,7 @@ def objective(trial: optuna.Trial) -> float:
                                     f"{fold_idx} chunk {i}: {e}. Pruning trial."
                                 )
                                 last_trial_error = e
-                            raise optuna.TrialPruned() from e
+                                raise optuna.TrialPruned() from e
 
                         # Combine chunk results
                         y_pred = np.concatenate(y_pred_chunks)

From 62d03ce4d3f475b41491136dd3c3154464fb176e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mar=C3=ADa=20Juaristi?=
 <127882282+juaristi22@users.noreply.github.com>
Date: Mon, 21 Sep 2026 11:54:46 +0100
Subject: [PATCH 5/5] Fix issues from review: restore lint compatibility and
 preserve matching failure metadata

---
 changelog.d/215-review.fixed.md             |  1 +
 microimpute/models/matching.py              |  1 +
 tests/test_models/test_matching_failures.py | 37 +++++++++++++++++++++
 3 files changed, 39 insertions(+)
 create mode 100644 changelog.d/215-review.fixed.md

diff --git a/changelog.d/215-review.fixed.md b/changelog.d/215-review.fixed.md
new file mode 100644
index 00000000..45e4520d
--- /dev/null
+++ b/changelog.d/215-review.fixed.md
@@ -0,0 +1 @@
+Include the unmatched-record count in the default Matching prediction frame's metadata, consistently with explicit-quantile results.
diff --git a/microimpute/models/matching.py b/microimpute/models/matching.py
index 3bbd3894..fe366a2b 100644
--- a/microimpute/models/matching.py
+++ b/microimpute/models/matching.py
@@ -420,6 +420,7 @@ def _process_matching_results(
                             prob_results[variable] = prob_df
 
                 imputations[q_default] = imputed_df
+                imputed_df.attrs["n_failed_records"] = self.n_failed_records
 
                 # Add probabilities to results if requested
                 if return_probs and prob_results:
diff --git a/tests/test_models/test_matching_failures.py b/tests/test_models/test_matching_failures.py
index 1a212ad9..9985f8c0 100644
--- a/tests/test_models/test_matching_failures.py
+++ b/tests/test_models/test_matching_failures.py
@@ -162,3 +162,40 @@ def backend(**kwargs):
     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
+
+
+@pytest.mark.parametrize("quantiles", [None, [0.25, 0.75]])
+@pytest.mark.parametrize("failed_rows", [0, 1])
+def test_prediction_frames_include_failure_metadata(
+    matching_class, quantiles, failed_rows
+):
+    """Every returned frame exposes failures without changing values or row labels."""
+
+    def backend(**kwargs):
+        result, _ = exact_backend(**kwargs)
+        if failed_rows:
+            result.iloc[-1, result.columns.get_loc("y")] = np.nan
+        return result, result.copy()
+
+    fitted = matching_class(backend).fit(donor_data(), ["x"], ["y"])
+    receiver = pd.DataFrame({"x": [1.0, 3.0, 5.0]}, index=[20, 10, 30])
+    result = fitted.predict(receiver, quantiles=quantiles)
+    frames = list(result.values()) if isinstance(result, dict) else [result]
+    expected = 2 * receiver.x + 1
+    if failed_rows:
+        expected.iloc[-1] = np.nan
+    for frame in frames:
+        assert frame.attrs["n_failed_records"] == failed_rows
+        assert frame.index.equals(receiver.index)
+        np.testing.assert_array_equal(frame.y, expected)
+    assert fitted.n_failed_records == failed_rows
+
+    fitted.matching_hotdeck = exact_backend
+    next_result = fitted.predict(receiver, quantiles=quantiles)
+    next_frames = (
+        list(next_result.values()) if isinstance(next_result, dict) else [next_result]
+    )
+    for frame in next_frames:
+        assert frame.attrs["n_failed_records"] == 0
+    assert fitted.n_failed_records == 0
+    assert all(frame.attrs["n_failed_records"] == failed_rows for frame in frames)