From a297d6391bf468a5334220c5d1500654408f25ff Mon Sep 17 00:00:00 2001 From: vahid-ahmadi Date: Wed, 16 Sep 2026 13:02:37 +0100 Subject: [PATCH 1/5] Give each per-variable QRF model its own seed Every per-variable model builds its own generator from the seed it is given, and all of them were handed self.seed. They therefore drew the same random quantiles in the same row order, so variables imputed together came out rank-comonotonic whatever their dependence in the donor: three targets with nil conditional dependence reproduced at 0.71 Spearman, 0.11 after this change. The offset follows the convention already used for the subsampling seed in _apply_max_train_samples. QRF also now accepts a seed argument; there was previously no way for a caller to vary the draws. Fixes #207 --- changelog.d/207.fixed.md | 1 + microimpute/models/qrf.py | 32 +++++++++++++++++++---- tests/test_models/test_qrf.py | 48 +++++++++++++++++++++++++++++++++++ 3 files changed, 76 insertions(+), 5 deletions(-) create mode 100644 changelog.d/207.fixed.md diff --git a/changelog.d/207.fixed.md b/changelog.d/207.fixed.md new file mode 100644 index 00000000..58db8524 --- /dev/null +++ b/changelog.d/207.fixed.md @@ -0,0 +1 @@ +Per-variable QRF models now derive distinct seeds, so variables imputed together no longer share one random quantile per row and come out comonotonic. `QRF` also accepts a `seed` argument. diff --git a/microimpute/models/qrf.py b/microimpute/models/qrf.py index 8edc5d08..da9f169d 100644 --- a/microimpute/models/qrf.py +++ b/microimpute/models/qrf.py @@ -10,7 +10,7 @@ from quantile_forest import RandomForestQuantileRegressor from sklearn.ensemble import RandomForestClassifier -from microimpute.config import VALIDATE_CONFIG +from microimpute.config import RANDOM_STATE, VALIDATE_CONFIG from microimpute.models.imputer import Imputer, ImputerResults try: @@ -607,6 +607,7 @@ def __init__( batch_size: Optional[int] = None, cleanup_interval: int = 10, max_train_samples: Optional[int] = None, + seed: Optional[int] = RANDOM_STATE, ) -> None: """Initialize the QRF model. @@ -618,8 +619,11 @@ def __init__( max_train_samples: If set, subsample X_train to at most this many rows before fitting. Reduces memory and training time while preserving sequential covariance structure. + seed: Base random seed. Each imputed variable is given a distinct + seed derived from it, so variables imputed together draw + independently. Pass None for non-reproducible draws. """ - super().__init__(log_level=log_level) + super().__init__(log_level=log_level, seed=seed) self.models = {} self.log_level = log_level self.memory_efficient = memory_efficient @@ -695,20 +699,38 @@ def _encode_imputed_variable( return data + def _seed_for_variable(self, variable: str) -> Optional[int]: + """Derive a distinct seed for one imputed variable. + + Each per-variable model builds its own generator from the seed it is + given. Handing every variable the same seed makes them draw the same + random quantiles in the same row order, so variables imputed together + come out comonotonic regardless of their dependence in the donor. The + offset follows the same convention as the subsampling seed below. + """ + if self.seed is None: + return None + try: + variable_offset = (self.imputed_variables or []).index(variable) + except ValueError: + variable_offset = 0 + return self.seed + variable_offset + def _create_model_for_variable(self, variable: str, **kwargs) -> Any: """Create the appropriate model (classifier or regressor) based on variable type.""" categorical_targets = getattr(self, "categorical_targets", {}) boolean_targets = getattr(self, "boolean_targets", {}) + seed = self._seed_for_variable(variable) if variable in categorical_targets: # Use classifier for categorical targets - return _RandomForestClassifierModel(seed=self.seed, logger=self.logger) + return _RandomForestClassifierModel(seed=seed, logger=self.logger) elif variable in boolean_targets: # Use classifier for boolean targets - return _RandomForestClassifierModel(seed=self.seed, logger=self.logger) + return _RandomForestClassifierModel(seed=seed, logger=self.logger) else: # Use QRF for numeric targets - return _QRFModel(seed=self.seed, logger=self.logger) + return _QRFModel(seed=seed, logger=self.logger) def _fit_model( self, diff --git a/tests/test_models/test_qrf.py b/tests/test_models/test_qrf.py index e3099f75..3e2f0b8a 100644 --- a/tests/test_models/test_qrf.py +++ b/tests/test_models/test_qrf.py @@ -1649,3 +1649,51 @@ def test_qrf_fit_predict_with_max_train_samples() -> None: assert result.shape == (n_test, 1) assert not result.isna().any().any() + + +def test_per_variable_models_draw_independently() -> None: + """Variables imputed together must not share one random quantile per row. + + Every per-variable model builds its own generator from the seed it is + given. When they all receive the same seed they draw the same quantiles in + the same row order, so the imputed variables come out comonotonic whatever + their dependence in the donor. + """ + rng = np.random.default_rng(7) + n = 1500 + predictors = pd.DataFrame( + {"inc": rng.normal(30, 8, n), "age": rng.normal(45, 12, n)} + ) + targets = ["savings", "property_wealth", "corporate_wealth"] + train = predictors.copy() + for variable in targets: + # Shared signal, independent shocks: the conditional dependence is nil. + train[variable] = 0.5 * predictors["inc"] + rng.normal(0, 10, n) + + test = pd.DataFrame({"inc": rng.normal(30, 8, 600), "age": rng.normal(45, 12, 600)}) + + model = QRF(log_level="WARNING") + imputations = model.fit(train, ["inc", "age"], targets).predict(test) + + ranks = imputations[targets].rank().corr() + off_diagonal = [ + abs(ranks.loc[a, b]) for i, a in enumerate(targets) for b in targets[i + 1 :] + ] + assert max(off_diagonal) < 0.4, ( + "imputed variables are comonotonic; per-variable models are sharing " + f"a seed (rank correlations {off_diagonal})" + ) + + +def test_seed_is_configurable_and_reproducible() -> None: + """QRF should accept a seed, and the same seed should reproduce draws.""" + rng = np.random.default_rng(3) + n = 400 + train = pd.DataFrame({"x": rng.normal(size=n)}) + train["y"] = train["x"] + rng.normal(0, 1, n) + test = pd.DataFrame({"x": rng.normal(size=120)}) + + first = QRF(log_level="WARNING", seed=1234).fit(train, ["x"], ["y"]).predict(test) + second = QRF(log_level="WARNING", seed=1234).fit(train, ["x"], ["y"]).predict(test) + + np.testing.assert_allclose(first["y"], second["y"]) From 4dd7caf316e7f8da52f65e82cd4f9887c7e21bff Mon Sep 17 00:00:00 2001 From: vahid-ahmadi Date: Wed, 16 Sep 2026 13:04:34 +0100 Subject: [PATCH 2/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 2862e6fd993b12c7399110e0f738605119d1b365 Mon Sep 17 00:00:00 2001 From: vahid-ahmadi Date: Wed, 16 Sep 2026 13:13:36 +0100 Subject: [PATCH 3/5] Fix the Lint job, which is failing on main The Lint job installs ruff>=0.9.0 with no upper bound. ruff 0.16.7 formats Python inside markdown code blocks, which earlier versions left alone, so five documentation files under docs/ became unformatted without anyone changing them. make check-format fails on an untouched checkout of main, and therefore on every open pull request. Reformats the five files and gives the constraint an upper bound, so a future ruff release changes the lint result only when someone chooses to move the pin. --- .github/workflows/pr_code_changes.yaml | 2 +- changelog.d/lint-on-main.fixed.md | 1 + docs/imputation-benchmarking/cross-validation.md | 12 ++++++------ docs/imputation-benchmarking/preprocessing.md | 8 ++++---- docs/imputation-benchmarking/visualizations.md | 8 ++------ docs/models/imputer/implement-new-model.md | 4 +--- docs/use_cases/index.md | 12 ++++++------ 7 files changed, 21 insertions(+), 26 deletions(-) create mode 100644 changelog.d/lint-on-main.fixed.md diff --git a/.github/workflows/pr_code_changes.yaml b/.github/workflows/pr_code_changes.yaml index 592dbe51..3afaf568 100644 --- a/.github/workflows/pr_code_changes.yaml +++ b/.github/workflows/pr_code_changes.yaml @@ -16,7 +16,7 @@ jobs: uses: astral-sh/setup-uv@v8.1.0 - name: Install relevant dependencies run: | - uv pip install "ruff>=0.9.0" --system + uv pip install "ruff>=0.9.0,<0.17.0" --system - name: Check code formatting run: make check-format diff --git a/changelog.d/lint-on-main.fixed.md b/changelog.d/lint-on-main.fixed.md new file mode 100644 index 00000000..592ead5c --- /dev/null +++ b/changelog.d/lint-on-main.fixed.md @@ -0,0 +1 @@ +Reformats five documentation files so the Lint job passes again, and bounds the ruff version the job installs. 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 From eabe40b002fc80adbdfa77a3b9ae758a4fa82fa4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mar=C3=ADa=20Juaristi?= <127882282+juaristi22@users.noreply.github.com> Date: Wed, 16 Sep 2026 15:03:08 +0200 Subject: [PATCH 4/5] Fix pre-submission imputation and evaluation correctness --- changelog.d/202.fixed.md | 1 + changelog.d/203.fixed.md | 1 + changelog.d/204.fixed.md | 1 + changelog.d/205.fixed.md | 1 + changelog.d/206-distributions.fixed.md | 1 + changelog.d/206-matching-scoring.breaking.md | 1 + changelog.d/206-models.fixed.md | 1 + changelog.d/206-pipeline.fixed.md | 1 + changelog.d/206-qrf-quantiles.breaking.md | 1 + changelog.d/208.fixed.md | 1 + changelog.d/209-target-types.breaking.md | 1 + changelog.d/209.fixed.md | 1 + changelog.d/212.fixed.md | 1 + changelog.d/213-validation.fixed.md | 1 + microimpute/comparisons/autoimpute.py | 90 ++-- microimpute/comparisons/autoimpute_helpers.py | 153 +++--- microimpute/comparisons/imputations.py | 28 +- microimpute/comparisons/metrics.py | 136 +++-- microimpute/comparisons/validation.py | 21 + microimpute/config.py | 4 +- microimpute/evaluations/cross_validation.py | 323 +++++------ microimpute/evaluations/predictor_analysis.py | 166 +++--- microimpute/models/imputer.py | 35 +- microimpute/models/matching.py | 502 ++++++------------ microimpute/models/ols.py | 77 ++- microimpute/models/qrf.py | 208 ++++++-- microimpute/models/quantreg.py | 61 ++- microimpute/models/zero_inflated.py | 294 ++++++---- microimpute/utils/data.py | 50 +- microimpute/utils/statmatch_hotdeck.py | 82 ++- microimpute/utils/type_handling.py | 100 ++-- tests/test_autoimpute.py | 24 +- tests/test_dashboard_formatter.py | 8 +- tests/test_data_preprocessing.py | 65 ++- tests/test_metrics.py | 47 +- tests/test_models/test_imputers.py | 20 +- tests/test_models/test_matching.py | 156 +++--- .../test_models/test_matching_correctness.py | 129 +++++ tests/test_models/test_qrf.py | 86 ++- tests/test_models/test_qrf_distribution.py | 234 ++++++++ .../test_regression_correctness.py | 147 +++++ tests/test_models/test_statmatch_bridge.py | 184 +++++++ .../test_zero_inflated_quantiles.py | 190 +++++++ tests/test_predictor_analysis.py | 94 ++++ tests/test_publication_pipeline.py | 316 +++++++++++ tests/test_quantile_comparison.py | 14 +- tests/test_type_handling.py | 4 +- 47 files changed, 2833 insertions(+), 1229 deletions(-) create mode 100644 changelog.d/202.fixed.md create mode 100644 changelog.d/203.fixed.md create mode 100644 changelog.d/204.fixed.md create mode 100644 changelog.d/205.fixed.md create mode 100644 changelog.d/206-distributions.fixed.md create mode 100644 changelog.d/206-matching-scoring.breaking.md create mode 100644 changelog.d/206-models.fixed.md create mode 100644 changelog.d/206-pipeline.fixed.md create mode 100644 changelog.d/206-qrf-quantiles.breaking.md create mode 100644 changelog.d/208.fixed.md create mode 100644 changelog.d/209-target-types.breaking.md create mode 100644 changelog.d/209.fixed.md create mode 100644 changelog.d/212.fixed.md create mode 100644 changelog.d/213-validation.fixed.md create mode 100644 tests/test_models/test_matching_correctness.py create mode 100644 tests/test_models/test_qrf_distribution.py create mode 100644 tests/test_models/test_regression_correctness.py create mode 100644 tests/test_models/test_statmatch_bridge.py create mode 100644 tests/test_models/test_zero_inflated_quantiles.py create mode 100644 tests/test_publication_pipeline.py diff --git a/changelog.d/202.fixed.md b/changelog.d/202.fixed.md new file mode 100644 index 00000000..b43d6822 --- /dev/null +++ b/changelog.d/202.fixed.md @@ -0,0 +1 @@ +Fit preprocessing statistics on donor training rows and reuse them for receivers, including single-row predictions. Cross-validation fits transformations within each fold and scores inverse-transformed predictions on the original target scale. Returned fitted models now replay donor preprocessing on future raw-data predictions and restore target units. diff --git a/changelog.d/203.fixed.md b/changelog.d/203.fixed.md new file mode 100644 index 00000000..1e04d39e --- /dev/null +++ b/changelog.d/203.fixed.md @@ -0,0 +1 @@ +Compute deterministic zero-inflated quantiles by inverting the ordered negative, zero, and positive mixture CDF. Preserve stochastic draws when quantiles are omitted, and reject unsupported sequential marginal quantiles or component predictions outside their sign support. diff --git a/changelog.d/204.fixed.md b/changelog.d/204.fixed.md new file mode 100644 index 00000000..71a654d4 --- /dev/null +++ b/changelog.d/204.fixed.md @@ -0,0 +1 @@ +Run core autoimpute tests without optional Matching or MDN dependencies instead of failing with an undefined Matching name. diff --git a/changelog.d/205.fixed.md b/changelog.d/205.fixed.md new file mode 100644 index 00000000..09a2e913 --- /dev/null +++ b/changelog.d/205.fixed.md @@ -0,0 +1 @@ +Apply the configured QRF defaults with at least 20 training observations per leaf and full leaf distributions. Compute survey-weighted conditional quantiles using each tree's normalized leaf weights and bootstrap donor multiplicities, and query sampled quantiles without discretizing them onto a truncated grid. diff --git a/changelog.d/206-distributions.fixed.md b/changelog.d/206-distributions.fixed.md new file mode 100644 index 00000000..40e1f5b3 --- /dev/null +++ b/changelog.d/206-distributions.fixed.md @@ -0,0 +1 @@ +Validate and propagate numeric zero-inflated sample weights to gate classifiers and sign-specific component fits, including aligned Series and array weights. Give numeric components reproducible independent random seeds and forward component fit parameters. diff --git a/changelog.d/206-matching-scoring.breaking.md b/changelog.d/206-matching-scoring.breaking.md new file mode 100644 index 00000000..b677caae --- /dev/null +++ b/changelog.d/206-matching-scoring.breaking.md @@ -0,0 +1 @@ +Matching donor draws are no longer presented as quantile forecasts or class probabilities. autoimpute excludes Matching from default distributional model selection and skips its unsupported scores if explicitly supplied. Direct Matching.predict without quantiles remains available. Rerun comparisons that previously scored replicated donor draws or fabricated probabilities. diff --git a/changelog.d/206-models.fixed.md b/changelog.d/206-models.fixed.md new file mode 100644 index 00000000..e5f028ef --- /dev/null +++ b/changelog.d/206-models.fixed.md @@ -0,0 +1 @@ +Allow QuantReg to fit requested prediction quantiles lazily on its original donor data; keep intercept columns on homogeneous receivers and reject nonfinite OLS/QuantReg inputs. Matching now explicitly rejects conditional quantile and class-probability requests instead of relabeling donor draws, retains and subsets donor weights during tuning, uses the documented weighted StatMatch API, and reports failed-record counts on returned data frames. Failed tuning trials remain pruned rather than scored using fallback values. Normalize OLS survey weights so arbitrary weight units cannot change predictive quantiles. Matching's R donor draws now use reproducible advancing child seeds while preserving the caller's R RNG state. diff --git a/changelog.d/206-pipeline.fixed.md b/changelog.d/206-pipeline.fixed.md new file mode 100644 index 00000000..a88e7b82 --- /dev/null +++ b/changelog.d/206-pipeline.fixed.md @@ -0,0 +1 @@ +Honor autoimpute train_size by reproducibly sampling donor rows, forward random_state to model fitting, evaluate requested quantile grids, and retune final parameters on all training rows instead of selecting the luckiest outer test fold. Report the standard deviation of fold-average losses correctly. Distributional QRF comparisons use independent target fits so reported quantiles are conditional on the original predictors; stochastic sequential donor draws remain available through direct QRF use. diff --git a/changelog.d/206-qrf-quantiles.breaking.md b/changelog.d/206-qrf-quantiles.breaking.md new file mode 100644 index 00000000..035331cd --- /dev/null +++ b/changelog.d/206-qrf-quantiles.breaking.md @@ -0,0 +1 @@ +Add QRF(sequential=False) to estimate each target's conditional marginal quantiles using the original predictors only, with target-order-independent seeds. Sequential multi-target QRF now rejects explicit quantiles because same-quantile chaining does not calculate marginal quantiles; stochastic sequential draws remain supported. Distributional comparison and imputation helpers use independent QRF fits. Numeric hyperparameter tuning now evaluates the actual predicted median instead of scoring a random draw as a median. diff --git a/changelog.d/208.fixed.md b/changelog.d/208.fixed.md new file mode 100644 index 00000000..72f0efe6 --- /dev/null +++ b/changelog.d/208.fixed.md @@ -0,0 +1 @@ +Draw one independent OLS residual per receiver from a persistent random generator, advancing across targets and prediction calls while preserving reproducibility from the same seed. QuantReg random grid sampling also advances independently across calls and targets. diff --git a/changelog.d/209-target-types.breaking.md b/changelog.d/209-target-types.breaking.md new file mode 100644 index 00000000..f5237592 --- /dev/null +++ b/changelog.d/209-target-types.breaking.md @@ -0,0 +1 @@ +Numeric columns, including integer counts and 0/1 integers, now remain numeric regardless of cardinality or the values present in a fold. Declare categorical targets with target_types={"column": "categorical"}, pandas categorical dtype, or boolean dtype for binary categories. Log-loss comparisons require the probabilities returned by predict(return_probs=True). diff --git a/changelog.d/209.fixed.md b/changelog.d/209.fixed.md new file mode 100644 index 00000000..4a0eff5e --- /dev/null +++ b/changelog.d/209.fixed.md @@ -0,0 +1 @@ +Compute categorical comparison and cross-validation log loss from actual model probabilities. Reject class-label inputs rather than fabricating 0.99/0.01 probabilities, and align classes when a cross-validation fold lacks a category. diff --git a/changelog.d/212.fixed.md b/changelog.d/212.fixed.md new file mode 100644 index 00000000..3e4bac46 --- /dev/null +++ b/changelog.d/212.fixed.md @@ -0,0 +1 @@ +Compute normalized mutual information from consistent paired discretizations and natural-log entropies, with configurable continuous-variable bins and explicit zero information for constants. Predictor analysis now scores categorical forecasts from their actual probabilities. diff --git a/changelog.d/213-validation.fixed.md b/changelog.d/213-validation.fixed.md new file mode 100644 index 00000000..b875fc60 --- /dev/null +++ b/changelog.d/213-validation.fixed.md @@ -0,0 +1 @@ +Reject duplicate predictor/target names, predictor-target overlap, duplicate DataFrame columns, incompatible numeric/string predictor dtypes, and infinite sample weights with actionable errors. Treat numeric int/float predictor dtypes as compatible. diff --git a/microimpute/comparisons/autoimpute.py b/microimpute/comparisons/autoimpute.py index 2121c6a7..e5d5ff17 100644 --- a/microimpute/comparisons/autoimpute.py +++ b/microimpute/comparisons/autoimpute.py @@ -17,6 +17,7 @@ evaluate_model, fit_and_predict_model, prepare_data_for_imputation, + preprocessing_aware_model, select_best_model_dual_metrics, validate_autoimpute_inputs, ) @@ -28,11 +29,10 @@ ) from microimpute.models import OLS, QRF, Imputer, QuantReg from microimpute.utils.data import ( - un_asinh_transform_predictions, - unlog_transform_predictions, unnormalize_predictions, + reverse_transformations, ) -from microimpute.utils.type_handling import VariableTypeDetector +from microimpute.utils.type_handling import VariableTypeDetector, declare_target_types try: from microimpute.models import Matching @@ -76,22 +76,12 @@ def _reverse_transformations( return unnormalize_predictions(imputations, params) elif transform_type == "preprocessing": - # New preprocessing format with multiple transformation types - result = imputations - - # Reverse normalization if any - if params.get("normalization"): - result = unnormalize_predictions(result, params["normalization"]) - - # Reverse log transform if any - if params.get("log_transform"): - result = unlog_transform_predictions(result, params["log_transform"]) - - # Reverse asinh transform if any - if params.get("asinh_transform"): - result = un_asinh_transform_predictions(result, params["asinh_transform"]) - - return result + return { + q: reverse_transformations(frame, params) + if isinstance(frame, pd.DataFrame) + else frame + for q, frame in imputations.items() + } else: log.warning(f"Unknown transform type: {transform_type}") @@ -121,16 +111,17 @@ class AutoImputeResult(BaseModel): receiver_data : pd.DataFrame Copy of the receiver data with the median-quantile imputations of the best performing model attached. fitted_models : Dict[str, Any] - Mapping model name → fitted Imputer instance. + Mapping model name → fitted model. Returned models accept raw receiver data + and replay donor-fitted preprocessing, returning predictions in original units. cv_results : Dict[str, Dict[str, Any]] Cross-validation results with separate quantile_loss and log_loss metrics for each model. """ model_config = ConfigDict(arbitrary_types_allowed=True) - imputations: Union[ - Dict[str, Dict[float, pd.DataFrame]], Dict[str, pd.DataFrame] - ] = Field(...) + imputations: Dict[str, Union[pd.DataFrame, Dict[Union[float, str], Any]]] = Field( + ... + ) receiver_data: pd.DataFrame = Field(...) fitted_models: Dict[str, Any] = Field(...) cv_results: Dict[str, Dict[str, Any]] = Field(...) @@ -198,7 +189,6 @@ def _setup_logging(log_level: str) -> int: } numeric_level = level_map[log_level] log.setLevel(numeric_level) - warnings.filterwarnings("ignore") return numeric_level @@ -214,6 +204,8 @@ def _evaluate_models_parallel( tune_hyperparameters: bool, hyperparameters: Optional[Dict[str, Dict[str, Any]]], n_jobs: int = -1, + preprocessing: Optional[Dict[str, str]] = None, + target_types: Optional[Dict[str, str]] = None, ) -> Tuple[Dict[str, Dict[str, Any]], Optional[Dict[str, Any]]]: """Evaluate multiple models in parallel using cross-validation with dual metrics. @@ -248,6 +240,8 @@ def _evaluate_models_parallel( random_state, tune_hyperparameters, model_hyperparams, + preprocessing, + target_types, ) ) @@ -291,6 +285,9 @@ def _generate_imputations_for_all_models( hyperparams: Optional[Dict[str, Any]], log_level: str, preprocessing: Optional[Dict[str, str]] = None, + random_state: int = RANDOM_STATE, + target_types: Optional[Dict[str, str]] = None, + quantiles: Optional[List[float]] = None, ) -> Tuple[Dict[str, pd.DataFrame], Dict[str, Any]]: """Generate imputations for all models when impute_all=True. @@ -307,7 +304,7 @@ def _generate_imputations_for_all_models( for model_class in model_classes: model_name = model_class.__name__ - if model_name == best_method: + if model_name == best_method or model_name == "Matching": continue # Skip the best method as it's already done # Check if model can handle the variable types using original data @@ -345,13 +342,18 @@ def _generate_imputations_for_all_models( imputation_q, model_hyperparams, log_level, + random_state=random_state, + target_types=target_types, + quantiles=quantiles, ) # Reverse transformations if needed final_imputations = _reverse_transformations(imputations, transform_params) final_imputations_dict[model_name] = final_imputations[imputation_q] - fitted_models_dict[model_name] = fitted_model + fitted_models_dict[model_name] = preprocessing_aware_model( + fitted_model, transform_params + ) return final_imputations_dict, fitted_models_dict @@ -375,6 +377,7 @@ def autoimpute( k_folds: Optional[int] = 5, force_retrain: Optional[bool] = False, log_level: Optional[str] = "WARNING", + target_types: Optional[Dict[str, str]] = None, ) -> AutoImputeResult: """Automatically select and apply the best imputation model. @@ -414,8 +417,11 @@ def autoimpute( 'numerical': select based on quantile loss only 'categorical': select based on log loss only 'combined': weighted average of both metrics + target_types : Optional mapping of target names to "numeric", "categorical", + or "bool". Numeric dtypes remain numeric unless explicitly declared. random_state : Random seed for reproducibility - train_size : Proportion of data to use for training in preprocessing + train_size : Fraction of donor rows used for comparison and final fitting; + sampled without replacement using random_state. Use 1.0 for all rows. k_folds : Number of folds for cross-validation. Defaults to 5. force_retrain : If True, forces MDN models to retrain instead of using cached models. Defaults to False. @@ -452,7 +458,9 @@ def autoimpute( receiver_data = receiver_data.copy() # Use provided quantiles or defaults - quantiles = imputation_quantiles if imputation_quantiles else QUANTILES + quantiles = ( + imputation_quantiles if imputation_quantiles is not None else QUANTILES + ) # Validate all inputs validate_autoimpute_inputs( @@ -473,6 +481,14 @@ def autoimpute( f"with predictors {predictors}." ) + donor_data = declare_target_types(donor_data, imputed_variables, target_types) + if train_size is None or not 0 < train_size <= 1: + raise ValueError("train_size must be greater than 0 and at most 1") + if train_size < 1: + donor_data = donor_data.sample(frac=train_size, random_state=random_state) + if len(donor_data) < k_folds: + raise ValueError("train_size leaves fewer donor rows than k_folds") + # Step 1: Data preparation if numeric_log_level <= logging.INFO: log.info("Preprocessing data...") @@ -501,8 +517,6 @@ def autoimpute( # Get model classes if not models: model_classes: List[Type[Imputer]] = [QRF, OLS, QuantReg] - if HAS_MATCHING: - model_classes.append(Matching) if HAS_MDN: model_classes.append(MDN) else: @@ -534,7 +548,7 @@ def autoimpute( # Evaluate models in parallel method_results, best_hyperparams = _evaluate_models_parallel( model_classes, - training_data, + donor_data, predictors, imputed_variables, weight_col, @@ -543,6 +557,8 @@ def autoimpute( random_state, tune_hyperparameters, hyperparameters, + preprocessing=preprocessing, + target_types=target_types, ) # Step 3: Model selection @@ -609,6 +625,9 @@ def autoimpute( imputation_q, model_hyperparams, log_level, + random_state=random_state, + target_types=target_types, + quantiles=sorted(set((imputation_quantiles or []) + [0.5])), ) # Reverse transformations if needed @@ -635,7 +654,11 @@ def autoimpute( else final_imputations ) } - fitted_models_dict = {"best_method": best_fitted_model} + fitted_models_dict = { + "best_method": preprocessing_aware_model( + best_fitted_model, transform_params + ) + } # Step 5: Generate imputations for all models if requested if impute_all: @@ -665,6 +688,9 @@ def autoimpute( merged_hyperparams if merged_hyperparams else None, log_level, preprocessing=preprocessing, + random_state=random_state, + target_types=target_types, + quantiles=sorted(set((imputation_quantiles or []) + [0.5])), ) final_imputations_dict.update(other_imputations) fitted_models_dict.update(other_models) diff --git a/microimpute/comparisons/autoimpute_helpers.py b/microimpute/comparisons/autoimpute_helpers.py index 8cc0f321..37379c39 100644 --- a/microimpute/comparisons/autoimpute_helpers.py +++ b/microimpute/comparisons/autoimpute_helpers.py @@ -25,7 +25,13 @@ ) from microimpute.evaluations import cross_validate_model from microimpute.models import Imputer -from microimpute.utils.data import preprocess_data +from microimpute.models.imputer import create_distributional_model +from microimpute.config import RANDOM_STATE +from microimpute.utils.data import ( + preprocess_data, + apply_transformations, + reverse_transformations, +) log = logging.getLogger(__name__) @@ -64,7 +70,7 @@ def validate_autoimpute_inputs( raise ValueError(error_msg) # Validate quantiles if provided - if quantiles: + if quantiles is not None: validate_quantiles(quantiles) # Validate data and columns @@ -154,23 +160,9 @@ def prepare_data_for_imputation( asinh_transform=asinh_cols if asinh_cols else False, ) - # Apply same transformations to predictors in imputing data - predictor_normalize = [c for c in normalize_cols if c in predictors] - predictor_log = [c for c in log_cols if c in predictors] - predictor_asinh = [c for c in asinh_cols if c in predictors] - - if predictor_normalize or predictor_log or predictor_asinh: - transformed_imputing, _ = preprocess_data( - imputing_data[predictors], - full_data=True, - train_size=train_size, - test_size=test_size, - normalize=(predictor_normalize if predictor_normalize else False), - log_transform=predictor_log if predictor_log else False, - asinh_transform=predictor_asinh if predictor_asinh else False, - ) - else: - transformed_imputing = imputing_data[predictors].copy() + transformed_imputing = apply_transformations( + imputing_data[predictors], transform_result + ) training_data = transformed_training if weight_col: @@ -197,20 +189,12 @@ def prepare_data_for_imputation( }, } - # Only return params if there are transformations to reverse - has_transforms = any( - imputed_transform_params[key] - for key in ["normalization", "log_transform", "asinh_transform"] - ) - - if has_transforms: - transform_params = { - "type": "preprocessing", - "params": imputed_transform_params, - } - return training_data, imputing_data, transform_params - else: - return training_data, imputing_data, None + transform_params = { + "type": "preprocessing", + "params": imputed_transform_params, + "all_params": transform_result, + } + return training_data, imputing_data, transform_params else: # No transformation needed @@ -247,6 +231,8 @@ def evaluate_model( random_state: int, tune_hyperparams: bool, hyperparameters: Optional[Dict[str, Any]], + preprocessing: Optional[Dict[str, str]] = None, + target_types: Optional[Dict[str, str]] = None, ) -> tuple: """Evaluate a single imputation model with cross-validation. @@ -280,6 +266,8 @@ def evaluate_model( random_state=random_state, tune_hyperparameters=tune_hyperparams, model_hyperparams=hyperparameters, + preprocessing=preprocessing, + target_types=target_types, ) if tune_hyperparams and isinstance(cv_result, tuple) and len(cv_result) == 2: @@ -299,6 +287,9 @@ def fit_and_predict_model( quantile: float, hyperparams: Optional[Dict[str, Any]] = None, log_level: str = "WARNING", + random_state: int = RANDOM_STATE, + target_types: Optional[Dict[str, str]] = None, + quantiles: Optional[List[float]] = None, ) -> Tuple[Any, Dict[float, pd.DataFrame]]: """Fit a model and generate predictions. @@ -317,7 +308,9 @@ def fit_and_predict_model( Tuple of (fitted_model, predictions_dict) """ model_name = model_class.__name__ - model = model_class(log_level=log_level) + model = create_distributional_model( + model_class, log_level=log_level, seed=random_state + ) # Check for categorical variables from microimpute.comparisons.metrics import get_metric_for_variable_type @@ -341,40 +334,18 @@ def fit_and_predict_model( log.error(error_msg) raise ValueError(error_msg) - # Fit the model + requested_quantiles = quantiles if quantiles is not None else [quantile] + params = dict(hyperparams or {}) + params["target_types"] = target_types if model_name == "QuantReg": - # QuantReg needs explicit quantiles during fitting - fitted_model = model.fit( - training_data, - predictors, - imputed_variables, - weight_col=weight_col, - quantiles=[quantile], - ) - elif hyperparams and model_name in ["Matching", "QRF", "MDN"]: - # Apply hyperparameters for specific models - fitted_model = model.fit( - training_data, - predictors, - imputed_variables, - weight_col=weight_col, - **hyperparams, - ) - else: - fitted_model = model.fit( - training_data, - predictors, - imputed_variables, - weight_col=weight_col, - ) - - # Generate predictions with return_probs for categorical variables - if has_categorical: - imputations = fitted_model.predict( - imputing_data, quantiles=[quantile], return_probs=True - ) - else: - imputations = fitted_model.predict(imputing_data, quantiles=[quantile]) + params["quantiles"] = requested_quantiles + fitted_model = model.fit( + training_data, predictors, imputed_variables, weight_col=weight_col, **params + ) + has_categorical = bool(model.categorical_targets or model.boolean_targets) + imputations = fitted_model.predict( + imputing_data, quantiles=requested_quantiles, return_probs=has_categorical + ) # Handle case where predict returns a DataFrame directly if isinstance(imputations, pd.DataFrame): @@ -550,3 +521,51 @@ def select_best_model_dual_metrics( log.info(f"Selected {best_method} based on combined metric: {best_score:.6f}") return best_method, model_metrics[best_method] + + +class PreprocessedImputerResults: + """A fitted model that accepts raw receiver data and returns original units. + + The wrapped learner remains available through ``fitted_model``; its public + metadata (predictors, seed, models, etc.) is forwarded unchanged. + """ + + def __init__(self, fitted_model: Any, transform_params: dict): + self.fitted_model = fitted_model + self.transform_params = transform_params + + def __getattr__(self, name: str) -> Any: + fitted_model = self.__dict__.get("fitted_model") + if fitted_model is None: + raise AttributeError(name) + return getattr(fitted_model, name) + + def predict( + self, + X_test: pd.DataFrame, + quantiles: Optional[List[float]] = None, + return_probs: bool = False, + **kwargs: Any, + ) -> Any: + receiver = X_test.drop(columns=self.imputed_variables, errors="ignore") + transformed = apply_transformations(receiver, self.transform_params) + predictions = self.fitted_model.predict( + transformed, quantiles=quantiles, return_probs=return_probs, **kwargs + ) + if isinstance(predictions, pd.DataFrame): + return reverse_transformations(predictions, self.transform_params) + return { + q: reverse_transformations(frame, self.transform_params) + if isinstance(frame, pd.DataFrame) + else frame + for q, frame in predictions.items() + } + + +def preprocessing_aware_model( + fitted_model: Any, transform_params: Optional[dict] +) -> Any: + """Attach fitted preprocessing to models returned after initial imputation.""" + if transform_params and any(transform_params.get("all_params", {}).values()): + return PreprocessedImputerResults(fitted_model, transform_params["all_params"]) + return fitted_model diff --git a/microimpute/comparisons/imputations.py b/microimpute/comparisons/imputations.py index f24ecd36..defdb44e 100644 --- a/microimpute/comparisons/imputations.py +++ b/microimpute/comparisons/imputations.py @@ -8,6 +8,7 @@ import logging from typing import Any, Dict, List, Optional, Type +import numpy as np import pandas as pd from pydantic import validate_call @@ -16,7 +17,10 @@ validate_quantiles, ) from microimpute.config import QUANTILES, VALIDATE_CONFIG +from microimpute.comparisons.metrics import get_metric_for_variable_type from microimpute.models.quantreg import QuantReg +from microimpute.models.imputer import create_distributional_model +from microimpute.utils.type_handling import declare_target_types log = logging.getLogger(__name__) @@ -29,7 +33,8 @@ def get_imputations( predictors: List[str], imputed_variables: List[str], quantiles: Optional[List[float]] = QUANTILES, -) -> Dict[str, Dict[float, pd.DataFrame]]: + target_types: Optional[Dict[str, str]] = None, +) -> Dict[str, dict]: """Generate imputations using multiple model classes for the specified variables. Args: @@ -63,6 +68,7 @@ def get_imputations( if quantiles: validate_quantiles(quantiles) + X_train = declare_target_types(X_train, imputed_variables, target_types) log.info(f"Generating imputations for {len(model_classes)} model classes") log.info( f"Training data shape: {X_train.shape}, Test data shape: {X_test.shape}" @@ -82,7 +88,7 @@ def get_imputations( try: # Instantiate the model - model = model_class() + model = create_distributional_model(model_class) # Handle QuantReg which needs quantiles during fitting if model_class == QuantReg: @@ -99,7 +105,23 @@ def get_imputations( # Get predictions log.info(f"Generating predictions with {model_name}") - imputations = fitted_model.predict(X_test, quantiles) + imputations = fitted_model.predict( + X_test, + quantiles, + return_probs=any( + get_metric_for_variable_type(X_train[var], var) == "log_loss" + for var in imputed_variables + ), + ) + for variable, info in model.constant_targets.items(): + if ( + get_metric_for_variable_type(X_train[variable], variable) + == "log_loss" + ): + imputations.setdefault("probabilities", {})[variable] = { + "probabilities": np.ones((len(X_test), 1)), + "classes": np.asarray([info["value"]]), + } method_imputations[model_name] = imputations except (TypeError, AttributeError, ValueError) as model_error: diff --git a/microimpute/comparisons/metrics.py b/microimpute/comparisons/metrics.py index 92f8fd4a..d8f0c34d 100644 --- a/microimpute/comparisons/metrics.py +++ b/microimpute/comparisons/metrics.py @@ -23,7 +23,7 @@ validate_quantiles, ) from microimpute.config import QUANTILES, VALIDATE_CONFIG -from microimpute.utils.type_handling import VariableTypeDetector +from microimpute.utils.type_handling import VariableTypeDetector, declare_target_types log = logging.getLogger(__name__) @@ -74,71 +74,46 @@ def log_loss( normalize: bool = True, labels: Optional[np.ndarray] = None, ) -> float: - """Calculate log loss for categorical predictions. + """Calculate log loss from genuine probabilities, never class labels. - Args: - y_true: True labels (can be class indices or one-hot encoded). - y_pred: Predicted probabilities. Shape should be (n_samples,) for binary - or (n_samples, n_classes) for multiclass. - If class labels are provided instead of probabilities, they will be - converted to high-confidence probabilities (0.99/0.01) with a warning. - normalize: If True, return the mean loss. If False, return sum. - labels: List of labels to include in the loss computation. - - Returns: - Log loss value. - - Note: - For more accurate metrics, models should provide predicted probabilities - rather than class labels. Use model.predict_proba() instead of model.predict() - when available. + Binary vectors must have floating dtype. Probability matrices must have + one column per class, finite values in [0, 1], and rows summing to one. """ - try: - # Handle case where predictions are class labels instead of probabilities - if len(y_pred.shape) == 1 or (len(y_pred.shape) == 2 and y_pred.shape[1] == 1): - # Binary case or class predictions - if labels is None: - labels = np.unique(y_true) - - # Convert to probabilities if needed - if np.all(np.isin(y_pred.flatten(), labels)): - # These are class predictions, not probabilities - log.info( - "Converting class labels to probabilities for log loss computation. " - "For more accurate metrics, please provide predicted probabilities " - "using model.predict_proba() or equivalent method instead of class predictions. " - "Class labels are being converted to high-confidence probabilities (0.99/0.01)." - ) - - # Create one-hot encoded probabilities with high confidence - n_samples = len(y_true) - n_classes = len(labels) - - if n_classes == 2: - # Binary case - y_pred_proba = np.zeros(n_samples) - y_pred_proba[y_pred.flatten() == labels[1]] = 0.99 - y_pred_proba[y_pred.flatten() == labels[0]] = 0.01 - else: - # Multiclass case - y_pred_proba = np.full( - (n_samples, n_classes), 0.01 / (n_classes - 1) - ) - for i, label in enumerate(labels): - mask = y_pred.flatten() == label - y_pred_proba[mask, i] = 0.99 - - y_pred = y_pred_proba - - log.info( - f"Converted {n_samples} class predictions to probabilities " - f"for {n_classes}-class classification." - ) - - return sklearn_log_loss(y_true, y_pred, normalize=normalize, labels=labels) - except Exception as e: - log.error(f"Error computing log loss: {str(e)}") - raise RuntimeError(f"Failed to compute log loss: {str(e)}") from e + y_pred = np.asarray(y_pred) + if not np.issubdtype(y_pred.dtype, np.number): + raise ValueError("Log loss requires predicted probabilities, not class labels") + if y_pred.ndim == 1 and not np.issubdtype(y_pred.dtype, np.floating): + raise ValueError( + "Log loss requires floating predicted probabilities, not class labels" + ) + if ( + y_pred.ndim not in (1, 2) + or not np.isfinite(y_pred).all() + or ((y_pred < 0) | (y_pred > 1)).any() + ): + raise ValueError("Predicted probabilities must be finite and between 0 and 1") + if y_pred.ndim == 2 and not np.allclose(y_pred.sum(axis=1), 1): + raise ValueError("Predicted probability rows must sum to one") + y_true = np.asarray(y_true) + if y_pred.ndim == 2 and y_true.ndim == 1: + # Matrix columns follow sklearn's sorted-label convention. A fitted + # model may lack a class appearing in held-out rows: its probability + # for that class is exactly zero, not an invented confidence level. + model_labels = ( + np.sort(np.asarray(labels)) if labels is not None else np.unique(y_true) + ) + if len(model_labels) != y_pred.shape[1]: + raise ValueError("Probability columns must match the supplied class labels") + all_labels = np.union1d(model_labels, np.unique(y_true)) + if len(all_labels) == 1: + return 0.0 + if len(all_labels) != len(model_labels): + aligned = np.zeros((len(y_pred), len(all_labels)), dtype=float) + for index, label in enumerate(model_labels): + aligned[:, np.flatnonzero(all_labels == label)[0]] = y_pred[:, index] + y_pred = aligned + labels = all_labels + return sklearn_log_loss(y_true, y_pred, normalize=normalize, labels=labels) def order_probabilities_alphabetically( @@ -334,16 +309,16 @@ def _compute_method_losses( # Get values as numpy arrays (handles Arrow-backed dtypes) test_values = np.asarray(test_y[variable]) - pred_values = np.asarray(imputation[quantile][variable]) - - # Get unique labels from test data - labels = np.unique(test_values) - - # Compute loss - # Note: If pred_values contains class labels instead of probabilities, - # they will be converted with a warning + info = imputation.get("probabilities", {}).get(variable) + if info is None: + raise ValueError( + f"Log loss for '{variable}' requires predicted probabilities; call predict(return_probs=True)" + ) + probabilities, labels = order_probabilities_alphabetically( + np.asarray(info["probabilities"]), np.asarray(info["classes"]) + ) _, mean_loss = compute_loss( - test_values, pred_values, "log_loss", labels=labels + test_values, probabilities, "log_loss", labels=labels ) categorical_losses.append(mean_loss) @@ -431,8 +406,10 @@ def _compute_method_losses( @validate_call(config=VALIDATE_CONFIG) def compare_metrics( test_y: pd.DataFrame, - method_imputations: Dict[str, Dict[float, pd.DataFrame]], + method_imputations: Dict[str, dict], imputed_variables: List[str], + quantiles: Optional[List[float]] = None, + target_types: Optional[Dict[str, str]] = None, ) -> pd.DataFrame: """Compare metrics across different imputation methods. @@ -463,6 +440,7 @@ def compare_metrics( # Validate inputs validate_columns_exist(test_y, imputed_variables, "test_y") + test_y = declare_target_types(test_y, imputed_variables, target_types) # Detect metric type for each variable variable_metrics = {} for var in imputed_variables: @@ -475,12 +453,20 @@ def compare_metrics( # Process each method for method, imputation in method_imputations.items(): + method_quantiles = ( + quantiles + if quantiles is not None + else [q for q in imputation if isinstance(q, (float, int))] + ) + if not method_quantiles: + raise ValueError(f"No prediction quantiles for {method}") + validate_quantiles(method_quantiles) method_results = _compute_method_losses( method, imputation, test_y, imputed_variables, - QUANTILES, + method_quantiles, variable_metrics, ) all_results.extend(method_results) diff --git a/microimpute/comparisons/validation.py b/microimpute/comparisons/validation.py index 06e2f6b0..a6234b73 100644 --- a/microimpute/comparisons/validation.py +++ b/microimpute/comparisons/validation.py @@ -25,6 +25,8 @@ def validate_quantiles(quantiles: List[float]) -> None: Raises: ValueError: If any quantile is outside [0, 1] range. """ + if not quantiles: + raise ValueError("At least one quantile is required") invalid_quantiles = [q for q in quantiles if not 0 <= q <= 1] if invalid_quantiles: error_msg = f"Invalid quantiles (must be between 0 and 1): {invalid_quantiles}" @@ -101,6 +103,17 @@ def validate_imputation_inputs( Raises: ValueError: If validation fails. """ + for name, columns in [ + ("predictors", predictors), + ("imputed_variables", imputed_variables), + ]: + if len(columns) != len(set(columns)): + raise ValueError(f"Duplicate column names in {name}") + overlap = set(predictors) & set(imputed_variables) + if overlap: + raise ValueError(f"Predictors and imputed_variables overlap: {sorted(overlap)}") + if not donor_data.columns.is_unique or not receiver_data.columns.is_unique: + raise ValueError("Duplicate DataFrame column names are not supported") # Validate donor data has all required columns validate_columns_exist(donor_data, predictors, "donor data") validate_columns_exist(donor_data, imputed_variables, "donor data") @@ -108,6 +121,14 @@ def validate_imputation_inputs( # Validate receiver data has predictor columns validate_columns_exist(receiver_data, predictors, "receiver data") + for column in predictors: + donor_numeric = pd.api.types.is_numeric_dtype(donor_data[column]) + receiver_numeric = pd.api.types.is_numeric_dtype(receiver_data[column]) + if donor_numeric != receiver_numeric: + raise ValueError( + f"Incompatible predictor dtype for '{column}': donor {donor_data[column].dtype}, receiver {receiver_data[column].dtype}" + ) + # Validate weight column if provided if weight_col: validate_columns_exist(donor_data, [weight_col], "donor data") diff --git a/microimpute/config.py b/microimpute/config.py index ba5896a2..b2118f45 100644 --- a/microimpute/config.py +++ b/microimpute/config.py @@ -46,7 +46,9 @@ "n_estimators": 100, "max_depth": None, "min_samples_split": 2, - "min_samples_leaf": 1, + "min_samples_leaf": 20, + # Retain the leaf distribution instead of one randomly selected donor. + "max_samples_leaf": None, "max_features": 1.0, }, "quantreg": { diff --git a/microimpute/evaluations/cross_validation.py b/microimpute/evaluations/cross_validation.py index ffe2682d..c36cffff 100644 --- a/microimpute/evaluations/cross_validation.py +++ b/microimpute/evaluations/cross_validation.py @@ -22,12 +22,19 @@ validate_quantiles, ) from microimpute.config import QUANTILES, RANDOM_STATE, VALIDATE_CONFIG +from microimpute.utils.data import ( + preprocess_data, + apply_transformations, + reverse_transformations, +) +from microimpute.utils.type_handling import declare_target_types try: from microimpute.models.matching import Matching except ImportError: # optional dependency Matching = None from microimpute.models.quantreg import QuantReg +from microimpute.models.imputer import create_distributional_model log = logging.getLogger(__name__) @@ -43,6 +50,9 @@ def _process_single_fold( model_hyperparams: Optional[dict], tune_hyperparameters: bool, variable_metrics: Dict[str, str], + preprocessing: Optional[Dict[str, str]] = None, + target_types: Optional[Dict[str, str]] = None, + random_state: int = RANDOM_STATE, ) -> Tuple[ int, Dict, @@ -63,8 +73,14 @@ def _process_single_fold( train_y = {var: train_data[var].values for var in imputed_variables} test_y = {var: test_data[var].values for var in imputed_variables} - # Instantiate and fit the model - model = model_class() + transform_params = {} + if preprocessing: + train_data, transform_params = preprocess_data( + train_data, full_data=True, **_preprocessing_kwargs(preprocessing) + ) + test_data = apply_transformations(test_data, transform_params) + # Instantiate with the caller's seed for fitting and prediction sampling. + model = create_distributional_model(model_class, seed=random_state) fold_tuned_params = None # Fit model with appropriate parameters @@ -78,6 +94,7 @@ def _process_single_fold( quantiles, model_hyperparams, tune_hyperparameters, + target_types, ) # Check if model fitting failed (incompatible with variable types) @@ -106,6 +123,25 @@ def _process_single_fold( fold_test_imputations = fitted_model.predict(test_data, quantiles) fold_train_imputations = fitted_model.predict(train_data, quantiles) + if has_categorical: + for predictions, frame in [ + (fold_test_imputations, test_data), + (fold_train_imputations, train_data), + ]: + probabilities = predictions.setdefault("probabilities", {}) + for variable, info in model.constant_targets.items(): + if variable_metrics[variable] == "log_loss": + probabilities[variable] = { + "probabilities": np.ones((len(frame), 1)), + "classes": np.asarray([info["value"]]), + } + if transform_params: + for predictions in [fold_test_imputations, fold_train_imputations]: + for quantile in quantiles: + predictions[quantile] = reverse_transformations( + predictions[quantile], transform_params + ) + return ( fold_idx, fold_test_imputations, @@ -126,106 +162,54 @@ def _fit_model_for_fold( quantiles: List[float], model_hyperparams: Optional[dict], tune_hyperparameters: bool, + target_types: Optional[Dict[str, str]] = None, ) -> Tuple[Any, Optional[dict]]: """Fit a model for a single fold with appropriate parameters. Returns None for fitted_model if the model cannot handle the variable types. """ model_name = model_class.__name__ - fold_tuned_params = None - - # Special handling for QuantReg with categorical variables - if model_name == "QuantReg": - # Check if any imputed variables are categorical - from microimpute.comparisons.metrics import ( - get_metric_for_variable_type, - ) - - for var in imputed_variables: - if get_metric_for_variable_type(train_data[var], var) == "log_loss": - log.warning( - f"QuantReg does not support categorical variable '{var}'. " - f"Skipping QuantReg for this fold." - ) - return None, None - - # Handle model-specific hyperparameters - if model_hyperparams: - try: - log.info(f"Fitting {model_name} with hyperparameters: {model_hyperparams}") - fitted_model = model.fit( - X_train=train_data, - predictors=predictors, - imputed_variables=imputed_variables, - weight_col=weight_col, - **model_hyperparams, - ) - except ValueError as e: - # Check if it's due to categorical incompatibility - if "QuantReg does not support categorical" in str(e): - log.warning(f"{model_name} incompatible with variable types: {str(e)}") - return None, None - raise e - except TypeError as e: - log.warning( - f"Invalid hyperparameters for {model_name}, using defaults: {str(e)}" - ) - fitted_model = model.fit( - X_train=train_data, - predictors=predictors, - imputed_variables=imputed_variables, - weight_col=weight_col, - ) - raise ValueError(f"Invalid hyperparameters for {model_name}") from e - - # Handle QuantReg which needs explicit quantiles - elif model_class == QuantReg: - try: - log.info(f"Fitting QuantReg model with explicit quantiles") - fitted_model = model.fit( - train_data, - predictors, - imputed_variables, - weight_col=weight_col, - quantiles=quantiles, - ) - except ValueError as e: - if "QuantReg does not support categorical" in str(e): - log.warning(f"QuantReg incompatible with variable types: {str(e)}") - return None, None - raise e - - # Handle hyperparameter tuning for QRF, Matching, and MDN - elif tune_hyperparameters and model_name in ["QRF", "Matching", "MDN"]: - log.info(f"Tuning {model_name} hyperparameters during fitting") - fitted_model, fold_tuned_params = model.fit( - train_data, - predictors, - imputed_variables, - weight_col=weight_col, - tune_hyperparameters=True, + if model_name == "Matching": + log.warning( + "Matching provides donor samples, not quantiles or class probabilities; skipping distributional scoring" ) + return None, None + metric_types = { + var: ("quantile_loss" if target_types[var] == "numeric" else "log_loss") + if target_types and var in target_types + else get_metric_for_variable_type(train_data[var], var) + for var in imputed_variables + } + if model_name == "QuantReg" and "log_loss" in metric_types.values(): + log.warning("QuantReg does not support categorical targets; skipping") + return None, None + params = dict(model_hyperparams or {}) + params["target_types"] = target_types + if model_name == "QuantReg": + params["quantiles"] = quantiles + if tune_hyperparameters and model_name in ["QRF", "MDN"]: + params["tune_hyperparameters"] = True + fitted = model.fit( + train_data, predictors, imputed_variables, weight_col=weight_col, **params + ) + if isinstance(fitted, tuple): + return fitted + return fitted, None - # Default fitting - else: - try: - log.info(f"Fitting {model_name} model with default parameters") - fitted_model = model.fit( - train_data, - predictors, - imputed_variables, - weight_col=weight_col, - ) - except ValueError as e: - if ( - "QuantReg does not support categorical" in str(e) - and model_name == "QuantReg" - ): - log.warning(f"QuantReg incompatible with variable types: {str(e)}") - return None, None - raise e - return fitted_model, fold_tuned_params +def _preprocessing_kwargs(preprocessing: Dict[str, str]) -> dict: + valid = { + "normalize": "normalize", + "log": "log_transform", + "asinh": "asinh_transform", + } + if set(preprocessing.values()) - set(valid): + raise ValueError("Unknown preprocessing transformation") + return { + argument: [col for col, transform in preprocessing.items() if transform == name] + or False + for name, argument in valid.items() + } def _compute_fold_loss_by_metric( @@ -278,70 +262,32 @@ def _compute_fold_loss_by_metric( result["quantile_loss"]["variables"].append(var) else: # log_loss - # Use probabilities if available, otherwise use class predictions - if test_probabilities and test_probabilities[var][fold_idx] is not None: - # Get probabilities and classes for this variable - test_prob_info = test_probabilities[var][fold_idx] - train_prob_info = train_probabilities[var][fold_idx] - - if ( - isinstance(test_prob_info, dict) - and "probabilities" in test_prob_info - ): - # Extract probabilities and classes - test_probs = test_prob_info["probabilities"] - train_probs = train_prob_info["probabilities"] - model_classes = test_prob_info["classes"] - - # Import the ordering function - from microimpute.comparisons.metrics import ( - order_probabilities_alphabetically, - ) - - # Order probabilities alphabetically - test_probs_ordered, alphabetical_labels = ( - order_probabilities_alphabetically(test_probs, model_classes) - ) - train_probs_ordered, _ = order_probabilities_alphabetically( - train_probs, model_classes - ) - - # Compute log loss with properly ordered probabilities - _, test_loss = compute_loss( - test_y_var, - test_probs_ordered, - "log_loss", - labels=alphabetical_labels, - ) - _, train_loss = compute_loss( - train_y_var, - train_probs_ordered, - "log_loss", - labels=alphabetical_labels, + if not test_probabilities or test_probabilities[var][fold_idx] is None: + raise ValueError( + f"Log loss for '{var}' requires predicted probabilities" + ) + losses = [] + labels = np.unique(np.concatenate([test_y_var, train_y_var])) + for truth, info in [ + (test_y_var, test_probabilities[var][fold_idx]), + (train_y_var, train_probabilities[var][fold_idx]), + ]: + classes = np.asarray(info["classes"]) + all_labels = np.union1d(labels, classes) + probabilities = np.zeros((len(truth), len(all_labels))) + for idx, label in enumerate(classes): + probabilities[:, np.flatnonzero(all_labels == label)[0]] = ( + np.asarray(info["probabilities"])[:, idx] ) + if len(all_labels) == 1: + losses.append(0.0) else: - # Fallback for old format or if probabilities not available - log.warning( - f"Probabilities not in expected format for variable {var}, using class predictions" - ) - labels = np.unique(np.concatenate([test_y_var, train_y_var])) - labels = np.sort(labels) # Ensure alphabetical order - _, test_loss = compute_loss( - test_y_var, test_pred_var, "log_loss", labels=labels - ) - _, train_loss = compute_loss( - train_y_var, train_pred_var, "log_loss", labels=labels + losses.append( + compute_loss( + truth, probabilities, "log_loss", labels=all_labels + )[1] ) - else: - # Fall back to using class predictions (less accurate) - labels = np.unique(np.concatenate([test_y_var, train_y_var])) - labels = np.sort(labels) # Ensure alphabetical order - _, test_loss = compute_loss( - test_y_var, test_pred_var, "log_loss", labels=labels - ) - _, train_loss = compute_loss( - train_y_var, train_pred_var, "log_loss", labels=labels - ) + test_loss, train_loss = losses if result["log_loss"]["test"] is None: result["log_loss"]["test"] = [] @@ -459,6 +405,8 @@ def cross_validate_model( random_state: Optional[int] = RANDOM_STATE, model_hyperparams: Optional[dict] = None, tune_hyperparameters: Optional[bool] = False, + preprocessing: Optional[Dict[str, str]] = None, + target_types: Optional[Dict[str, str]] = None, ) -> Union[Dict[str, Any], Tuple[Dict[str, Any], Dict]]: """Perform cross-validation with dual metric support. @@ -491,8 +439,10 @@ def cross_validate_model( validate_columns_exist(data, imputed_variables, "data") if weight_col: validate_columns_exist(data, [weight_col], "data") - if quantiles: - validate_quantiles(quantiles) + quantiles = QUANTILES if quantiles is None else quantiles + validate_quantiles(quantiles) + + data = declare_target_types(data, imputed_variables, target_types) # Set up parallel processing n_jobs = 1 if (Matching is not None and model_class == Matching) else -1 @@ -530,6 +480,9 @@ def cross_validate_model( model_hyperparams, tune_hyperparameters, variable_metrics, + preprocessing, + target_types, + random_state, ) for i, fold_pair in enumerate(fold_indices) ) @@ -668,8 +621,25 @@ def cross_validate_model( # Calculate means and stds across all quantiles mean_test = combined_df.loc["test"].mean() mean_train = combined_df.loc["train"].mean() - std_test = std_df.loc["test"].mean() - std_train = std_df.loc["train"].mean() + std_test = float( + np.std( + np.mean( + [metric_results[metric_type]["test"][q] for q in quantiles], + axis=0, + ) + ) + ) + std_train = float( + np.std( + np.mean( + [ + metric_results[metric_type]["train"][q] + for q in quantiles + ], + axis=0, + ) + ) + ) final_results[metric_type] = { "results": combined_df, # Single DataFrame with train/test rows @@ -699,26 +669,25 @@ def cross_validate_model( # Return results with optional hyperparameters if tune_hyperparameters and tuned_hyperparameters: - # Select best hyperparameters based on primary metric - primary_metric = ( - "quantile_loss" - if len(final_results["quantile_loss"]["variables"]) - >= len(final_results["log_loss"]["variables"]) - else "log_loss" + # Outer test folds estimate performance only. Select final parameters + # with a fresh internal tuning run on all available training rows. + tuning_data = data + if preprocessing: + tuning_data, _ = preprocess_data( + data, full_data=True, **_preprocessing_kwargs(preprocessing) + ) + _, best_hyperparams = _fit_model_for_fold( + create_distributional_model(model_class, seed=random_state), + model_class, + tuning_data, + predictors, + imputed_variables, + weight_col, + quantiles, + model_hyperparams, + True, + target_types, ) - - # Use median quantile (0.5) for selection - best_fold = 0 - best_loss = float("inf") - - if 0.5 in quantiles: - for fold_idx in range(n_splits): - fold_loss = metric_results[primary_metric]["test"][0.5][fold_idx] - if fold_loss < best_loss: - best_loss = fold_loss - best_fold = fold_idx - - best_hyperparams = tuned_hyperparameters.get(best_fold) return final_results, best_hyperparams else: return final_results diff --git a/microimpute/evaluations/predictor_analysis.py b/microimpute/evaluations/predictor_analysis.py index 938c08be..f30e449e 100644 --- a/microimpute/evaluations/predictor_analysis.py +++ b/microimpute/evaluations/predictor_analysis.py @@ -12,10 +12,7 @@ import pandas as pd from pydantic import validate_call from scipy.stats import spearmanr -from sklearn.feature_selection import ( - mutual_info_classif, - mutual_info_regression, -) +from sklearn.metrics import mutual_info_score from sklearn.model_selection import train_test_split from sklearn.preprocessing import LabelEncoder from tqdm.auto import tqdm @@ -23,6 +20,7 @@ from microimpute.comparisons.metrics import ( compute_loss, get_metric_for_variable_type, + order_probabilities_alphabetically, ) from microimpute.config import ( QUANTILES, @@ -31,6 +29,7 @@ VALIDATE_CONFIG, ) from microimpute.models import Imputer, ImputerResults +from microimpute.models.imputer import create_distributional_model from microimpute.utils.type_handling import ( DummyVariableProcessor, VariableTypeDetector, @@ -45,6 +44,7 @@ def compute_predictor_correlations( predictors: List[str], imputed_variables: Optional[List[str]] = None, method: str = "all", + n_bins: int = 10, ) -> Dict[str, pd.DataFrame]: """Compute correlation matrices between predictors using multiple methods. @@ -66,6 +66,12 @@ def compute_predictor_correlations( - "pearson": Only Pearson correlation - "spearman": Only Spearman correlation - "mutual_info": Only mutual information + n_bins: Maximum number of equal-frequency bins for continuous variables + in mutual information (default 10). Categorical levels are retained. + MI and both entropies use the same discretization and natural logs. + The score is MI / min(H(X), H(Y)), not Pearson correlation. It depends + on the binning and is not adjusted for chance. Constant variables + have score zero, including on the diagonal. Missing pairs are omitted. Returns: Dictionary containing correlation matrices: @@ -105,6 +111,8 @@ def compute_predictor_correlations( valid_methods = ["all", "pearson", "spearman", "mutual_info"] if method not in valid_methods: raise ValueError(f"Invalid method. Choose from: {valid_methods}") + if n_bins < 2: + raise ValueError("n_bins must be at least 2") # Prepare data - encode categorical variables detector = VariableTypeDetector() @@ -150,25 +158,14 @@ def compute_predictor_correlations( for i, pred1 in enumerate(predictors): for j, pred2 in enumerate(predictors): - if i == j: - mi_matrix.iloc[i, j] = 1.0 - elif j > i: - # Compute MI between pred1 and pred2 - mi_value = _compute_mutual_information( - data_encoded[pred1].values, - data_encoded[pred2].values, + if j >= i: + mi_normalized = _normalized_mutual_information( + data[pred1].values, + data[pred2].values, + categorical_mask[pred1], categorical_mask[pred2], + n_bins, ) - # Normalize by max possible MI (min of entropies) - # This makes it comparable to correlation coefficients - max_mi = min( - _compute_entropy(data_encoded[pred1].values), - _compute_entropy(data_encoded[pred2].values), - ) - if max_mi > 0: - mi_normalized = mi_value / max_mi - else: - mi_normalized = 0.0 mi_matrix.iloc[i, j] = mi_normalized mi_matrix.iloc[j, i] = mi_normalized @@ -187,7 +184,6 @@ def compute_predictor_correlations( ) # Prepare target variables - encode if categorical - targets_encoded = {} target_is_categorical = {} for target in imputed_variables: @@ -198,32 +194,17 @@ def compute_predictor_correlations( "bool", ] - if target_is_categorical[target]: - # Encode categorical targets - le = LabelEncoder() - targets_encoded[target] = le.fit_transform(data[target].astype(str)) - else: - targets_encoded[target] = data[target].values - # Compute MI between each predictor and each target for pred in predictors: for target in imputed_variables: - # Use encoded predictor values - pred_values = data_encoded[pred].values - target_values = targets_encoded[target] - - # Compute mutual information - mi_value = _compute_mutual_information( - pred_values, target_values, target_is_categorical[target] + mi_normalized = _normalized_mutual_information( + data[pred].values, + data[target].values, + categorical_mask[pred], + target_is_categorical[target], + n_bins, ) - # Optionally normalize by target entropy for comparability - target_entropy = _compute_entropy(target_values) - if target_entropy > 0: - mi_normalized = mi_value / target_entropy - else: - mi_normalized = 0.0 - pred_target_mi.loc[pred, target] = mi_normalized results["predictor_target_mi"] = pred_target_mi @@ -526,32 +507,40 @@ def progressive_predictor_inclusion( # Helper functions -def _compute_mutual_information( - x: np.ndarray, y: np.ndarray, y_is_categorical: bool +def _normalized_mutual_information( + x: np.ndarray, + y: np.ndarray, + x_is_categorical: bool, + y_is_categorical: bool, + n_bins: int, ) -> float: - """Compute mutual information between two variables.""" - # Remove any rows where either variable is NaN + """Normalize empirical MI using entropies of the same paired discretization.""" mask = ~(pd.isna(x) | pd.isna(y)) - x_clean = x[mask] - y_clean = y[mask] - - if len(x_clean) == 0: + if not np.any(mask): return 0.0 - # Reshape for sklearn - x_clean = x_clean.reshape(-1, 1) - - # Use appropriate MI function based on target type - if y_is_categorical: - mi = mutual_info_classif(x_clean, y_clean, random_state=RANDOM_STATE)[0] - else: - mi = mutual_info_regression(x_clean, y_clean, random_state=RANDOM_STATE)[0] - - return mi + def discretize(values: np.ndarray, categorical: bool) -> np.ndarray: + if categorical: + return pd.factorize(values)[0] + numeric = np.asarray(values, dtype=float) + if not np.isfinite(numeric).all(): + raise ValueError("Mutual information requires finite numeric values") + unique = np.unique(numeric) + if len(unique) <= n_bins: + return pd.factorize(numeric)[0] + # Equal-frequency bins are invariant to increasing unit transformations. + return np.asarray(pd.qcut(numeric, q=n_bins, labels=False, duplicates="drop")) + + x_codes = discretize(x[mask], x_is_categorical) + y_codes = discretize(y[mask], y_is_categorical) + normalizer = min(_compute_entropy(x_codes), _compute_entropy(y_codes)) + if normalizer <= 0: + return 0.0 + return float(np.clip(mutual_info_score(x_codes, y_codes) / normalizer, 0.0, 1.0)) def _compute_entropy(x: np.ndarray) -> float: - """Compute entropy of a variable.""" + """Compute empirical discrete entropy in nats.""" # Remove NaN values x_clean = x[~pd.isna(x)] @@ -563,7 +552,7 @@ def _compute_entropy(x: np.ndarray) -> float: probs = counts / counts.sum() # Compute entropy - entropy = -np.sum(probs * np.log2(probs + 1e-10)) + entropy = -np.sum(probs * np.log(probs)) return entropy @@ -581,7 +570,8 @@ def _evaluate_model_performance( """Train a model and evaluate its performance.""" try: # Initialize and fit the model - model = model_class() + model = create_distributional_model(model_class) + model.seed = random_state fitted_model = model.fit( X_train=train_data, predictors=predictors, @@ -590,7 +580,24 @@ def _evaluate_model_performance( ) # Get predictions - predictions = fitted_model.predict(test_data, quantiles) + predictions = fitted_model.predict(test_data, quantiles, return_probs=True) + metrics_by_variable = { + var: get_metric_for_variable_type(train_data[var], var) + for var in imputed_variables + } + # Constant targets bypass a classifier, but their predictive distribution + # is an exact point mass. Unseen evaluation classes receive probability 0. + for var, info in model.constant_targets.items(): + if metrics_by_variable[var] == "log_loss": + classes = np.unique( + np.concatenate([train_data[var].values, test_data[var].values]) + ) + predictions.setdefault("probabilities", {})[var] = { + "probabilities": np.tile( + (classes == info["value"]).astype(float), (len(test_data), 1) + ), + "classes": classes, + } # Compute losses losses = _compute_losses_from_predictions( @@ -598,6 +605,7 @@ def _evaluate_model_performance( true_data=test_data, imputed_variables=imputed_variables, quantiles=quantiles, + metrics_by_variable=metrics_by_variable, ) return losses @@ -608,10 +616,11 @@ def _evaluate_model_performance( def _compute_losses_from_predictions( - predictions: Dict[float, pd.DataFrame], + predictions: Dict[Any, Any], true_data: pd.DataFrame, imputed_variables: List[str], quantiles: List[float], + metrics_by_variable: Optional[Dict[str, str]] = None, ) -> Dict[str, float]: """Compute losses from model predictions.""" quantile_losses = [] @@ -630,14 +639,35 @@ def _compute_losses_from_predictions( pred_values = predictions[quantile][var] # Determine appropriate loss metric - metric_type = get_metric_for_variable_type(true_values, var) + metric_type = ( + metrics_by_variable[var] + if metrics_by_variable is not None + else get_metric_for_variable_type(true_values, var) + ) + + labels = None + if metric_type == "log_loss": + info = predictions.get("probabilities", {}).get(var) + if ( + not isinstance(info, dict) + or not {"probabilities", "classes"} <= info.keys() + ): + raise ValueError( + f"Log loss for {var!r} requires predicted probabilities and classes" + ) + pred_values, labels = order_probabilities_alphabetically( + np.asarray(info["probabilities"]), np.asarray(info["classes"]) + ) + else: + pred_values = pred_values.values # Compute loss (returns tuple of element-wise losses and mean) _, mean_loss = compute_loss( test_y=true_values.values, - imputations=pred_values.values, + imputations=pred_values, metric=metric_type, q=quantile, + labels=labels, ) if metric_type == "quantile_loss": diff --git a/microimpute/models/imputer.py b/microimpute/models/imputer.py index dae46fa6..732e59bf 100644 --- a/microimpute/models/imputer.py +++ b/microimpute/models/imputer.py @@ -20,6 +20,7 @@ from microimpute.utils.type_handling import ( DummyVariableProcessor, VariableTypeDetector, + declare_target_types, ) @@ -253,6 +254,7 @@ def fit( target_filters: Optional[ Dict[str, Union[str, np.ndarray, pd.Series, List[bool], Tuple[bool, ...]]] ] = None, + target_types: Optional[Dict[str, str]] = None, **kwargs: Any, ) -> Any: # Returns ImputerResults """Fit the model to the training data. @@ -261,7 +263,7 @@ def fit( X_train: DataFrame containing the training data. predictors: List of column names to use as predictors. imputed_variables: List of column names to impute. - weight_col: Optional name of the column or column array/series containing sampling weights. When provided, `X_train` will be sampled with replacement using this column as selection probabilities before fitting the model. + weight_col: Optional name or array of positive finite sample weights, passed to the learner's native weighted fit. skip_missing: If True, skip variables missing from training data with warning. If False, raise error for missing variables. not_numeric_categorical: Optional list of variable names that should be treated as numeric even if they would normally be detected as @@ -283,6 +285,20 @@ def fit( RuntimeError: If model fitting fails. NotImplementedError: If method is not implemented by subclass. """ + if len(predictors) != len(set(predictors)) or len(imputed_variables) != len( + set(imputed_variables) + ): + raise ValueError("Duplicate predictor or imputed variable names") + if not X_train.columns.is_unique: + raise ValueError("Duplicate DataFrame column names are not supported") + X_train = declare_target_types(X_train, imputed_variables, target_types) + not_numeric_categorical = list(not_numeric_categorical or []) + [ + name for name, kind in (target_types or {}).items() if kind == "numeric" + ] + self.categorical_targets = {} + self.boolean_targets = {} + self.numeric_targets = [] + self.constant_targets = {} original_predictors = predictors.copy() target_filters = target_filters or {} unknown_target_filters = set(target_filters) - set(imputed_variables) @@ -366,7 +382,7 @@ def fit( # weights — those then propagated into .sample() as NaN # probabilities or corrupted sample_weight passed to learners. weights_arr = np.asarray(weights, dtype=float) - invalid_mask = np.isnan(weights_arr) | (weights_arr <= 0) + invalid_mask = ~np.isfinite(weights_arr) | (weights_arr <= 0) if invalid_mask.any(): raise ValueError( "Weights must be positive and finite; found " @@ -577,6 +593,8 @@ def preprocess_data_types( processor = DummyVariableProcessor(self.logger) # This will only encode predictors in test data return processor.preprocess_predictors(data, predictors, []) + except ValueError: + raise except Exception as e: self.logger.error(f"Error during test data preprocessing: {str(e)}") raise RuntimeError("Failed to preprocess data types") from e @@ -654,3 +672,16 @@ def _predict( NotImplementedError: If method is not implemented by subclass. """ raise NotImplementedError("Subclasses must implement the predict method") + + +def create_distributional_model(model_class: type, **kwargs: Any) -> Imputer: + """Construct a model for marginal distribution evaluation. + + QRF chains model joint donor draws by default. Evaluating target-specific + quantiles instead requires independent fits conditional on original predictors. + """ + from microimpute.models.qrf import QRF + + if issubclass(model_class, QRF): + kwargs["sequential"] = False + return model_class(**kwargs) diff --git a/microimpute/models/matching.py b/microimpute/models/matching.py index 5b6af9a5..07fb1117 100644 --- a/microimpute/models/matching.py +++ b/microimpute/models/matching.py @@ -8,7 +8,14 @@ from microimpute.config import RANDOM_STATE, VALIDATE_CONFIG from microimpute.models.imputer import Imputer, ImputerResults -from microimpute.utils.statmatch_hotdeck import nnd_hotdeck_using_rpy2 + + +def nnd_hotdeck_using_rpy2(*args, **kwargs): + """Load the optional R bridge only when the default matcher is called.""" + from microimpute.utils.statmatch_hotdeck import nnd_hotdeck_using_rpy2 as match + + return match(*args, **kwargs) + MatchingHotdeckFn = Callable[ [ @@ -71,33 +78,50 @@ def __init__( self.matching_hotdeck = matching_hotdeck self.donor_data = donor_data self.hyperparameters = hyperparameters + self._rng = np.random.default_rng(seed) self.categorical_targets = categorical_targets or {} self.boolean_targets = boolean_targets or {} self.dummy_processor = dummy_processor + def _matching_kwargs(self) -> Dict[str, Any]: + """Advance a reproducible child-seed stream for the optional R bridge.""" + kwargs = dict(self.hyperparameters or {}) + if self.matching_hotdeck is nnd_hotdeck_using_rpy2: + kwargs["random_state"] = int(self._rng.integers(0, np.iinfo(np.int32).max)) + return kwargs + @validate_call(config=VALIDATE_CONFIG) def _predict( self, X_test: pd.DataFrame, quantiles: Optional[List[float]] = None, return_probs: bool = False, - ) -> Dict[float, pd.DataFrame]: + ) -> pd.DataFrame: """Predict imputed values using the matching model. Args: X_test: DataFrame containing the recipient data. - quantiles: List of quantiles to predict. - return_probs: If True, return one-hot probability vectors for matched categories. + quantiles: Unsupported; Matching returns donor draws. + return_probs: Unsupported; Matching does not estimate probabilities. Returns: - Dictionary mapping quantiles to imputed values. - If return_probs=True, includes 'probabilities' key with one-hot encodings. + DataFrame of donor draws, with n_failed_records in its attrs. Raises: ValueError: If model is not properly set up or input data is invalid. RuntimeError: If matching or prediction fails. + NotImplementedError: If quantiles or probabilities are requested. """ + if quantiles is not None: + raise NotImplementedError( + "Matching returns donor draws, not conditional quantiles. " + "Call predict without quantiles, or use QRF, OLS, or QuantReg." + ) + if return_probs: + raise NotImplementedError( + "Matching does not estimate class probabilities. Use QRF or OLS." + ) try: self.logger.info(f"Performing matching for {len(X_test)} recipient records") @@ -153,29 +177,24 @@ def _predict_single( X_test_copy: pd.DataFrame, quantiles: Optional[List[float]] = None, return_probs: bool = False, - ) -> Dict[float, pd.DataFrame]: + ) -> pd.DataFrame: """Perform matching on the full dataset without chunking.""" try: self.logger.info("Calling R-based hot deck matching function") - if self.hyperparameters: - fused0, fused1 = self.matching_hotdeck( - receiver=X_test_copy, - donor=self.donor_data, - matching_variables=self.predictors, - z_variables=self.imputed_variables, - **self.hyperparameters, - ) - else: - fused0, fused1 = self.matching_hotdeck( - receiver=X_test_copy, - donor=self.donor_data, - matching_variables=self.predictors, - z_variables=self.imputed_variables, - ) + fused0, fused1 = self.matching_hotdeck( + receiver=X_test_copy, + donor=self.donor_data, + matching_variables=self.predictors, + z_variables=self.imputed_variables, + **self._matching_kwargs(), + ) except Exception as matching_error: self.logger.error(f"Error in hot deck matching: {str(matching_error)}") raise RuntimeError("Hot deck matching failed") from matching_error + self.n_failed_records = int( + fused0[self.imputed_variables].isna().any(axis=1).sum() + ) return self._process_matching_results( fused0, X_test_copy, quantiles, return_probs ) @@ -186,7 +205,7 @@ def _predict_chunked( quantiles: Optional[List[float]], chunk_size: int, return_probs: bool = False, - ) -> Dict[float, pd.DataFrame]: + ) -> pd.DataFrame: """Perform matching using chunking for large datasets.""" all_results = [] @@ -202,21 +221,13 @@ def _predict_chunked( try: # Perform matching for this chunk - if self.hyperparameters: - fused0, fused1 = self.matching_hotdeck( - receiver=chunk_data, - donor=self.donor_data, - matching_variables=self.predictors, - z_variables=self.imputed_variables, - **self.hyperparameters, - ) - else: - fused0, fused1 = self.matching_hotdeck( - receiver=chunk_data, - donor=self.donor_data, - matching_variables=self.predictors, - z_variables=self.imputed_variables, - ) + fused0, fused1 = self.matching_hotdeck( + receiver=chunk_data, + donor=self.donor_data, + matching_variables=self.predictors, + z_variables=self.imputed_variables, + **self._matching_kwargs(), + ) # Store results with original indices chunk_results = pd.DataFrame(index=chunk_data.index) @@ -239,7 +250,6 @@ def _predict_chunked( # Combine all chunk results, preserving original order if all_results: 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 @@ -259,159 +269,32 @@ def _predict_chunked( else: raise RuntimeError("No chunk results were produced") - def _generate_one_hot_probabilities( - self, - variable: str, - matched_values: np.ndarray, - index: pd.Index, - categorical_targets: Dict, - boolean_targets: Dict, - ) -> Optional[Dict]: - """Generate one-hot probability matrix for categorical/boolean variables. - - Args: - variable: Name of the variable - matched_values: Array of matched category values - index: Index for the output DataFrame - categorical_targets: Dictionary of categorical target info - boolean_targets: Dictionary of boolean target info - - Returns: - Dict with 'probabilities' and 'classes' keys - """ - if variable not in categorical_targets and variable not in boolean_targets: - return None - - # Determine categories - if variable in boolean_targets: - categories = [False, True] - else: - categories = categorical_targets[variable].get("categories", []) - - if not categories: - return None - - # Create probability matrix (one-hot encoding) - n_samples = len(matched_values) - n_categories = len(categories) - prob_matrix = np.zeros((n_samples, n_categories)) - - # Set 1.0 for matched category - for idx, val in enumerate(matched_values): - try: - cat_idx = categories.index(val) - prob_matrix[idx, cat_idx] = 1.0 - except ValueError: - # If value not found in categories, default to first category - prob_matrix[idx, 0] = 1.0 - - return {"probabilities": prob_matrix, "classes": np.array(categories)} - def _process_matching_results( self, fused0: pd.DataFrame, X_test_copy: pd.DataFrame, quantiles: Optional[List[float]], return_probs: bool = False, - ) -> Dict[float, pd.DataFrame]: - """Process matching results into the expected output format.""" - try: - # Verify imputed variables exist in the result - missing_imputed = [ - var for var in self.imputed_variables if var not in fused0.columns - ] - if missing_imputed: - self.logger.error( - f"Imputed variables missing from matching result: {missing_imputed}" - ) - raise ValueError( - f"Matching failed to produce these variables: {missing_imputed}" - ) - - self.logger.info( - f"Matching completed, fused dataset has {len(fused0)} records" - ) - except Exception as convert_error: - self.logger.error( - f"Error converting matching results: {str(convert_error)}" + ) -> pd.DataFrame: + """Return donor draws, with an explicit count of unsuccessful matches.""" + if quantiles is not None: + raise NotImplementedError( + "Matching does not estimate conditional quantiles" ) - raise RuntimeError("Failed to process matching results") from convert_error - - # Create output dictionary with results - imputations: Dict[float, pd.DataFrame] = {} - prob_results = {} if return_probs else None - - # Get target type information if available - categorical_targets = getattr(self, "categorical_targets", {}) - boolean_targets = getattr(self, "boolean_targets", {}) - - try: - if quantiles: - self.logger.info(f"Creating imputations for {len(quantiles)} quantiles") - # For each quantile, return a DataFrame with all imputed variables - for q in quantiles: - imputed_df = pd.DataFrame(index=X_test_copy.index) - for variable in self.imputed_variables: - self.logger.debug( - f"Adding result for imputed variable {variable} at quantile {q}" - ) - imputed_df[variable] = fused0[variable].values - - # Generate one-hot probabilities if requested - if return_probs and prob_results is not None: - prob_df = self._generate_one_hot_probabilities( - variable, - fused0[variable].values, - X_test_copy.index, - categorical_targets, - boolean_targets, - ) - if prob_df is not None: - prob_results[variable] = prob_df - - imputations[q] = imputed_df - - # Add probabilities to results if requested - if return_probs and prob_results: - imputations["probabilities"] = prob_results - - return imputations - else: - # If no quantiles specified, use a default one - q_default = 0.5 - self.logger.info( - f"Creating imputation for default quantile {q_default}" - ) - imputed_df = pd.DataFrame(index=X_test_copy.index) - for variable in self.imputed_variables: - self.logger.info(f"Imputing variable {variable}") - imputed_df[variable] = fused0[variable].values - - # Generate one-hot probabilities if requested - if return_probs and prob_results is not None: - prob_df = self._generate_one_hot_probabilities( - variable, - fused0[variable].values, - X_test_copy.index, - categorical_targets, - boolean_targets, - ) - if prob_df is not None: - prob_results[variable] = prob_df - - imputations[q_default] = imputed_df - - # Add probabilities to results if requested - if return_probs and prob_results: - # Return dict with both quantile predictions and probabilities - imputations["probabilities"] = prob_results - return imputations - else: - # Return just the DataFrame for the single quantile - return imputations[q_default] - except Exception as output_error: - self.logger.error(f"Error creating output imputations: {str(output_error)}") - raise RuntimeError("Failed to create output imputations") from output_error + if return_probs: + raise NotImplementedError("Matching does not estimate class probabilities") + missing = [v for v in self.imputed_variables if v not in fused0] + if missing: + raise ValueError(f"Matching failed to produce these variables: {missing}") + if len(fused0) != len(X_test_copy): + raise ValueError("Matching must return one record per receiver") + output = pd.DataFrame( + {v: fused0[v].to_numpy() for v in self.imputed_variables}, + index=X_test_copy.index, + ) + self.n_failed_records = int(output.isna().any(axis=1).sum()) + output.attrs["n_failed_records"] = self.n_failed_records + return output class Matching(Imputer): @@ -427,6 +310,7 @@ def __init__( self, matching_hotdeck: MatchingHotdeckFn = nnd_hotdeck_using_rpy2, log_level: Optional[str] = "WARNING", + seed: int = RANDOM_STATE, ) -> None: """Initialize the matching model. @@ -437,7 +321,7 @@ def __init__( Raises: ValueError: If matching_hotdeck is not callable """ - super().__init__(log_level=log_level) + super().__init__(seed=seed, log_level=log_level) self.log_level = log_level self.logger.debug("Initializing Matching imputer") @@ -472,9 +356,9 @@ def _fit( imputed_variables: List of column names to impute. sample_weight: Optional per-row sample weights for the donor dataset. When provided, weights are passed to R StatMatch's - ``NND.hotdeck`` via ``weight.don`` so that donor records are - matched in proportion to their survey weights rather than - uniformly. + ``RANDwNND.hotdeck`` via a donor weight column. By default, + donors tied at the minimum distance are sampled in proportion + to their weights. matching_kwargs: Additional keyword arguments for hyperparameter tuning of the matching function. @@ -500,6 +384,9 @@ def _fit( data=X_train, predictors=predictors, imputed_variables=imputed_variables, + matching_kwargs=matching_kwargs, + categorical_targets=categorical_targets, + boolean_targets=boolean_targets, ) self.logger.info(f"Best hyperparameters: {best_params}") @@ -516,7 +403,7 @@ def _fit( dummy_processor=getattr(self, "dummy_processor", None), seed=self.seed, log_level=self.log_level, - hyperparameters=best_params, + hyperparameters={**matching_kwargs, **best_params}, ), best_params, ) @@ -553,190 +440,97 @@ def _tune_hyperparameters( data: pd.DataFrame, predictors: List[str], imputed_variables: List[str], + matching_kwargs: Optional[Dict[str, Any]] = None, + categorical_targets: Optional[Dict[str, Dict]] = None, + boolean_targets: Optional[Dict[str, Dict]] = None, ) -> Dict[str, Any]: - """Tune hyperparameters for the Matching model using Optuna with CV. + """Tune donor-draw accuracy; failed matches prune the entire trial. - Uses cross-validation and quantile loss for robust hyperparameter selection. - - Args: - data: DataFrame containing the training data. - predictors: List of column names to use as predictors. - imputed_variables: List of column names to impute. - - Returns: - Dictionary of tuned hyperparameters. + Numeric donor draws are assessed with absolute error normalized by the + donor training standard deviation; categorical draws use misclassification + rate. These are internal tuning criteria, not predictive distribution scores. """ import optuna from sklearn.model_selection import KFold - from microimpute.comparisons.metrics import compute_loss - - optuna.logging.set_verbosity(optuna.logging.WARNING) - - # Use 3-fold CV with 10 trials - n_cv_folds = 3 - n_trials = 10 - - # Set up CV folds - kf = KFold(n_splits=n_cv_folds, shuffle=True, random_state=self.seed) - - self.logger.info( - f"Tuning Matching hyperparameters with {n_cv_folds}-fold CV and {n_trials} trials" - ) + kf = KFold(n_splits=3, shuffle=True, random_state=self.seed) + fixed_kwargs = dict(matching_kwargs or {}) + weights = fixed_kwargs.pop("donor_sample_weight", None) + discrete_targets = set(categorical_targets or {}) | set(boolean_targets or {}) def objective(trial: optuna.Trial) -> float: + # NND.hotdeck's k controls donor re-use only under constrained + # matching; it is not a nearest-neighbor count. Do not tune a no-op. params = { "dist_fun": trial.suggest_categorical( "dist_fun", - [ - "Manhattan", - "Euclidean", - "Mahalanobis", - "Gower", - "minimax", - ], - ), - "k": trial.suggest_int("k", 1, 10), + ["Manhattan", "Euclidean", "Mahalanobis", "Gower", "minimax"], + ) } - - # Detect variable types for appropriate metric selection - from microimpute.comparisons.metrics import ( - get_metric_for_variable_type, - ) - - variable_metrics = {} - for var in imputed_variables: - variable_metrics[var] = get_metric_for_variable_type(data[var], var) - - # Track errors across CV folds fold_errors = [] - - # Perform CV + # Common random numbers across trials make parameter comparisons + # reproducible without rewarding a different random donor sequence. + trial_rng = np.random.default_rng(self.seed) for fold_idx, (train_idx, val_idx) in enumerate(kf.split(data)): - X_train_fold = data.iloc[train_idx] - X_val_fold = data.iloc[val_idx] - - # Track errors for all variables in this fold - var_errors = [] - + donor = data.iloc[train_idx] + receiver = data.iloc[val_idx].drop(columns=imputed_variables) + call_kwargs = {**fixed_kwargs, **params} + if weights is not None: + call_kwargs["donor_sample_weight"] = np.asarray(weights)[train_idx] + predicted = [] + for start in range(0, len(receiver), 1000): + chunk = receiver.iloc[start : start + 1000] + if self.matching_hotdeck is nnd_hotdeck_using_rpy2: + call_kwargs["random_state"] = int( + trial_rng.integers(0, np.iinfo(np.int32).max) + ) + try: + fused, _ = self.matching_hotdeck( + receiver=chunk, + donor=donor, + matching_variables=predictors, + z_variables=imputed_variables, + **call_kwargs, + ) + if ( + len(fused) != len(chunk) + or fused[imputed_variables].isna().any().any() + ): + raise ValueError("Matching returned incomplete predictions") + predicted.append( + fused[imputed_variables].reset_index(drop=True) + ) + except Exception as error: + self.logger.warning( + f"Matching failed on fold {fold_idx} chunk {start}: {error}. Pruning trial." + ) + raise optuna.TrialPruned() from error + predictions = pd.concat(predicted, ignore_index=True) + errors = [] for var in imputed_variables: - y_val = X_val_fold[var] - X_val_var = X_val_fold.copy().drop(var, axis=1) - - # Determine if chunking is needed for hyperparameter tuning - chunk_size = 1000 # Smaller chunks for tuning - total_size = len(X_train_fold) * len(X_val_var) - use_chunking = ( - len(X_val_var) > chunk_size - or total_size > 25_000_000 # Lower threshold for tuning - ) - - if use_chunking: - # Perform chunked matching for hyperparameter tuning - y_pred_chunks = [] - y_val_chunks = [] - - for i in range(0, len(X_val_var), chunk_size): - chunk_end = min(i + chunk_size, len(X_val_var)) - chunk_data = X_val_var.iloc[i:chunk_end] - chunk_y_val = y_val.iloc[i:chunk_end] - - try: - fused0, fused1 = self.matching_hotdeck( - receiver=chunk_data, - donor=X_train_fold, - matching_variables=predictors, - z_variables=[var], - **params, - ) - y_pred_chunks.append(fused0[var].values) - 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) - y_val_combined = np.concatenate(y_val_chunks) + actual = data.iloc[val_idx][var].to_numpy() + estimate = predictions[var].to_numpy() + if var in discrete_targets: + errors.append(float(np.mean(actual != estimate))) else: - # Perform single matching - try: - fused0, fused1 = self.matching_hotdeck( - receiver=X_val_var, - donor=X_train_fold, - matching_variables=predictors, - z_variables=[var], - **params, + if not np.isfinite(estimate.astype(float)).all(): + raise optuna.TrialPruned( + "Matching returned nonfinite numeric predictions" ) - y_pred = fused0[var].values - 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] - - if metric == "quantile_loss": - _, loss_value = compute_loss( - y_val_combined.flatten(), - y_pred.flatten(), - "quantile_loss", - q=0.5, - ) - # Normalize by variable's standard deviation - std = np.std(y_val_combined.flatten()) - normalized_loss = loss_value / std if std > 0 else loss_value - else: # log_loss for categorical/boolean - _, loss_value = compute_loss( - y_val_combined.flatten(), - y_pred.flatten(), - "log_loss", + scale = float(donor[var].std(ddof=0)) + errors.append( + float(np.mean(np.abs(actual - estimate))) / (scale or 1.0) ) - # Log loss is already normalized - normalized_loss = loss_value - - var_errors.append(normalized_loss) - - # Average across variables for this fold - if var_errors: - fold_errors.append(np.mean(var_errors)) - - # Return mean error across all CV folds - return np.mean(fold_errors) if fold_errors else float("inf") + fold_errors.append(float(np.mean(errors))) + return float(np.mean(fold_errors)) study = optuna.create_study( - direction="minimize", - sampler=optuna.samplers.TPESampler(seed=self.seed), + direction="minimize", sampler=optuna.samplers.TPESampler(seed=self.seed) ) - - # Suppress warnings during optimization - import os - - os.environ["PYTHONWARNINGS"] = "ignore" - - study.optimize(objective, n_trials=n_trials) - - best_value = study.best_value + study.optimize(objective, n_trials=10) + if not any(t.state == optuna.trial.TrialState.COMPLETE for t in study.trials): + raise ValueError("No matching hyperparameter trial succeeded") self.logger.info( - f"Matching - Lowest average normalized quantile loss ({n_cv_folds}-fold CV): {best_value}" + f"Matching best normalized donor-draw error: {study.best_value}" ) - - best_params = study.best_params - self.logger.info(f"Matching - Best hyperparameters found: {best_params}") - - return best_params + return study.best_params diff --git a/microimpute/models/ols.py b/microimpute/models/ols.py index 6cd6affa..df6f0bf4 100644 --- a/microimpute/models/ols.py +++ b/microimpute/models/ols.py @@ -9,7 +9,7 @@ from scipy.stats import norm from sklearn.linear_model import LogisticRegression -from microimpute.config import VALIDATE_CONFIG +from microimpute.config import RANDOM_STATE, VALIDATE_CONFIG from microimpute.models.imputer import Imputer, ImputerResults @@ -167,17 +167,38 @@ def fit( the weights. """ self.output_column = y.name - X_with_const = sm.add_constant(X) + if ( + not np.isfinite(X.to_numpy(dtype=float)).all() + or not np.isfinite(y.to_numpy(dtype=float)).all() + ): + raise ValueError("OLS training predictors and targets must be finite") + X_with_const = sm.add_constant(X, has_constant="add") if sample_weight is not None: weights = np.asarray(sample_weight, dtype=float) + if ( + weights.shape != (len(y),) + or not np.isfinite(weights).all() + or (weights <= 0).any() + ): + raise ValueError( + "OLS sample weights must be positive, finite, and aligned to training rows" + ) + # weight_col represents relative survey mass, not observation + # precision. Unit-mean weights leave coefficient covariance unchanged + # but put residual scale back in outcome units, invariant to w -> c*w. + weights = weights / weights.mean() self.model = sm.WLS(y, X_with_const, weights=weights).fit() else: self.model = sm.OLS(y, X_with_const).fit() self.scale = self.model.scale + if not np.isfinite(self.scale): + raise ValueError( + "OLS residual variance is undefined; provide more observations than fitted parameters" + ) def predict(self, X: pd.DataFrame) -> np.ndarray: """Predict using OLS model.""" - X_with_const = sm.add_constant(X) + X_with_const = sm.add_constant(X, has_constant="add") return self.model.predict(X_with_const) @@ -234,7 +255,9 @@ def _predict_variable( # the residual std and under-dispersed imputations for test # rows far from the training centroid; at extreme quantiles # (0.01, 0.99) the under-dispersion is material. - X_test_with_const = sm.add_constant(X_test[self.predictors]) + X_test_with_const = sm.add_constant( + X_test[self.predictors], has_constant="add" + ) prediction = model.model.get_prediction(X_test_with_const) # var_pred_mean is the leverage term (x' (X'X)^-1 x) * scale; # adding model.scale (residual variance) gives the prediction @@ -291,6 +314,7 @@ def __init__( log_level, ) self.models = models + self.rng = np.random.default_rng(seed) self.categorical_targets = categorical_targets or {} self.boolean_targets = boolean_targets or {} self.constant_targets = constant_targets or {} @@ -319,6 +343,8 @@ def _predict( Raises: RuntimeError: If prediction fails. """ + if not np.isfinite(X_test[self.predictors].to_numpy(dtype=float)).all(): + raise ValueError("OLS prediction predictors must be finite") try: # Create output dictionary with results imputations: Dict[float, pd.DataFrame] = {} @@ -327,7 +353,7 @@ def _predict( if quantiles: if random_quantile_sample: self.logger.warning( - f"Predicting at random quantiles sampled from a beta distribution is not possible when specified quantiles are provided." + "Explicit quantiles take precedence; ignoring random_quantile_sample." ) self.logger.info( f"Predicting at {len(quantiles)} quantiles: {quantiles}" @@ -341,7 +367,7 @@ def _predict( variable, X_test, q, - random_quantile_sample, + False, return_probs, prob_results, ) @@ -387,7 +413,6 @@ def _predict_quantile( se: Any, mean_quantile: float, random_sample: bool, - count_samples: int = 10, ) -> pd.Series: """Predict values at a specified quantile. @@ -399,8 +424,6 @@ def _predict_quantile( mean_quantile: Quantile to predict (the quantile affects the center of the beta distribution from which to sample when imputing each data point). random_sample: If True, use random quantile sampling for prediction. - count_samples: Number of quantile samples to generate when - random_sample is True. Returns: Series of predicted values at the specified quantile, indexed to @@ -423,23 +446,14 @@ def _predict_quantile( self.logger.info( f"Predicting at random quantiles sampled from a beta distribution with mean quantile {q_clipped}" ) - random_generator = np.random.default_rng(self.seed) - - # Calculate alpha parameter for beta distribution (q is - # safely in (0,1) after clipping). + # A persistent stream advances across rows, variables and calls. + # At q=0.5, Beta(1, 1) yields independent normal residuals. a = q_clipped / (1 - q_clipped) - - # Generate count_samples beta distributed values with parameter a - beta_samples = random_generator.beta(a, 1, size=count_samples) - - # Convert to normal quantiles using norm.ppf - normal_quantiles = norm.ppf(beta_samples) - - # For each mean prediction, randomly select one of the quantiles - sampled_indices = random_generator.integers( - 0, count_samples, size=len(mean_preds) + beta_samples = self.rng.beta(a, 1, size=len(mean_preds)) + beta_samples = np.clip( + beta_samples, np.finfo(float).eps, 1 - np.finfo(float).eps ) - selected_quantiles = normal_quantiles[sampled_indices] + selected_quantiles = norm.ppf(beta_samples) # Adjust each mean prediction by the sampled quantile # times its per-row SE (or scalar SE if se is a float). @@ -470,9 +484,11 @@ class OLS(Imputer): distributed residuals. """ - def __init__(self, log_level: Optional[str] = "WARNING") -> None: + def __init__( + self, log_level: Optional[str] = "WARNING", seed: int = RANDOM_STATE + ) -> None: """Initialize the OLS model.""" - super().__init__(log_level=log_level) + super().__init__(seed=seed, log_level=log_level) self.model = None self.log_level = log_level self.logger.debug("Initializing OLS imputer") @@ -507,6 +523,15 @@ def _fit( Raises: RuntimeError: If model fitting fails. """ + if X_train[predictors + imputed_variables].isna().any().any(): + raise ValueError( + "OLS training predictors and targets must not contain missing values" + ) + numeric_data = X_train[predictors + imputed_variables].select_dtypes( + include=["number", "bool"] + ) + if not np.isfinite(numeric_data.to_numpy(dtype=float)).all(): + raise ValueError("OLS training predictors and targets must be finite") try: self.logger.info(f"Fitting OLS model with {len(predictors)} predictors") diff --git a/microimpute/models/qrf.py b/microimpute/models/qrf.py index da9f169d..d9e71aff 100644 --- a/microimpute/models/qrf.py +++ b/microimpute/models/qrf.py @@ -1,6 +1,7 @@ """Quantile Regression Forest imputation model with sequential imputation.""" import gc +import hashlib import time from typing import Any, Dict, List, Optional, Tuple @@ -10,7 +11,7 @@ from quantile_forest import RandomForestQuantileRegressor from sklearn.ensemble import RandomForestClassifier -from microimpute.config import RANDOM_STATE, VALIDATE_CONFIG +from microimpute.config import DEFAULT_MODEL_PARAMS, RANDOM_STATE, VALIDATE_CONFIG from microimpute.models.imputer import Imputer, ImputerResults try: @@ -25,6 +26,7 @@ def _get_sequential_predictors( predictors: List[str], imputed_variables: List[str], current_variable_index: int, + sequential: bool = True, ) -> List[str]: """Get the predictor set for sequential imputation. @@ -36,7 +38,9 @@ def _get_sequential_predictors( Returns: List of predictor columns including previously imputed variables """ - return predictors + imputed_variables[:current_variable_index] + return predictors + ( + imputed_variables[:current_variable_index] if sequential else [] + ) class _RandomForestClassifierModel: @@ -165,6 +169,7 @@ def __init__(self, seed: int, logger): self.seed = seed self.logger = logger self.qrf = None + self._weighted_leaves = None self.output_column = None self.feature_columns: List[str] = [] # Create the RNG once at construction so that repeated predict() @@ -186,10 +191,12 @@ def fit( Args: X: Predictor DataFrame (preprocessed). y: Target Series. - sample_weight: Optional per-row sample weights, passed directly to - the underlying ``RandomForestQuantileRegressor.fit`` so each - row contributes to the weighted-survey estimator rather than - being treated as a bootstrap-resample probability. + sample_weight: Optional positive per-row survey weights. Weights + affect tree splits and the conditional empirical CDF. Each + tree normalizes its bootstrap donor weights within the leaf; + quantiles invert the average of these tree CDFs without + interpolation. All in-bag donors are retained in the weighted + CDF even if max_samples_leaf limits the upstream leaf storage. """ self.output_column = y.name @@ -197,7 +204,7 @@ def fit( # we set them explicitly below. qrf_kwargs_filtered = { k: v - for k, v in qrf_kwargs.items() + for k, v in {**DEFAULT_MODEL_PARAMS["qrf"], **qrf_kwargs}.items() if k not in ("random_state", "sample_weight") } @@ -210,6 +217,80 @@ def fit( fit_kwargs["sample_weight"] = np.asarray(sample_weight, dtype=float) self.feature_columns = list(X.columns) self.qrf.fit(X, y.values.ravel(), **fit_kwargs) + self._weighted_leaves = None + if sample_weight is not None: + self._fit_weighted_leaves(X, y, np.asarray(sample_weight, dtype=float)) + + def _fit_weighted_leaves( + self, X: pd.DataFrame, y: pd.Series, sample_weight: np.ndarray + ) -> None: + """Store survey-weighted empirical leaf distributions. + + quantile-forest 1.4 with sklearn <1.9 uses positive sample weights + for splits but not their magnitudes for quantiles. Retain its in-bag + donor multiplicities, normalize survey weights within each leaf, and + average the tree CDFs (Meinshausen, 2006, equations 4--6). + """ + self._weighted_y = np.asarray(y, dtype=float) + self._weighted_order = np.argsort(self._weighted_y) + leaves = self.qrf.apply(X) + self._weighted_leaves = [] + for tree, sampled in enumerate(self.qrf.estimators_samples_): + unique, counts = np.unique(sampled, return_counts=True) + tree_leaves = leaves[unique, tree] + masses = sample_weight[unique] * counts + mapping = {} + for leaf in np.unique(tree_leaves): + mask = tree_leaves == leaf + leaf_mass = masses[mask] + mapping[int(leaf)] = (unique[mask], leaf_mass / leaf_mass.sum()) + self._weighted_leaves.append(mapping) + + def predict_quantiles_per_row( + self, X: pd.DataFrame, quantiles: np.ndarray + ) -> pd.Series: + """Evaluate one exact conditional quantile per row, without sampling.""" + X = self._align_features(X) + quantiles = np.asarray(quantiles, dtype=float) + if quantiles.shape != (len(X),) or not np.isfinite(quantiles).all(): + raise ValueError("Provide one finite quantile per prediction row") + if ((quantiles < 0) | (quantiles > 1)).any(): + raise ValueError("Quantiles must be between zero and one") + values = np.empty(len(X)) + if self._weighted_leaves is not None: + query_leaves = self.qrf.apply(X) + # Cache consecutive identical leaf signatures (common in homogeneous + # prediction groups), without retaining an unbounded query cache. + previous = None + for row, signature in enumerate(query_leaves): + if previous is None or not np.array_equal(previous, signature): + mass = np.zeros(len(self._weighted_y)) + for tree, leaf in enumerate(signature): + indices, weights = self._weighted_leaves[tree][int(leaf)] + mass[indices] += weights + ordered_mass = mass[self._weighted_order] + supported = ordered_mass > 0 + support = self._weighted_y[self._weighted_order][supported] + cumulative = np.cumsum(ordered_mass[supported]) + cumulative /= cumulative[-1] + previous = signature.copy() + position = np.searchsorted(cumulative, quantiles[row], side="left") + values[row] = support[min(position, len(support) - 1)] + else: + # The upstream API accepts a shared grid, so bound intermediate + # storage by querying small batches then selecting each row's q. + for start in range(0, len(X), 128): + stop = min(start + 128, len(X)) + grid, columns = np.unique(quantiles[start:stop], return_inverse=True) + predicted = np.asarray( + self.qrf.predict( + X.iloc[start:stop], + quantiles=grid.tolist(), + weighted_leaves=True, + ) + ).reshape(stop - start, -1) + values[start:stop] = predicted[np.arange(stop - start), columns] + return pd.Series(values, index=X.index, name=self.output_column) def _align_features(self, X: pd.DataFrame) -> pd.DataFrame: """Reorder prediction features to the fitted QRF column contract.""" @@ -253,9 +334,9 @@ def predict( # QRF directly so that for any row i, # prediction(q_low) <= prediction(q_mid) <= prediction(q_high). if exact_quantile is not None: - pred = self.qrf.predict(X, quantiles=[float(exact_quantile)]) - pred = np.asarray(pred).reshape(len(X), -1)[:, 0] - return pd.Series(pred, index=X.index, name=self.output_column) + return self.predict_quantiles_per_row( + X, np.full(len(X), float(exact_quantile)) + ) # Stochastic path: draw one continuous quantile per row from a Beta # distribution centred at ``mean_quantile`` (Beta(a,1) with @@ -272,29 +353,7 @@ def predict( # from ``self.seed``, collapsing variance to zero). continuous_quantiles = self._rng.beta(a, 1, size=len(X)) - # Bucket continuous quantiles onto a fine symmetric grid covering the - # full open interval (0, 1). Using round() (not floor) keeps the - # mapping centred on the intended quantile, so the empirical mean of - # mapped quantiles ≈ ``mean_quantile``. We avoid exact 0 and 1 because - # QRF cannot extrapolate beyond observed extremes. - grid_size = max(int(count_samples), 101) - eps = 1.0 / (grid_size + 1) - quantile_grid = np.linspace(eps, 1.0 - eps, grid_size) - # Round (not floor) onto the grid to eliminate the low-side bias. - grid_indices = np.clip( - np.rint(continuous_quantiles * (grid_size - 1)).astype(int), - 0, - grid_size - 1, - ) - - pred = self.qrf.predict(X, quantiles=list(quantile_grid)) - pred = np.asarray(pred) - if pred.ndim == 2: - predictions = pred[np.arange(len(X)), grid_indices] - else: - predictions = pred[np.arange(len(X)), :, grid_indices] - - return pd.Series(predictions, index=X.index, name=self.output_column) + return self.predict_quantiles_per_row(X, continuous_quantiles) class QRFResults(ImputerResults): @@ -317,6 +376,7 @@ def __init__( constant_targets: Optional[Dict[str, Dict]] = None, dummy_processor: Optional[Any] = None, log_level: Optional[str] = "WARNING", + sequential: bool = True, ) -> None: """Initialize the QRF results. @@ -342,6 +402,7 @@ def __init__( log_level, ) self.models = models + self.sequential = sequential self.categorical_targets = categorical_targets or {} self.boolean_targets = boolean_targets or {} self.constant_targets = constant_targets or {} @@ -388,6 +449,24 @@ def _encode_imputed_variable( return data + def _predict_quantiles_per_row( + self, X_test: pd.DataFrame, variable: str, quantiles: np.ndarray + ) -> np.ndarray: + """Exact row-specific quantiles for a single numeric component model.""" + from microimpute.models.imputer import _ConstantValueModel + + if self.imputed_variables != [variable]: + raise NotImplementedError("Row-specific quantiles require a single target") + prepared, _ = self.preprocess_data_types( + X_test, self.original_predictors, self.dummy_processor + ) + model = self.models[variable] + if isinstance(model, _ConstantValueModel): + return model.predict(prepared).to_numpy(dtype=float) + if not isinstance(model, _QRFModel): + raise NotImplementedError("Quantile components must be numeric") + return model.predict_quantiles_per_row(prepared, quantiles).to_numpy() + @validate_call(config=VALIDATE_CONFIG) def _predict( self, @@ -400,8 +479,9 @@ def _predict( Args: X_test: DataFrame containing the test data. - quantiles: List of quantiles to predict (the quantile affects the - center of the beta distribution from which to sample when imputing each data point). + quantiles: Exact quantiles conditional on the original predictors. + Multiple targets require fitting with sequential=False: chaining + conditional quantiles is not a marginal quantile calculation. mean_quantile: The mean quantile to used for prediction if quantiles are not provided. return_probs: If True, return probability distributions for categorical variables. @@ -413,6 +493,15 @@ def _predict( Raises: RuntimeError: If prediction fails. """ + if ( + quantiles is not None + and self.sequential + and len(self.imputed_variables) > 1 + ): + raise NotImplementedError( + "Marginal quantiles for sequential multi-target QRF are not available; " + "fit QRF(sequential=False) or fit a single target" + ) try: # Create output dictionary with results imputations: Dict[float, pd.DataFrame] = {} @@ -453,7 +542,10 @@ def _predict( # Build predictor set: original predictors + previously imputed variables var_predictors = _get_sequential_predictors( - self.predictors, self.imputed_variables, i + self.predictors, + self.imputed_variables, + i, + sequential=self.sequential, ) # Get properly encoded predictor columns @@ -517,7 +609,7 @@ def _predict( # (the user wants to inspect specific quantiles, e.g. # for prediction intervals), we query the QRF at # exactly ``q`` per row — NO beta sampling. This - # guarantees row-level monotonicity across quantiles. + # is monotone for fixed predictors (single target). # Otherwise, sample stochastically around ``q`` (the # beta-mean default for imputation variance). if quantiles: @@ -608,6 +700,7 @@ def __init__( cleanup_interval: int = 10, max_train_samples: Optional[int] = None, seed: Optional[int] = RANDOM_STATE, + sequential: bool = True, ) -> None: """Initialize the QRF model. @@ -619,12 +712,16 @@ def __init__( max_train_samples: If set, subsample X_train to at most this many rows before fitting. Reduces memory and training time while preserving sequential covariance structure. + sequential: If True, condition each target on earlier targets for + stochastic joint draws. If False, fit each target using only the + original predictors, supporting marginal conditional quantiles. seed: Base random seed. Each imputed variable is given a distinct seed derived from it, so variables imputed together draw independently. Pass None for non-reproducible draws. """ super().__init__(log_level=log_level, seed=seed) self.models = {} + self.sequential = sequential self.log_level = log_level self.memory_efficient = memory_efficient self.batch_size = batch_size @@ -710,6 +807,15 @@ def _seed_for_variable(self, variable: str) -> Optional[int]: """ if self.seed is None: return None + if not self.sequential: + # A stable variable-name seed makes independent fits invariant to + # target ordering and to adding/removing other target variables. + identity = int.from_bytes( + hashlib.blake2s(variable.encode(), digest_size=4).digest(), "little" + ) + return int( + np.random.SeedSequence([self.seed, identity]).generate_state(1)[0] + ) try: variable_offset = (self.imputed_variables or []).index(variable) except ValueError: @@ -822,12 +928,7 @@ def _target_fit_data( self.max_train_samples is not None and len(target_train) > self.max_train_samples ): - try: - variable_offset = (self.imputed_variables or []).index(variable) - except ValueError: - variable_offset = 0 - seed = None if self.seed is None else self.seed + variable_offset - rng = np.random.default_rng(seed) + rng = np.random.default_rng(self._seed_for_variable(variable)) sel = rng.choice( len(target_train), size=self.max_train_samples, replace=False ) @@ -993,7 +1094,10 @@ def _fit( # Build predictor set: original predictors + previously imputed variables current_predictors = _get_sequential_predictors( - predictors, imputed_variables, i + predictors, + imputed_variables, + i, + sequential=self.sequential, ) # Get properly encoded predictor columns @@ -1084,6 +1188,7 @@ def _fit( constant_targets=constant_targets, dummy_processor=getattr(self, "dummy_processor", None), seed=self.seed, + sequential=self.sequential, ), qrf_kwargs, ) @@ -1156,7 +1261,7 @@ def _fit( # Build predictor set: original predictors + previously imputed variables current_predictors = _get_sequential_predictors( - predictors, imputed_variables, i + predictors, imputed_variables, i, sequential=self.sequential ) # Get properly encoded predictor columns @@ -1255,6 +1360,7 @@ def _fit( constant_targets=constant_targets, dummy_processor=getattr(self, "dummy_processor", None), seed=self.seed, + sequential=self.sequential, log_level=self.log_level, ) except Exception as e: @@ -1299,7 +1405,7 @@ def _fit_variable_batch( # Build predictor set: original predictors + previously imputed variables current_predictors = _get_sequential_predictors( - predictors, imputed_variables, i + predictors, imputed_variables, i, sequential=self.sequential ) dummy_processor = getattr(self, "dummy_processor", None) encoded_predictors = self._get_encoded_predictors( @@ -1338,7 +1444,7 @@ def _fit_variable_batch( self.logger.info(f" ✓ Success: {variable} fitted in {var_time:.2f}s") # Get model complexity metrics if available - if hasattr(model.qrf, "n_estimators"): + if hasattr(getattr(model, "qrf", None), "n_estimators"): self.logger.info( f" Model complexity: {model.qrf.n_estimators} trees" ) @@ -1472,7 +1578,7 @@ def objective(trial: optuna.Trial) -> float: for i, var in enumerate(all_imputed_vars): # Build predictor set: original predictors + previously imputed variables current_predictors = _get_sequential_predictors( - predictors, all_imputed_vars, i + predictors, all_imputed_vars, i, sequential=self.sequential ) # Get properly encoded predictor columns @@ -1495,7 +1601,9 @@ def objective(trial: optuna.Trial) -> float: ) # Predict - y_pred = model.predict(X_val_augmented[encoded_predictors]) + y_pred = model.predict( + X_val_augmented[encoded_predictors], exact_quantile=0.5 + ) # Add predictions to augmented datasets for next variable X_train_augmented[var] = model.predict( @@ -1629,7 +1737,7 @@ def objective(trial: optuna.Trial) -> float: for i, var in enumerate(all_imputed_vars): # Build predictor set: original predictors + previously imputed variables current_predictors = _get_sequential_predictors( - predictors, all_imputed_vars, i + predictors, all_imputed_vars, i, sequential=self.sequential ) # Get properly encoded predictor columns diff --git a/microimpute/models/quantreg.py b/microimpute/models/quantreg.py index c767d387..aa3bb17a 100644 --- a/microimpute/models/quantreg.py +++ b/microimpute/models/quantreg.py @@ -9,7 +9,7 @@ from pydantic import validate_call from statsmodels.tools.sm_exceptions import IterationLimitWarning -from microimpute.config import VALIDATE_CONFIG +from microimpute.config import RANDOM_STATE, VALIDATE_CONFIG from microimpute.models.imputer import Imputer, ImputerResults warnings.filterwarnings("ignore", category=IterationLimitWarning) @@ -58,11 +58,27 @@ def __init__( log_level, ) self.models = models + self.default_quantiles = list(models[imputed_variables[0]]) + self.rng = np.random.default_rng(seed) self.quantiles_specified = quantiles_specified self.boolean_targets = boolean_targets or {} self.constant_targets = constant_targets or {} self.dummy_processor = dummy_processor + def _ensure_quantiles(self, quantiles: List[float]) -> None: + """Fit and cache missing quantiles using the original donor design matrix.""" + if any(not 0 < q < 1 for q in quantiles): + raise ValueError("QuantReg quantiles must be strictly between 0 and 1") + for variable in self.imputed_variables: + for q in quantiles: + if q in self.models[variable]: + continue + fitted = next(iter(self.models[variable].values())) + if variable in self.constant_targets: + self.models[variable][q] = fitted + else: + self.models[variable][q] = fitted.model.fit(q=q) + @validate_call(config=VALIDATE_CONFIG) def _predict( self, @@ -84,9 +100,16 @@ def _predict( Dictionary mapping quantiles to predicted values. Raises: - ValueError: If a requested quantile was not fitted during training. + ValueError: If a quantile is not strictly between zero and one. + + Quantiles absent from the original fit are fitted lazily on the retained + donor data and cached. Prediction data never participate in fitting. RuntimeError: If prediction fails. """ + if not np.isfinite(X_test[self.predictors].to_numpy(dtype=float)).all(): + raise ValueError("QuantReg prediction predictors must be finite") + if quantiles is not None: + self._ensure_quantiles(quantiles) # Log warning if return_probs is used with QuantReg if return_probs: self.logger.warning( @@ -99,7 +122,9 @@ def _predict( # Store original quantiles parameter to determine return type quantiles_param = quantiles - X_test_with_const = sm.add_constant(X_test[self.predictors]) + X_test_with_const = sm.add_constant( + X_test[self.predictors], has_constant="add" + ) self.logger.info(f"Prepared test data with {len(X_test)} samples") if quantiles is not None: @@ -147,7 +172,7 @@ def _predict( imputed_df[variable] = predictions imputations[q] = imputed_df else: - quantiles = list(self.models[self.imputed_variables[0]].keys()) + quantiles = self.default_quantiles if random_quantile_sample: self.logger.info("Sampling random quantiles for each prediction") mean_quantile = np.mean(quantiles) @@ -189,15 +214,14 @@ def _predict( # lookups and silently demotes numeric predictions to # object dtype — a major contributor to issue #96 # (OOM with many variables). - rng = np.random.default_rng(self.seed) index = random_q_imputations[quantiles[0]].index n_rows = len(index) quantiles_arr = np.asarray(quantiles) - # Sampled quantile index per row. - sampled_idx = rng.integers(0, len(quantiles_arr), size=n_rows) - result_df = pd.DataFrame(index=index) for variable in self.imputed_variables: + sampled_idx = self.rng.integers( + 0, len(quantiles_arr), size=n_rows + ) # Stack predictions for this variable across all # quantiles into an (n_rows, n_quantiles) array, # then select per-row with np.take_along_axis so @@ -276,9 +300,11 @@ class QuantReg(Imputer): directly predict specific quantiles. """ - def __init__(self, log_level: Optional[str] = "WARNING") -> None: + def __init__( + self, log_level: Optional[str] = "WARNING", seed: int = RANDOM_STATE + ) -> None: """Initialize the Quantile Regression model.""" - super().__init__(log_level=log_level) + super().__init__(seed=seed, log_level=log_level) self.models: Dict[str, Any] = {} self.log_level = log_level self.logger.debug("Initializing QuantReg imputer") @@ -341,7 +367,20 @@ def _fit( f"Values will be thresholded at 0.5 during prediction." ) + if X_train[predictors + imputed_variables].isna().any().any(): + raise ValueError( + "QuantReg training predictors and targets must not contain missing values" + ) + if not np.isfinite( + X_train[predictors + imputed_variables].to_numpy(dtype=float) + ).all(): + raise ValueError("QuantReg training predictors and targets must be finite") + if quantiles is not None and ( + not quantiles or any(not 0 < q < 1 for q in quantiles) + ): + raise ValueError("QuantReg quantiles must be strictly between 0 and 1") try: + self.models = {} for variable in imputed_variables: self.models[variable] = {} @@ -358,7 +397,7 @@ def _fit( f"Fitting QuantReg models for {len(quantiles)} quantiles: {quantiles}" ) - X_with_const = sm.add_constant(X_train[predictors]) + X_with_const = sm.add_constant(X_train[predictors], has_constant="add") self.logger.info( f"Prepared training data with {len(X_train)} samples, {len(predictors)} predictors" ) diff --git a/microimpute/models/zero_inflated.py b/microimpute/models/zero_inflated.py index fc9b5179..ac18c668 100644 --- a/microimpute/models/zero_inflated.py +++ b/microimpute/models/zero_inflated.py @@ -51,6 +51,7 @@ from microimpute.config import RANDOM_STATE, VALIDATE_CONFIG from microimpute.models.imputer import Imputer, ImputerResults from microimpute.models.qrf import QRF +from microimpute.utils.type_handling import declare_target_types # Regime labels. Kept as module-level constants so downstream code can # match on them without magic strings. @@ -160,6 +161,8 @@ def __init__( self.base_imputer_class = base_imputer_class or QRF self.base_imputer_kwargs = dict(base_imputer_kwargs or {}) self.zero_atol = float(zero_atol) + if not np.isfinite(self.zero_atol) or self.zero_atol < 0: + raise ValueError("zero_atol must be finite and nonnegative") self.classifier_type = classifier_type self.sequential = bool(sequential) @@ -187,6 +190,7 @@ def fit( weight_col: Optional[Union[str, np.ndarray, pd.Series]] = None, skip_missing: bool = False, not_numeric_categorical: Optional[List[str]] = None, + target_types: Optional[Dict[str, str]] = None, **kwargs: Any, ) -> Any: """Fit the regime-aware wrapper. @@ -199,7 +203,32 @@ def fit( Returns a ``ZeroInflatedImputerResults`` that routes predictions through each target's regime-specific pipeline. """ + X_train = declare_target_types(X_train, imputed_variables, target_types) + not_numeric_categorical = list(not_numeric_categorical or []) + not_numeric_categorical.extend( + name for name, kind in (target_types or {}).items() if kind == "numeric" + ) + if skip_missing: + imputed_variables = self._handle_missing_variables( + X_train, imputed_variables + ) self._validate_data(X_train, predictors + imputed_variables) + if set(predictors) & set(imputed_variables): + raise ValueError("Predictors and imputed variables must be distinct") + sample_weight = None + if isinstance(weight_col, str): + if weight_col not in X_train: + raise ValueError(f"Weight column {weight_col!r} not found") + sample_weight = X_train[weight_col].to_numpy(dtype=float) + elif isinstance(weight_col, pd.Series): + sample_weight = weight_col.reindex(X_train.index).to_numpy(dtype=float) + elif weight_col is not None: + sample_weight = np.asarray(weight_col, dtype=float) + if sample_weight is not None: + if sample_weight.shape != (len(X_train),): + raise ValueError("Weights must have one value per training row") + if not np.isfinite(sample_weight).all() or (sample_weight <= 0).any(): + raise ValueError("Weights must be positive and finite") # Classify target variables as numeric / categorical / boolean / # constant using the base Imputer's detector. @@ -259,6 +288,8 @@ def fit( regime=regime, y=y, not_numeric_categorical=nested_not_numeric_categorical, + sample_weight=sample_weight, + fit_kwargs=kwargs, ) bundle["predictors"] = list(seq_predictors) self._per_variable[var] = bundle @@ -307,117 +338,62 @@ def _fit_single_numeric( regime: str, y: np.ndarray, not_numeric_categorical: Optional[List[str]] = None, + sample_weight: Optional[np.ndarray] = None, + fit_kwargs: Optional[Dict[str, Any]] = None, ) -> Dict[str, Any]: - """Fit the gate and base imputer(s) for one numeric target. - - Returns a bundle dict with the regime, the gate classifier - (or None), and the base imputer(s) keyed by their role. - """ + """Fit weighted regime probabilities and weighted sign components.""" + if not np.isfinite(y).all(): + raise ValueError(f"Numeric target {variable!r} must contain finite values") X_pred = X_train[predictors].to_numpy(dtype=float, copy=False) - + if not np.isfinite(X_pred).all(): + raise ValueError("Zero-inflated predictors must contain finite values") if regime == REGIME_DEGENERATE_ZERO: return {"kind": "constant", "value": 0.0} - if regime in (REGIME_POSITIVE_ONLY, REGIME_NEGATIVE_ONLY): - # No gate; single base imputer on the full training set. - return { - "kind": "single", - "base": self._fit_base_single( - X_train, - predictors, - variable, - not_numeric_categorical=not_numeric_categorical, - ), - } - - if regime == REGIME_ZI_POSITIVE: - labels = (y > self.zero_atol).astype(int) - clf = _make_classifier(self.classifier_type, self.seed) - clf.fit(X_pred, labels) - pos_mask = y > self.zero_atol - pos_base = self._fit_base_single( - X_train.loc[pos_mask], - predictors, - variable, - not_numeric_categorical=not_numeric_categorical, + def fit_component(mask: np.ndarray, offset: int) -> ImputerResults: + weights = sample_weight[mask] if sample_weight is not None else None + seed = ( + None + if self.seed is None + else (self.seed + 3 * self.imputed_variables.index(variable) + offset) ) - return { - "kind": "zi_positive", - "classifier": clf, - "positive_base": pos_base, - } - - if regime == REGIME_ZI_NEGATIVE: - labels = (y < -self.zero_atol).astype(int) - clf = _make_classifier(self.classifier_type, self.seed) - clf.fit(X_pred, labels) - neg_mask = y < -self.zero_atol - neg_base = self._fit_base_single( - X_train.loc[neg_mask], + return self._fit_base_single( + X_train.loc[mask], predictors, variable, not_numeric_categorical=not_numeric_categorical, + sample_weight=weights, + seed=seed, + fit_kwargs=fit_kwargs, ) - return { - "kind": "zi_negative", - "classifier": clf, - "negative_base": neg_base, - } - if regime == REGIME_SIGN_ONLY: - # No zero class, but both signs present. Binary sign gate - # plus a base imputer per sign. - labels = (y > 0).astype(int) - clf = _make_classifier(self.classifier_type, self.seed) - clf.fit(X_pred, labels) - pos_mask = y > 0 - neg_mask = ~pos_mask + if regime in (REGIME_POSITIVE_ONLY, REGIME_NEGATIVE_ONLY): return { - "kind": "sign_only", - "classifier": clf, - "positive_base": self._fit_base_single( - X_train.loc[pos_mask], - predictors, - variable, - not_numeric_categorical=not_numeric_categorical, - ), - "negative_base": self._fit_base_single( - X_train.loc[neg_mask], - predictors, - variable, - not_numeric_categorical=not_numeric_categorical, - ), + "kind": "single", + "base": fit_component(np.ones(len(y), dtype=bool), 0), } + positive = y > self.zero_atol + negative = y < -self.zero_atol if regime == REGIME_THREE_SIGN: - # 0 / neg / pos three-way gate + two base imputers. - labels = np.where( - y > self.zero_atol, - 2, - np.where(y < -self.zero_atol, 0, 1), - ) - clf = _make_classifier(self.classifier_type, self.seed) - clf.fit(X_pred, labels) - pos_mask = y > self.zero_atol - neg_mask = y < -self.zero_atol - return { - "kind": "three_sign", - "classifier": clf, - "positive_base": self._fit_base_single( - X_train.loc[pos_mask], - predictors, - variable, - not_numeric_categorical=not_numeric_categorical, - ), - "negative_base": self._fit_base_single( - X_train.loc[neg_mask], - predictors, - variable, - not_numeric_categorical=not_numeric_categorical, - ), - } - - raise ValueError(f"Unhandled regime {regime!r}") + labels = np.where(positive, 2, np.where(negative, 0, 1)) + kind = "three_sign" + elif regime == REGIME_SIGN_ONLY: + labels, kind = positive.astype(int), "sign_only" + elif regime == REGIME_ZI_POSITIVE: + labels, kind = positive.astype(int), "zi_positive" + elif regime == REGIME_ZI_NEGATIVE: + labels, kind = negative.astype(int), "zi_negative" + else: + raise ValueError(f"Unhandled regime {regime!r}") + classifier = _make_classifier(self.classifier_type, self.seed) + classifier.fit(X_pred, labels, sample_weight=sample_weight) + bundle = {"kind": kind, "classifier": classifier} + if positive.any(): + bundle["positive_base"] = fit_component(positive, 1) + if negative.any(): + bundle["negative_base"] = fit_component(negative, 2) + return bundle def _fit_base_single( self, @@ -425,18 +401,27 @@ def _fit_base_single( predictors: List[str], variable: str, not_numeric_categorical: Optional[List[str]] = None, + sample_weight: Optional[np.ndarray] = None, + seed: Optional[int] = None, + fit_kwargs: Optional[Dict[str, Any]] = None, ) -> ImputerResults: - """Fit a single base Imputer on a (possibly filtered) slice.""" - imputer = self.base_imputer_class( - log_level="ERROR", + """Fit a component with its own seed and aligned conditional weights.""" + constructor_kwargs = { + "log_level": "ERROR", + "seed": seed, **self.base_imputer_kwargs, - ) - return imputer.fit( + } + imputer = self.base_imputer_class(**constructor_kwargs) + result = imputer.fit( X_train=X_train, predictors=predictors, imputed_variables=[variable], + weight_col=sample_weight, not_numeric_categorical=not_numeric_categorical, + **(fit_kwargs or {}), ) + # The fit API returns (result, params) when tuning is requested. + return result[0] if isinstance(result, tuple) else result class ZeroInflatedImputerResults(ImputerResults): @@ -477,21 +462,32 @@ def predict( ) -> Union[pd.DataFrame, Dict[float, pd.DataFrame]]: """Predict imputed values, routing per-variable by regime. - For numeric targets, the gate assigns each record to zero, - positive, or negative regime (depending on the detected - regime), and the base imputer for that regime produces the - nonzero draw. Zeros are set exactly to 0.0 (no stochastic - smearing). + Explicit quantiles invert the fitted signed mixture CDF, including + its zero atom, deterministically. Without quantiles, the gate samples + a regime and its component produces a stochastic draw. Quantiles + across multiple sequential targets are unsupported because chaining + conditional quantiles does not produce marginal quantiles. For non-numeric targets (categorical / boolean / constant), delegation is to the single auxiliary base imputer fit at training time. """ + self._validate_quantiles(quantiles) + if return_probs: + raise NotImplementedError( + "ZeroInflatedImputer does not expose categorical probabilities" + ) if quantiles is not None: - # Quantile grid not currently supported in the wrapper; the - # regime routing only produces a single stochastic draw per - # call. Deterministic-quantile support would require the - # caller to specify quantile conditional on regime. + if not quantiles: + raise ValueError("quantiles must not be empty") + if any( + set(bundle.get("predictors", [])) & set(self._regimes) + for bundle in self._per_variable.values() + ): + raise NotImplementedError( + "Marginal quantiles for sequential multi-target imputation " + "are not available; fit with sequential=False or a single target" + ) return { q: self._predict_single_draw(X_test, quantile=q, **kwargs) for q in quantiles @@ -563,6 +559,11 @@ def _predict_single_variable( dtype=float, copy=False ) + if quantile is not None: + return self._mixture_quantile( + X_test, X_pred, variable, bundle, quantile, **kwargs + ) + if kind == "zi_positive": clf = bundle["classifier"] draw = self._bernoulli_gate_draw(clf, X_pred) @@ -649,6 +650,79 @@ def _predict_single_variable( raise ValueError(f"Unhandled bundle kind {kind!r}") + def _mixture_quantile( + self, + X_test: pd.DataFrame, + X_pred: np.ndarray, + variable: str, + bundle: Dict[str, Any], + quantile: float, + **kwargs: Any, + ) -> np.ndarray: + """Invert the ordered negative / zero / positive mixture CDF. + + Negative quantiles use q / p_neg; the zero atom occupies + (p_neg, p_neg + p_zero]; positive quantiles use + (q - p_neg - p_zero) / p_pos. Endpoints select the first/last + nonempty component, including rows with degenerate gate probabilities. + """ + classifier = bundle["classifier"] + probabilities = classifier.predict_proba(X_pred) + by_class = { + label: probabilities[:, i] for i, label in enumerate(classifier.classes_) + } + zero = np.zeros(len(X_test)) + if bundle["kind"] == "three_sign": + p_neg, p_zero, p_pos = (by_class.get(k, zero) for k in (0, 1, 2)) + elif bundle["kind"] == "sign_only": + p_neg, p_zero, p_pos = by_class.get(0, zero), zero, by_class.get(1, zero) + elif bundle["kind"] == "zi_positive": + p_neg, p_zero, p_pos = zero, by_class.get(0, zero), by_class.get(1, zero) + else: + p_neg, p_zero, p_pos = by_class.get(1, zero), by_class.get(0, zero), zero + values = np.zeros(len(X_test)) + negative = (p_neg > 0) & (quantile <= p_neg) + if quantile == 1: + negative &= (p_zero == 0) & (p_pos == 0) + positive = (p_pos > 0) & ( + (quantile == 1) + | (quantile > p_neg + p_zero) + | ((quantile == 0) & (p_neg + p_zero == 0)) + ) + for sign, mask, offset, mass in ( + ("negative", negative, zero, p_neg), + ("positive", positive, p_neg + p_zero, p_pos), + ): + if not mask.any(): + continue + conditional_q = np.clip((quantile - offset[mask]) / mass[mask], 0, 1) + if quantile == 1: + conditional_q[:] = 1 + result = bundle[f"{sign}_base"] + subset = X_test.loc[mask] + if hasattr(result, "_predict_quantiles_per_row") and not kwargs: + component = result._predict_quantiles_per_row( + subset, variable, conditional_q + ) + else: + component = np.empty(mask.sum()) + for q in np.unique(conditional_q): + selected = conditional_q == q + predictions = self._invoke_base( + result, subset.loc[selected], quantile=float(q), **kwargs + ) + component[selected] = predictions[variable].to_numpy(dtype=float) + if not np.isfinite(component).all(): + raise ValueError("Component quantiles must be finite") + invalid = component >= 0 if sign == "negative" else component <= 0 + if invalid.any(): + raise ValueError( + f"The {sign} component predicts outside its sign support; " + "use a base imputer that preserves the component support" + ) + values[mask] = component + return values + def _invoke_base( self, base_result: ImputerResults, diff --git a/microimpute/utils/data.py b/microimpute/utils/data.py index db2752cd..b276ff78 100644 --- a/microimpute/utils/data.py +++ b/microimpute/utils/data.py @@ -112,7 +112,7 @@ def normalize_data( std = data[numeric_cols].std(axis=0) # Check for constant columns (std=0) - constant_cols = std[std == 0].index.tolist() + constant_cols = std[(std == 0) | std.isna()].index.tolist() if constant_cols: logger.warning(f"Found constant columns (std=0): {constant_cols}") # Handle constant columns by setting std to 1 to avoid division by zero @@ -466,6 +466,22 @@ def preprocess_data( if missing_count > 0: logger.warning(f"Data contains {missing_count} missing values") + if not full_data: + train, test = train_test_split( + data, train_size=train_size, test_size=test_size, random_state=random_state + ) + transformed = preprocess_data( + train, + full_data=True, + normalize=normalize, + log_transform=log_transform, + asinh_transform=asinh_transform, + ) + if normalize_requested or log_transform_requested or asinh_transform_requested: + train, params = transformed + return train, apply_transformations(test, params), params + return train, test + # Apply normalization if requested normalization_params = {} if normalize_requested: @@ -679,3 +695,35 @@ def un_asinh_transform_predictions( ) return untransformed + + +def apply_transformations(data: pd.DataFrame, params: dict) -> pd.DataFrame: + """Apply transformations fitted on training data to any new rows.""" + result = data.copy() + for column, values in params.get("normalization", {}).items(): + if column in result: + result[column] = (result[column] - values["mean"]) / values["std"] + for column in params.get("log_transform", {}): + if column in result: + if (result[column] <= 0).any(): + raise ValueError(f"Column '{column}' contains non-positive values") + result[column] = np.log(result[column]) + for column in params.get("asinh_transform", {}): + if column in result: + result[column] = np.arcsinh(result[column]) + return result + + +def reverse_transformations(data: pd.DataFrame, params: dict) -> pd.DataFrame: + """Reverse fitted transformations only for the supplied columns.""" + result = data.copy() + for column, values in params.get("normalization", {}).items(): + if column in result: + result[column] = result[column] * values["std"] + values["mean"] + for column in params.get("log_transform", {}): + if column in result: + result[column] = np.exp(result[column]) + for column in params.get("asinh_transform", {}): + if column in result: + result[column] = np.sinh(result[column]) + return result diff --git a/microimpute/utils/statmatch_hotdeck.py b/microimpute/utils/statmatch_hotdeck.py index 9cbf14ec..db447a7d 100644 --- a/microimpute/utils/statmatch_hotdeck.py +++ b/microimpute/utils/statmatch_hotdeck.py @@ -5,6 +5,7 @@ """ import logging +from contextlib import contextmanager from typing import Any, List, Tuple import numpy as np @@ -48,6 +49,26 @@ def _get_statmatch(): return _statmatch_cache["StatMatch"] +@contextmanager +def _temporary_r_seed(seed): + """Seed a bridge call without changing the caller's R random stream.""" + if seed is None: + yield + return + had_state = ".Random.seed" in ro.globalenv + original_state = ( + ro.IntVector(list(ro.globalenv[".Random.seed"])) if had_state else None + ) + try: + ro.r["set.seed"](int(seed)) + yield + finally: + if had_state: + ro.globalenv[".Random.seed"] = original_state + elif ".Random.seed" in ro.globalenv: + del ro.globalenv[".Random.seed"] + + @validate_call(config=VALIDATE_CONFIG) def nnd_hotdeck_using_rpy2( receiver: pd.DataFrame, @@ -98,41 +119,54 @@ def nnd_hotdeck_using_rpy2( log.error(msg) raise ValueError(msg) + # NND.hotdeck has no weight argument. RANDwNND.hotdeck accepts + # weight.don as the NAME of a column in data.don, not a vector. + # cut.don="min" preserves nearest-distance matching and weights ties. + # https://search.r-project.org/CRAN/refmans/StatMatch/html/RANDwNND.hotdeck.html + r_kwargs = dict(matching_kwargs) + random_state = r_kwargs.pop("random_state", None) + donor_sample_weight = r_kwargs.pop("donor_sample_weight", None) + donor_for_matching = donor + matching_function = StatMatch.NND_hotdeck + if donor_sample_weight is not None: + weights = np.asarray(donor_sample_weight, dtype=float) + if weights.ndim != 1 or len(weights) != len(donor): + raise ValueError("Donor weights must contain one value per donor") + if not np.isfinite(weights).all() or (weights <= 0).any(): + raise ValueError("Donor weights must be positive and finite") + if r_kwargs.pop("constrained", False): + raise ValueError( + "Weighted constrained matching is not supported by RANDwNND.hotdeck" + ) + r_kwargs.pop("constr_alg", None) + if "k" in r_kwargs and "cut_don" not in r_kwargs: + raise ValueError( + "Weighted matching with k requires an explicit cut_don rule" + ) + weight_column = "__microimpute_donor_weight__" + while weight_column in donor.columns: + weight_column += "_" + donor_for_matching = donor.copy() + donor_for_matching[weight_column] = weights + r_kwargs["weight_don"] = weight_column + r_kwargs.setdefault("cut_don", "min") + matching_function = StatMatch.RANDwNND_hotdeck + with localconverter( default_converter + pandas2ri.converter + numpy2ri.converter ): r_receiver = conversion.py2rpy(receiver) - r_donor = conversion.py2rpy(donor) + r_donor = conversion.py2rpy(donor_for_matching) r_match = ro.StrVector(matching_variables) r_z = ro.StrVector(z_variables) - # Extract optional donor sample weights (threaded from Imputer.fit - # when weight_col was supplied). StatMatch accepts these via the - # ``weight.don`` R argument; we pop it from matching_kwargs so that - # other kwargs pass through unchanged. - r_kwargs = dict(matching_kwargs) - donor_sample_weight = r_kwargs.pop("donor_sample_weight", None) - if donor_sample_weight is not None: - with localconverter( - default_converter + pandas2ri.converter + numpy2ri.converter - ): - r_kwargs["weight_don"] = ro.FloatVector( - np.asarray(donor_sample_weight, dtype=float) - ) - - if r_kwargs: - out_NND = StatMatch.NND_hotdeck( + with _temporary_r_seed(random_state): + out_NND = matching_function( data_rec=r_receiver, data_don=r_donor, match_vars=r_match, **r_kwargs, ) - else: - out_NND = StatMatch.NND_hotdeck( - data_rec=r_receiver, - data_don=r_donor, - match_vars=r_match, - ) # Create the correct matching indices matrix for StatMatch.create_fused recipient_indices = np.arange(1, len(receiver) + 1) @@ -220,7 +254,7 @@ def nnd_hotdeck_using_rpy2( mtc_matrix = np.column_stack((recipient_indices, donor_indices_valid)) # Convert to R matrix mtc_ids = ro.r.matrix( - ro.IntVector(mtc_matrix.flatten()), + ro.IntVector(mtc_matrix.flatten(order="F")), nrow=len(recipient_indices), ncol=2, ) diff --git a/microimpute/utils/type_handling.py b/microimpute/utils/type_handling.py index fc4d2207..f3a1f556 100644 --- a/microimpute/utils/type_handling.py +++ b/microimpute/utils/type_handling.py @@ -17,27 +17,8 @@ class VariableTypeDetector: @staticmethod def is_boolean_variable(series: pd.Series) -> bool: - """Check if a series represents boolean data. - - A float series that happens to contain only {0.0, 1.0} is NOT - treated as boolean (#9); it could be a probability, a rescaled - indicator, or simply an accident of a small sample. Routing a - float column through a classifier would silently flip the model - type and destroy regression behaviour. Only genuine boolean - dtypes and integer columns with values in {0, 1} are classified - as boolean. - """ - if pd.api.types.is_bool_dtype(series): - return True - - unique_vals = set(series.dropna().unique()) - if pd.api.types.is_integer_dtype(series) and unique_vals <= {0, 1}: - return True - - # Deliberately NOT recognising floats with values {0.0, 1.0} as - # booleans — see docstring. Callers who really want a float - # 0/1 column treated as a boolean should cast it explicitly. - return False + """Recognize explicit boolean dtypes without reclassifying integer counts.""" + return pd.api.types.is_bool_dtype(series) @staticmethod def is_categorical_variable(series: pd.Series) -> bool: @@ -87,30 +68,22 @@ def categorize_variable( variable_type: 'bool', 'categorical', 'numeric_categorical', or 'numeric' categories: List of unique values for categorical types, None for numeric """ + if force_numeric: + if not pd.api.types.is_numeric_dtype(series): + raise ValueError( + f"Variable '{col_name}' declared numeric has nonnumeric dtype" + ) + return "numeric", None if VariableTypeDetector.is_boolean_variable(series): return "bool", None + if isinstance(series.dtype, pd.CategoricalDtype): + return "categorical", series.cat.categories.tolist() if VariableTypeDetector.is_categorical_variable(series): return "categorical", series.unique().tolist() - # Check if it would normally be numeric_categorical - if not force_numeric and VariableTypeDetector.is_numeric_categorical_variable( - series - ): - categories = [float(i) for i in series.unique().tolist()] - logger.info( - f"Treating numeric variable '{col_name}' as categorical due to low unique count and equal spacing" - ) - return "numeric_categorical", categories - - # If force_numeric is True or it's not numeric_categorical, treat as numeric - if force_numeric and VariableTypeDetector.is_numeric_categorical_variable( - series - ): - logger.info( - f"Variable '{col_name}' forced to be treated as numeric (override numeric_categorical detection)" - ) - + # Cardinality is not a semantic type: counts and continuous values remain + # numeric even when a training fold contains only a few distinct values. return "numeric", None @@ -151,6 +124,9 @@ def preprocess_predictors( Returns: Tuple of (processed_data, updated_predictors) """ + self.predictor_numeric = { + col: pd.api.types.is_numeric_dtype(data[col]) for col in predictors + } # Start with a copy containing all needed columns all_columns = list(set(predictors + imputed_variables)) data = data[all_columns].copy() @@ -373,6 +349,13 @@ def apply_dummy_encoding_to_test( data = data.copy() updated_predictors = predictors.copy() + if not data.columns.is_unique: + raise ValueError("Duplicate DataFrame column names are not supported") + for column, numeric in getattr(self, "predictor_numeric", {}).items(): + if column not in data: + raise ValueError(f"Missing predictor column: {column}") + if pd.api.types.is_numeric_dtype(data[column]) != numeric: + raise ValueError(f"Incompatible predictor dtype for '{column}'") # Apply dummy encoding based on stored mapping for orig_col, dummy_cols in self.dummy_mapping.items(): if orig_col in predictors and orig_col in data.columns: @@ -433,3 +416,42 @@ def apply_dummy_encoding_to_test( data[col] = data[col].astype("float64") return data, updated_predictors + + +def declare_target_types( + data: pd.DataFrame, + targets: List[str], + target_types: Optional[Dict[str, str]] = None, +) -> pd.DataFrame: + """Validate explicit target semantics before preprocessing or splitting.""" + target_types = target_types or {} + if set(target_types) - set(targets): + raise ValueError("target_types contains unknown imputed variables") + if any( + value not in {"numeric", "categorical", "bool"} + for value in target_types.values() + ): + raise ValueError("target_types values must be numeric, categorical, or bool") + result = data.copy() + for variable, declared_type in target_types.items(): + if declared_type == "categorical": + result[variable] = result[variable].astype("category") + elif declared_type == "bool": + if ( + result[variable].isna().any() + or not result[variable].isin([0, 1, False, True]).all() + ): + raise ValueError( + f"Boolean target '{variable}' must contain only 0 and 1" + ) + result[variable] = result[variable].astype(bool) + elif not pd.api.types.is_numeric_dtype(result[variable]): + if isinstance( + result[variable].dtype, pd.CategoricalDtype + ) and pd.api.types.is_numeric_dtype(result[variable].cat.categories): + result[variable] = pd.to_numeric(result[variable]) + else: + raise ValueError( + f"Variable '{variable}' declared numeric has nonnumeric dtype" + ) + return result diff --git a/tests/test_autoimpute.py b/tests/test_autoimpute.py index 24ad1b0e..f3f1dc7a 100644 --- a/tests/test_autoimpute.py +++ b/tests/test_autoimpute.py @@ -25,6 +25,10 @@ except ImportError: HAS_MDN = False +# These tests exercise the core API on every installation. Optional backends +# have dedicated test modules and must not gate core coverage. +CORE_MODELS = [QRF, QuantReg, OLS] + # === Fixtures === @@ -87,7 +91,7 @@ def test_autoimpute_basic_structure( receiver_data=diabetes_receiver, predictors=predictors, imputed_variables=imputed_variables, - models=[QRF, Matching, QuantReg, OLS] if not HAS_MDN else None, + models=CORE_MODELS, hyperparameters={ "QRF": {"n_estimators": 50}, "Matching": {"constrained": True}, @@ -137,7 +141,7 @@ def test_autoimpute_all_models( receiver_data=diabetes_receiver, predictors=predictors, imputed_variables=imputed_variables, - models=[QRF, Matching, QuantReg, OLS] if not HAS_MDN else None, + models=CORE_MODELS, impute_all=True, # Return results for all models log_level="WARNING", ) @@ -202,7 +206,7 @@ def test_autoimpute_with_hyperparameters(simple_data: tuple) -> None: receiver_data=receiver, predictors=["x1", "x2"], imputed_variables=["y1"], - models=[QRF, Matching, QuantReg, OLS] if not HAS_MDN else None, + models=CORE_MODELS, hyperparameters=hyperparameters, log_level="WARNING", ) @@ -224,7 +228,7 @@ def test_autoimpute_multiple_imputed_variables(simple_data: tuple) -> None: receiver_data=receiver, predictors=["x1", "x2"], imputed_variables=["y1", "y2"], # Multiple variables - models=[QRF, Matching, QuantReg, OLS] if not HAS_MDN else None, + models=CORE_MODELS, log_level="WARNING", ) @@ -246,7 +250,7 @@ def test_autoimpute_large_receiver() -> None: receiver_data=receiver, predictors=["x"], imputed_variables=["y"], - models=[QRF, Matching, QuantReg, OLS] if not HAS_MDN else None, + models=CORE_MODELS, log_level="WARNING", ) @@ -267,7 +271,7 @@ def test_autoimpute_best_method_selection(simple_data: tuple) -> None: receiver_data=receiver, predictors=["x1", "x2"], imputed_variables=["y1"], - models=[QRF, Matching, QuantReg, OLS] if not HAS_MDN else None, + models=CORE_MODELS, log_level="WARNING", ) @@ -326,7 +330,7 @@ def test_autoimpute_cv_results_structure(simple_data: tuple) -> None: receiver_data=receiver, predictors=["x1", "x2"], imputed_variables=["y1"], - models=[QRF, Matching, QuantReg, OLS] if not HAS_MDN else None, + models=CORE_MODELS, log_level="WARNING", ) @@ -378,7 +382,7 @@ def test_autoimpute_missing_predictors() -> None: receiver_data=receiver, predictors=["x1", "x2"], # x2 not in receiver imputed_variables=["y"], - models=[QRF, Matching, QuantReg, OLS] if not HAS_MDN else None, + models=CORE_MODELS, log_level="WARNING", ) @@ -461,7 +465,7 @@ def test_autoimpute_consistency(simple_data: tuple) -> None: receiver_data=receiver, predictors=["x1", "x2"], imputed_variables=["y1"], - models=[QRF, Matching, QuantReg, OLS] if not HAS_MDN else None, + models=CORE_MODELS, log_level="WARNING", ) @@ -470,7 +474,7 @@ def test_autoimpute_consistency(simple_data: tuple) -> None: receiver_data=receiver, predictors=["x1", "x2"], imputed_variables=["y1"], - models=[QRF, Matching, QuantReg, OLS] if not HAS_MDN else None, + models=CORE_MODELS, log_level="WARNING", ) diff --git a/tests/test_dashboard_formatter.py b/tests/test_dashboard_formatter.py index 1cf7f8e5..7984e999 100644 --- a/tests/test_dashboard_formatter.py +++ b/tests/test_dashboard_formatter.py @@ -974,14 +974,14 @@ def test_numeric_categorical_distribution_uses_categorical_rows(self): """Test numeric categorical variables produce categorical distributions.""" donor_data = pd.DataFrame( { - "rating": [1, 1, 2, 2, 3, 3], - "flag": [0, 1, 1, 0, 1, 0], + "rating": pd.Categorical([1, 1, 2, 2, 3, 3]), + "flag": np.array([0, 1, 1, 0, 1, 0], dtype=bool), } ) receiver_data = pd.DataFrame( { - "rating": [1, 2, 2, 3], - "flag": [1, 1, 0, 0], + "rating": pd.Categorical([1, 2, 2, 3]), + "flag": np.array([1, 1, 0, 0], dtype=bool), } ) diff --git a/tests/test_data_preprocessing.py b/tests/test_data_preprocessing.py index 5bb339cc..6282c737 100644 --- a/tests/test_data_preprocessing.py +++ b/tests/test_data_preprocessing.py @@ -22,8 +22,8 @@ def test_normalize_excludes_categorical_columns(self): data = pd.DataFrame( { "numeric_col": [1.0, 2.5, 3.7, 4.2, 5.9], # Non-equally spaced - "categorical_col": [1, 2, 3, 1, 2], - "boolean_col": [0, 1, 0, 1, 0], + "categorical_col": pd.Categorical([1, 2, 3, 1, 2]), + "boolean_col": np.array([0, 1, 0, 1, 0], dtype=bool), } ) @@ -52,8 +52,8 @@ def test_normalize_preserves_column_names(self): data = pd.DataFrame( { "age": [25, 30, 35, 40, 45], - "race": [1, 2, 3, 1, 2], - "is_female": [0, 1, 0, 1, 0], + "race": pd.Categorical([1, 2, 3, 1, 2]), + "is_female": np.array([0, 1, 0, 1, 0], dtype=bool), "income": [50000, 60000, 70000, 80000, 90000], } ) @@ -79,7 +79,7 @@ def test_normalize_correctly_normalizes_numeric_columns(self): 410.0, 505.0, ], # Non-equally spaced - "category": [1, 2, 1, 2, 1], + "category": pd.Categorical([1, 2, 1, 2, 1]), } ) @@ -108,12 +108,9 @@ def test_normalize_handles_constant_columns(self): normalized_data, norm_params = normalize_data(data) - # Constant columns are detected as numeric_categorical and excluded - # So they should remain unchanged - pd.testing.assert_series_equal(normalized_data["constant"], data["constant"]) - - # Only varying column should have normalization params - assert "constant" not in norm_params + # A constant numeric column stays numeric and uses unit scale. + np.testing.assert_array_equal(normalized_data["constant"], np.zeros(5)) + assert norm_params["constant"] == {"mean": 5.0, "std": 1.0} assert "varying" in norm_params def test_normalize_returns_copy(self): @@ -121,7 +118,7 @@ def test_normalize_returns_copy(self): data = pd.DataFrame( { "value": [1.3, 2.7, 3.2, 4.8, 5.1], # Non-equally spaced - "category": [1, 2, 1, 2, 1], + "category": pd.Categorical([1, 2, 1, 2, 1]), } ) original_data = data.copy() @@ -136,7 +133,12 @@ def test_normalize_returns_copy(self): def test_normalize_with_no_numeric_columns(self): """Test normalize with only categorical columns.""" - data = pd.DataFrame({"cat1": [1, 2, 3, 1, 2], "cat2": [0, 1, 0, 1, 0]}) + data = pd.DataFrame( + { + "cat1": pd.Categorical([1, 2, 3, 1, 2]), + "cat2": pd.Categorical([0, 1, 0, 1, 0]), + } + ) normalized_data, norm_params = normalize_data(data) @@ -161,8 +163,8 @@ def test_preprocess_data_excludes_categoricals_from_normalization(self): 40.9, 45.1, ], # Non-equally spaced floats - "race": [1, 2, 3, 1, 2], - "is_female": [0, 1, 0, 1, 0], + "race": pd.Categorical([1, 2, 3, 1, 2]), + "is_female": np.array([0, 1, 0, 1, 0], dtype=bool), "income": [ 50123.45, 60987.23, @@ -203,7 +205,7 @@ def test_categorical_columns_dont_get_weird_suffixes_when_dummified( """ data = pd.DataFrame( { - "race": [1, 2, 3, 1, 2, 3, 1, 2], + "race": pd.Categorical([1, 2, 3, 1, 2, 3, 1, 2]), "income": [ 50000, 60000, @@ -247,8 +249,8 @@ def test_log_transform_excludes_categorical_columns(self): data = pd.DataFrame( { "numeric_col": [1.0, 2.5, 3.7, 4.2, 5.9], - "categorical_col": [1, 2, 3, 1, 2], - "boolean_col": [0, 1, 0, 1, 0], + "categorical_col": pd.Categorical([1, 2, 3, 1, 2]), + "boolean_col": np.array([0, 1, 0, 1, 0], dtype=bool), } ) @@ -298,7 +300,7 @@ def test_log_transform_correctly_transforms_numeric_columns(self): 96.1, 102.4, ], - "category": [1, 2, 1, 2, 1, 2, 1, 2, 1, 2], + "category": pd.Categorical([1, 2, 1, 2, 1, 2, 1, 2, 1, 2]), } ) @@ -340,7 +342,7 @@ def test_log_transform_returns_copy(self): data = pd.DataFrame( { "value": [1.5, 2.7, 3.2, 4.8, 5.1, 6.3, 7.9, 8.4, 9.6, 10.2], - "category": [1, 2, 1, 2, 1, 2, 1, 2, 1, 2], + "category": pd.Categorical([1, 2, 1, 2, 1, 2, 1, 2, 1, 2]), } ) original_data = data.copy() @@ -355,7 +357,12 @@ def test_log_transform_returns_copy(self): def test_log_transform_with_no_numeric_columns(self): """Test log transform with only categorical columns.""" - data = pd.DataFrame({"cat1": [1, 2, 3, 1, 2], "cat2": [0, 1, 0, 1, 0]}) + data = pd.DataFrame( + { + "cat1": pd.Categorical([1, 2, 3, 1, 2]), + "cat2": pd.Categorical([0, 1, 0, 1, 0]), + } + ) log_data, log_params = log_transform_data(data) @@ -451,8 +458,8 @@ def test_preprocess_data_excludes_categoricals_from_log_transform(self): 65.7, 70.2, ], - "race": [1, 2, 3, 1, 2, 3, 1, 2, 3, 1], - "is_female": [0, 1, 0, 1, 0, 1, 0, 1, 0, 1], + "race": pd.Categorical([1, 2, 3, 1, 2, 3, 1, 2, 3, 1]), + "is_female": np.array([0, 1, 0, 1, 0, 1, 0, 1, 0, 1], dtype=bool), "income": [ 50123.45, 60987.23, @@ -894,8 +901,8 @@ def test_asinh_transform_excludes_categorical_columns(self): data = pd.DataFrame( { "numeric_col": [1.0, 2.5, 3.7, 4.2, 5.9], - "categorical_col": [1, 2, 3, 1, 2], - "boolean_col": [0, 1, 0, 1, 0], + "categorical_col": pd.Categorical([1, 2, 3, 1, 2]), + "boolean_col": np.array([0, 1, 0, 1, 0], dtype=bool), } ) @@ -945,7 +952,7 @@ def test_asinh_transform_correctly_transforms_numeric_columns(self): 100000.1, 1000000.7, ], - "category": [1, 2, 1, 2, 1, 2, 1, 2, 1, 2], + "category": pd.Categorical([1, 2, 1, 2, 1, 2, 1, 2, 1, 2]), } ) @@ -988,7 +995,7 @@ def test_asinh_transform_returns_copy(self): data = pd.DataFrame( { "value": [-10.5, -2.7, 0.0, 2.8, 10.1, 100.3, 1000.9, 10000.4], - "category": [1, 2, 1, 2, 1, 2, 1, 2], + "category": pd.Categorical([1, 2, 1, 2, 1, 2, 1, 2]), } ) original_data = data.copy() @@ -1079,8 +1086,8 @@ def test_preprocess_data_excludes_categoricals_from_asinh_transform(self): 65.7, 70.2, ], - "race": [1, 2, 3, 1, 2, 3, 1, 2, 3, 1], - "is_female": [0, 1, 0, 1, 0, 1, 0, 1, 0, 1], + "race": pd.Categorical([1, 2, 3, 1, 2, 3, 1, 2, 3, 1]), + "is_female": np.array([0, 1, 0, 1, 0, 1, 0, 1, 0, 1], dtype=bool), "income": [ -10000.0, 0.0, diff --git a/tests/test_metrics.py b/tests/test_metrics.py index 58ab02c3..e711c197 100644 --- a/tests/test_metrics.py +++ b/tests/test_metrics.py @@ -49,8 +49,10 @@ def mixed_type_data() -> pd.DataFrame: "num_target1": np.random.randn(n_samples) * 3, "num_target2": np.random.randn(n_samples) + 5, # Categorical targets - "binary_target": np.random.choice([0, 1], size=n_samples), - "multiclass_target": np.random.choice([0, 1, 2], size=n_samples), + "binary_target": np.random.choice([False, True], size=n_samples), + "multiclass_target": pd.Categorical( + np.random.choice([0, 1, 2], size=n_samples) + ), "string_target": np.random.choice(["A", "B", "C"], size=n_samples), } ) @@ -82,7 +84,7 @@ def test_metric_detection_numerical() -> None: def test_metric_detection_categorical() -> None: """Test that categorical variables are correctly identified.""" # Binary data - binary_series = pd.Series([0, 1, 0, 1, 1, 0, 1, 0]) + binary_series = pd.Series([0, 1, 0, 1, 1, 0, 1, 0], dtype=bool) assert get_metric_for_variable_type(binary_series, "binary_var") == "log_loss" # String categorical @@ -90,7 +92,7 @@ def test_metric_detection_categorical() -> None: assert get_metric_for_variable_type(string_series, "string_var") == "log_loss" # Low cardinality integer (categorical-like) - low_card_series = pd.Series([0, 1, 2, 0, 1, 2, 0, 1, 2]) + low_card_series = pd.Series([0, 1, 2, 0, 1, 2, 0, 1, 2], dtype="category") assert get_metric_for_variable_type(low_card_series, "low_card_var") == "log_loss" # Boolean type @@ -116,20 +118,24 @@ def test_log_loss_with_class_labels() -> None: y_true = np.array([0, 1, 0, 1, 1]) y_pred_labels = np.array([0, 1, 1, 1, 0]) # Class predictions - # Should convert to probabilities with a warning - loss = log_loss(y_true, y_pred_labels) - assert loss > 0 - # Loss should be higher since we're using high-confidence probabilities - assert loss > 1 + with pytest.raises(ValueError, match="probabilities"): + log_loss(y_true, y_pred_labels) def test_log_loss_multiclass() -> None: """Test log loss with multiclass data.""" y_true = np.array([0, 1, 2, 0, 1, 2]) - # Provide class predictions (should be converted) - y_pred_classes = np.array([0, 1, 2, 1, 1, 2]) - - loss = log_loss(y_true, y_pred_classes) + probabilities = np.array( + [ + [0.8, 0.1, 0.1], + [0.1, 0.8, 0.1], + [0.1, 0.1, 0.8], + [0.2, 0.7, 0.1], + [0.1, 0.8, 0.1], + [0.1, 0.1, 0.8], + ] + ) + loss = log_loss(y_true, probabilities) assert loss > 0 @@ -150,7 +156,7 @@ def test_compute_loss_quantile() -> None: def test_compute_loss_log() -> None: """Test compute_loss with log loss metric.""" y_true = np.random.choice([0, 1], size=50) - y_pred = np.random.choice([0, 1], size=50) + y_pred = np.random.default_rng(3).uniform(0.1, 0.9, size=50) losses, mean_loss = compute_loss( y_true, y_pred, "log_loss", q=0.5, labels=np.array([0, 1]) @@ -655,8 +661,6 @@ def test_autoimpute_with_all_models(mixed_type_data: pd.DataFrame) -> None: del receiver_data[target] models = [OLS, QRF, QuantReg] - if HAS_MATCHING: - models.append(Matching) result = autoimpute( donor_data=donor_data, @@ -957,15 +961,10 @@ def test_probability_ordering_with_real_model() -> None: assert not np.isinf(loss_ordered), "Log loss should not be infinite" assert loss_ordered > 0, "Log loss should be positive" - # Check if this is better than using dummy probabilities - # With dummy probabilities (converting class predictions to 0.99/0.01) + # Label predictions must not masquerade as calibrated probabilities. class_preds = predictions[0.5]["target"].values - _, loss_dummy = compute_loss(y_test, class_preds, "log_loss") - - # Real probabilities should give better (lower) loss than dummy probabilities - assert loss_ordered < loss_dummy, ( - "Real probabilities should give better loss than dummy probabilities" - ) + with pytest.raises(ValueError, match="probabilities"): + compute_loss(y_test, class_preds, "log_loss") # === Distribution Comparison Tests === diff --git a/tests/test_models/test_imputers.py b/tests/test_models/test_imputers.py index bd7016f6..a4fb6302 100644 --- a/tests/test_models/test_imputers.py +++ b/tests/test_models/test_imputers.py @@ -67,13 +67,8 @@ def data_with_edge_cases() -> pd.DataFrame: ALL_IMPUTER_MODELS = [OLS, QuantReg, QRF] CATEGORICAL_MODELS = [OLS, QRF] -try: - from microimpute.models.matching import Matching - - ALL_IMPUTER_MODELS.append(Matching) - CATEGORICAL_MODELS.append(Matching) -except ImportError: - pass +# These shared tests require conditional quantiles/probabilities. Matching's +# donor-draw contract is exercised in test_matching[_correctness].py. try: from microimpute.models.mdn import MDN @@ -265,7 +260,7 @@ def test_imputation_categorical_targets( assert pd.api.types.is_string_dtype(predictions["categorical"]) # Test probability predictions for models that support it - if model_class.__name__ in ["OLS", "QRF", "Matching"]: + if model_class.__name__ in ["OLS", "QRF"]: # Get predictions with probabilities using quantiles # (this ensures consistent return format across models) predictions_with_probs = fitted_model.predict( @@ -395,7 +390,8 @@ def test_multiple_targets( X_train, X_test = preprocess_data(data) - model = model_class() + # Marginal quantiles need independent targets; sequential QRF supports joint draws. + model = model_class(sequential=False) if model_class is QRF else model_class() if model_class.__name__ == "QuantReg": fitted = model.fit(X_train, predictors, imputed_variables, quantiles=[0.5]) @@ -728,12 +724,6 @@ def test_missing_predictors_in_test(model_class: Type[Imputer]) -> None: _REPRODUCIBILITY_MODELS = [OLS, QuantReg, QRF] -try: - from microimpute.models.matching import Matching as _Matching_for_repro - - _REPRODUCIBILITY_MODELS.append(_Matching_for_repro) -except ImportError: - pass @pytest.mark.parametrize( diff --git a/tests/test_models/test_matching.py b/tests/test_models/test_matching.py index 25b5403f..17e630ce 100644 --- a/tests/test_models/test_matching.py +++ b/tests/test_models/test_matching.py @@ -13,13 +13,17 @@ from microimpute.utils.data import preprocess_data from microimpute.visualizations import * -try: - from microimpute.models.matching import Matching +# The Matching class can load without R for injected Python callbacks; +# these integration tests specifically require the optional R bridge. +pytest.importorskip("rpy2.robjects") +from microimpute.models.matching import Matching +from microimpute.utils.statmatch_hotdeck import _get_statmatch +from rpy2.robjects.packages import PackageNotInstalledError - MATCHING_AVAILABLE = True -except ImportError: - MATCHING_AVAILABLE = False - pytest.skip("Matching model not available", allow_module_level=True) +try: + _get_statmatch() +except PackageNotInstalledError: + pytest.skip("R StatMatch package not available", allow_module_level=True) # === Fixtures === @@ -76,36 +80,20 @@ def test_matching_basic_fit_predict(diabetes_data: pd.DataFrame) -> None: model = Matching() fitted_model = model.fit(X_train, predictors, imputed_variables) - # Predict (matching uses same value for all quantiles) - predictions = fitted_model.predict(X_test, quantiles=[0.5]) + # Predict a donor draw for each recipient. + predictions = fitted_model.predict(X_test) # Validate predictions - assert isinstance(predictions, dict) - assert 0.5 in predictions - assert isinstance(predictions[0.5], pd.DataFrame) - assert predictions[0.5].shape == (len(X_test), len(imputed_variables)) - assert not predictions[0.5].isna().any().any() + assert isinstance(predictions, pd.DataFrame) + assert predictions.shape == (len(X_test), len(imputed_variables)) + assert not predictions.isna().any().any() -def test_matching_quantile_invariance(simple_data: pd.DataFrame) -> None: - """Test that Matching returns same values for different quantiles.""" - X_train, X_test = preprocess_data(simple_data) - - model = Matching() - fitted_model = model.fit(X_train, ["x1", "x2"], ["y"]) - - # Get predictions at different quantiles - predictions = fitted_model.predict(X_test, quantiles=[0.1, 0.5, 0.9]) - - # Matching should return same values for all quantiles - # (it doesn't model uncertainty) - for i in range(len(X_test)): - val_01 = predictions[0.1]["y"].iloc[i] - val_05 = predictions[0.5]["y"].iloc[i] - val_09 = predictions[0.9]["y"].iloc[i] - assert val_01 == val_05 == val_09, ( - "Matching should return same value for all quantiles" - ) +def test_matching_rejects_conditional_quantiles(simple_data: pd.DataFrame) -> None: + """A donor draw must not be mislabeled as several conditional quantiles.""" + fitted = Matching().fit(simple_data, ["x1", "x2"], ["y"]) + with pytest.raises(NotImplementedError, match="conditional quantiles"): + fitted.predict(simple_data, quantiles=[0.1, 0.5, 0.9]) def test_matching_donor_preservation(simple_data: pd.DataFrame) -> None: @@ -115,10 +103,10 @@ def test_matching_donor_preservation(simple_data: pd.DataFrame) -> None: model = Matching() fitted_model = model.fit(X_train, ["x1", "x2"], ["y"]) - predictions = fitted_model.predict(X_test[:1], quantiles=[0.5]) + predictions = fitted_model.predict(X_test[:1]) # The predicted value should be from the training set - predicted_value = predictions[0.5]["y"].iloc[0] + predicted_value = predictions["y"].iloc[0] assert predicted_value in X_train["y"].values, ( "Matched value should be from donor pool" ) @@ -146,14 +134,13 @@ def test_matching_different_distance_functions() -> None: model = Matching() fitted_model = model.fit(X_train, ["x1", "x2"], ["y"], dist_fun=dist_fun) - predictions = fitted_model.predict(X_test[:5], quantiles=[0.5]) + predictions = fitted_model.predict(X_test[:5]) - assert 0.5 in predictions - assert not predictions[0.5]["y"].isna().any() + assert not predictions["y"].isna().any() -def test_matching_k_neighbors() -> None: - """Test Matching with different k values.""" +def test_matching_donor_reuse_limit() -> None: + """NND k constrains donor reuse; it does not count nearest neighbors.""" np.random.seed(42) data = pd.DataFrame( { @@ -168,12 +155,13 @@ def test_matching_k_neighbors() -> None: # Test different k values for k in [1, 3, 5]: model = Matching() - fitted_model = model.fit(X_train, ["x1", "x2"], ["y"], k=k) + fitted_model = model.fit( + X_train, ["x1", "x2"], ["y"], k=k, constrained=True, constr_alg="lpSolve" + ) - predictions = fitted_model.predict(X_test[:5], quantiles=[0.5]) + predictions = fitted_model.predict(X_test[:5]) - assert 0.5 in predictions - assert not predictions[0.5]["y"].isna().any() + assert not predictions["y"].isna().any() # === Categorical Variables === @@ -203,10 +191,10 @@ def test_matching_mixed_types() -> None: ["target_numeric", "target_category"], ) - predictions = fitted_model.predict(X_test, quantiles=[0.5]) + predictions = fitted_model.predict(X_test) - assert predictions[0.5]["target_numeric"].dtype == np.float64 - assert pd.api.types.is_string_dtype(predictions[0.5]["target_category"]) + assert predictions["target_numeric"].dtype == np.float64 + assert pd.api.types.is_string_dtype(predictions["target_category"]) # === Edge Cases === @@ -221,13 +209,12 @@ def test_matching_single_donor(simple_data: pd.DataFrame) -> None: model = Matching() fitted_model = model.fit(X_train, ["x1", "x2"], ["y"]) - predictions = fitted_model.predict(X_test, quantiles=[0.5]) + predictions = fitted_model.predict(X_test) - assert 0.5 in predictions - assert not predictions[0.5]["y"].isna().any() + assert not predictions["y"].isna().any() # All predictions should be from the small donor pool - for val in predictions[0.5]["y"]: + for val in predictions["y"]: assert val in X_train["y"].values @@ -250,11 +237,10 @@ def test_matching_exact_match() -> None: model = Matching() fitted_model = model.fit(X_train, ["x1", "x2"], ["y"]) - predictions = fitted_model.predict(X_test, quantiles=[0.5]) + predictions = fitted_model.predict(X_test) # Check that predictions exist - assert 0.5 in predictions - assert not predictions[0.5].empty + assert predictions["y"].iloc[0] == 30 # === Constrained Matching === @@ -277,10 +263,9 @@ def test_matching_constrained_mode() -> None: model = Matching() fitted_model = model.fit(X_train, ["x1", "x2"], ["y"], constrained=True) - predictions = fitted_model.predict(X_test, quantiles=[0.5]) + predictions = fitted_model.predict(X_test) - assert 0.5 in predictions - assert not predictions[0.5]["y"].isna().any() + assert not predictions["y"].isna().any() # === Cross-Validation === @@ -303,14 +288,10 @@ def test_matching_cross_validation(diabetes_data: pd.DataFrame) -> None: assert "quantile_loss" in matching_results assert "log_loss" in matching_results - # Check quantile_loss results (for numerical variables) - ql_results = matching_results["quantile_loss"] - assert "results" in ql_results - assert isinstance(ql_results["results"], pd.DataFrame) - assert "train" in ql_results["results"].index - assert "test" in ql_results["results"].index - assert not ql_results["results"].isna().all().all() - assert ql_results["mean_test"] > 0 + # Matching does not estimate conditional distributions and cannot be + # ranked by the quantile/log-loss comparison API. + for metric in ["quantile_loss", "log_loss"]: + assert np.isnan(matching_results[metric]["mean_test"]) # === Hyperparameter Tuning === @@ -344,16 +325,16 @@ def test_matching_hyperparameter_tuning(diabetes_data: pd.DataFrame) -> None: ) # Make predictions - default_preds = default_fitted.predict(X_valid, quantiles=[0.5]) - tuned_preds = tuned_fitted.predict(X_valid, quantiles=[0.5]) + default_preds = default_fitted.predict(X_valid) + tuned_preds = tuned_fitted.predict(X_valid) # Calculate MSE default_mse = {} tuned_mse = {} for var in imputed_variables: - default_mse[var] = mean_squared_error(X_valid[var], default_preds[0.5][var]) - tuned_mse[var] = mean_squared_error(X_valid[var], tuned_preds[0.5][var]) + default_mse[var] = mean_squared_error(X_valid[var], default_preds[var]) + tuned_mse[var] = mean_squared_error(X_valid[var], tuned_preds[var]) # Both should produce valid results assert all(mse < np.inf for mse in default_mse.values()) @@ -391,12 +372,12 @@ def test_matching_multiple_targets(diabetes_data: pd.DataFrame) -> None: model = Matching() fitted_model = model.fit(X_train, predictors, imputed_variables) - predictions = fitted_model.predict(X_test, quantiles=[0.5]) + predictions = fitted_model.predict(X_test) - assert predictions[0.5].shape[1] == len(imputed_variables) + assert predictions.shape[1] == len(imputed_variables) for var in imputed_variables: - assert var in predictions[0.5].columns - assert not predictions[0.5][var].isna().any() + assert var in predictions.columns + assert not predictions[var].isna().any() def test_matching_preserves_relationships() -> None: @@ -420,15 +401,40 @@ def test_matching_preserves_relationships() -> None: model = Matching() fitted_model = model.fit(X_train, ["x"], ["y1", "y2"]) - predictions = fitted_model.predict(X_test, quantiles=[0.5]) + predictions = fitted_model.predict(X_test) # Check that the relationship between y1 and y2 is preserved # Since we're matching entire rows, y1 and y2 should maintain their relationship - pred_y1 = predictions[0.5]["y1"].values - pred_y2 = predictions[0.5]["y2"].values + pred_y1 = predictions["y1"].values + pred_y2 = predictions["y2"].values # Each prediction should come from the same donor row for i in range(len(pred_y1)): # Find which donor row was matched donor_mask = (X_train["y1"] == pred_y1[i]) & (X_train["y2"] == pred_y2[i]) assert donor_mask.any(), "Predictions should come from same donor row" + + +def test_matching_weights_change_tied_donor_selection(): + """Optional live-R check of RANDwNND's documented weighted tie selection.""" + donor = pd.DataFrame({"x": [1.0, 1.0], "y": [10.0, 20.0], "w": [1000.0, 1.0]}) + fitted = Matching().fit(donor, ["x"], ["y"], weight_col="w") + output = fitted.predict(pd.DataFrame({"x": np.ones(500)})) + assert (output.y == 10.0).mean() > 0.97 + + +def test_matching_seeded_draws_preserve_r_random_stream(): + """Optional live-R reproducibility and global RNG-isolation integration check.""" + import rpy2.robjects as ro + + donor = pd.DataFrame({"x": [1.0, 1.0], "y": [10.0, 20.0], "w": [1.0, 2.0]}) + receiver = pd.DataFrame({"x": np.ones(200)}) + first = Matching(seed=17).fit(donor, ["x"], ["y"], weight_col="w") + second = Matching(seed=17).fit(donor, ["x"], ["y"], weight_col="w") + ro.r["set.seed"](31) + expected_next_draws = np.asarray(ro.r["runif"](3)) + ro.r["set.seed"](31) + first_draw = first.predict(receiver) + np.testing.assert_array_equal(np.asarray(ro.r["runif"](3)), expected_next_draws) + pd.testing.assert_frame_equal(first_draw, second.predict(receiver)) + assert not first_draw.equals(first.predict(receiver)) diff --git a/tests/test_models/test_matching_correctness.py b/tests/test_models/test_matching_correctness.py new file mode 100644 index 00000000..9ec20a06 --- /dev/null +++ b/tests/test_models/test_matching_correctness.py @@ -0,0 +1,129 @@ +"""Matching API regression tests using an injected matcher, without optional R.""" + +import numpy as np +import pandas as pd +import pytest + +from microimpute.models.matching import Matching + + +@pytest.fixture +def donor(): + return pd.DataFrame( + {"x": np.arange(12.0), "y": np.arange(12.0) + 0.5, "w": np.arange(12.0) + 1} + ) + + +def first_donor(receiver, donor, matching_variables, z_variables, **kwargs): + result = receiver.copy() + for var in z_variables: + result[var] = donor[var].iloc[0] + return result, result + + +@pytest.mark.parametrize( + "options", [{"quantiles": [0.1, 0.5, 0.9]}, {"return_probs": True}] +) +def test_matching_rejects_unestimated_distributions(donor, options): + fitted = Matching(matching_hotdeck=first_donor).fit(donor, ["x"], ["y"]) + with pytest.raises( + NotImplementedError, match="conditional quantiles|probabilities" + ): + fitted.predict(donor[["x"]], **options) + + +def test_tuned_matching_preserves_donor_weights_and_fixed_options(donor, monkeypatch): + received = [] + + def match(**kwargs): + received.append(kwargs) + return first_donor(**kwargs) + + def tune(self, **kwargs): + assert np.array_equal(kwargs["matching_kwargs"]["donor_sample_weight"], donor.w) + return {"dist_fun": "Euclidean"} + + monkeypatch.setattr(Matching, "_tune_hyperparameters", tune) + fitted, params = Matching(matching_hotdeck=match).fit( + donor, ["x"], ["y"], weight_col="w", tune_hyperparameters=True, keep_t=True + ) + actual = fitted.predict(donor[["x"]]) + assert len(actual) == len(donor) + assert params == {"dist_fun": "Euclidean"} + assert received[0]["keep_t"] is True + np.testing.assert_array_equal(received[0]["donor_sample_weight"], donor.w) + + +def test_matching_cv_slices_weights_and_never_scores_failed_trials(donor, monkeypatch): + import optuna + + trials = [] + calls = [] + real_create_study = optuna.create_study + + def study_factory(*args, **kwargs): + study = real_create_study(*args, **kwargs) + trials.append(study) + return study + + def match(**kwargs): + calls.append(kwargs) + np.testing.assert_array_equal(kwargs["donor_sample_weight"], kwargs["donor"].w) + raise RuntimeError("injected matching failure") + + monkeypatch.setattr(optuna, "create_study", study_factory) + model = Matching(matching_hotdeck=match) + with pytest.raises(ValueError, match="No matching hyperparameter trial succeeded"): + model.fit(donor, ["x"], ["y"], weight_col="w", tune_hyperparameters=True) + assert calls + assert all(t.state == optuna.trial.TrialState.PRUNED for t in trials[0].trials) + + +def test_chunk_failures_report_counts_and_reset_on_success(donor): + def match(**kwargs): + if kwargs["receiver"].x.iloc[0] == 2: + raise RuntimeError("injected second-chunk failure") + return first_donor(**kwargs) + + fitted = Matching(matching_hotdeck=match).fit(donor, ["x"], ["y"]) + receiver = donor[["x"]].iloc[:5] + output = fitted._predict_chunked(receiver, quantiles=None, chunk_size=2) + assert fitted.n_failed_records == 2 + assert output.attrs["n_failed_records"] == 2 + assert output.y.isna().sum() == 2 + clean = fitted.predict(receiver.iloc[:1]) + assert fitted.n_failed_records == 0 + assert clean.attrs["n_failed_records"] == 0 + + +def test_r_matching_seed_stream_reproduces_and_advances(donor, monkeypatch): + """Default bridge gets fresh child seeds while equal model seeds reproduce.""" + import sys + from types import SimpleNamespace + + seeds = [] + + def fake_bridge(**kwargs): + seeds.append(kwargs.pop("random_state")) + return first_donor(**kwargs) + + monkeypatch.setitem( + sys.modules, + "microimpute.utils.statmatch_hotdeck", + SimpleNamespace(nnd_hotdeck_using_rpy2=fake_bridge), + ) + receiver = donor[["x"]].iloc[:5] + first = Matching(seed=17).fit(donor, ["x"], ["y"], weight_col="w") + first.predict(receiver) + first._predict_chunked(receiver, quantiles=None, chunk_size=2) + initial = seeds.copy() + assert len(initial) == 4 + assert len(set(initial)) == 4 + seeds.clear() + second = Matching(seed=17).fit(donor, ["x"], ["y"], weight_col="w") + second.predict(receiver) + second._predict_chunked(receiver, quantiles=None, chunk_size=2) + assert seeds == initial + third = Matching(seed=18).fit(donor, ["x"], ["y"], weight_col="w") + third.predict(receiver) + assert seeds[-1] != initial[0] diff --git a/tests/test_models/test_qrf.py b/tests/test_models/test_qrf.py index 3e2f0b8a..9016bb52 100644 --- a/tests/test_models/test_qrf.py +++ b/tests/test_models/test_qrf.py @@ -72,7 +72,7 @@ def test_qrf_basic_fit_predict(diabetes_data: pd.DataFrame) -> None: X_train, X_test = preprocess_data(data) # Initialize and fit model - model = QRF() + model = QRF(sequential=False) fitted_model = model.fit( X_train, predictors, @@ -111,15 +111,15 @@ def test_qrf_sequential_imputation(diabetes_data: pd.DataFrame) -> None: # Get predictions small_test = X_test.head(5).copy() - sequential_preds = fitted_model.predict(small_test, quantiles=[0.5])[0.5] + sequential_preds = fitted_model.predict(small_test) # Compare with parallel imputation (each variable independently) parallel_predictions = {} for var in imputed_variables: single_model = QRF() single_fitted = single_model.fit(X_train, predictors, [var], n_estimators=30) - single_pred = single_fitted.predict(small_test, quantiles=[0.5]) - parallel_predictions[var] = single_pred[0.5][var] + single_pred = single_fitted.predict(small_test) + parallel_predictions[var] = single_pred[var] # Sequential should differ from parallel for later variables differences_found = False @@ -137,7 +137,7 @@ def test_qrf_sequential_imputation(diabetes_data: pd.DataFrame) -> None: reversed_fitted = reversed_model.fit( X_train, predictors, imputed_variables[::-1], n_estimators=30 ) - reversed_preds = reversed_fitted.predict(small_test, quantiles=[0.5])[0.5] + reversed_preds = reversed_fitted.predict(small_test) # Middle variable should differ when imputed in different orders assert not np.allclose( @@ -234,15 +234,17 @@ def test_qrf_target_filters_fit_each_target_on_eligible_rows() -> None: """Target-specific masks should let each target ignore its own bad rows.""" train = pd.DataFrame( { - "x": [0.0, 1.0, 2.0, 3.0], - "y1": [10.0, 20.0, np.nan, np.nan], - "y2": [np.nan, np.nan, 30.0, 40.0], - "y1_observed": [True, True, False, False], - "y2_observed": [False, False, True, True], + "x": [0.0, 1.0, 2.0, 3.0, 4.0], + # y1 is observed in every row used to train the conditional y2 + # model, including rows excluded from y1's own training filter. + "y1": [10.0, 20.0, 30.0, 40.0, np.nan], + "y2": [np.nan, np.nan, 30.0, 40.0, np.nan], + "y1_observed": [True, True, False, False, False], + "y2_observed": [False, False, True, True, False], } ) - fitted = QRF().fit( + fitted = QRF(sequential=False).fit( train, predictors=["x"], imputed_variables=["y1", "y2"], @@ -760,10 +762,8 @@ def test_qrf_error_handling() -> None: # Try to predict with missing predictor test_data = pd.DataFrame({"z": [7, 8, 9]}) - try: - predictions = fitted_model.predict(test_data) - except Exception as e: - assert "none of" in str(e).lower() and "are in the" in str(e).lower() + with pytest.raises(ValueError, match="Missing predictor column: x"): + fitted_model.predict(test_data) # === Internal Model Tests === @@ -1256,10 +1256,10 @@ def test_qrf_sequential_imputation_discrete_numeric_categorical() -> None: def test_qrf_not_numeric_categorical_override() -> None: - """Test that not_numeric_categorical parameter correctly overrides automatic detection. + """Explicit target types distinguish category codes from numeric counts. - Variables with <10 unique equally-spaced values normally get treated as categorical, - but this parameter should force them to be treated as numeric. + Counts remain numeric by default; categorical labels must be declared. + The legacy force-numeric override remains accepted. """ np.random.seed(42) n_samples = 200 @@ -1286,7 +1286,7 @@ def test_qrf_not_numeric_categorical_override() -> None: {"predictor1": np.random.randn(50), "predictor2": np.random.randn(50)} ) - # Test 1: Default behavior - discrete vars should be treated as categorical + # Test 1: Explicitly declared category codes use classification model_default = QRF(log_level="WARNING") fitted_default = model_default.fit( X_train=donor_df, @@ -1294,14 +1294,15 @@ def test_qrf_not_numeric_categorical_override() -> None: imputed_variables=["discrete_var1", "discrete_var2", "continuous_var"], n_estimators=20, random_state=42, + target_types={"discrete_var1": "categorical", "discrete_var2": "categorical"}, ) # Check that discrete vars were treated as categorical assert "discrete_var1" in model_default.categorical_targets, ( - "discrete_var1 should be categorical by default" + "discrete_var1 should be categorical when explicitly declared" ) assert "discrete_var2" in model_default.categorical_targets, ( - "discrete_var2 should be categorical by default" + "discrete_var2 should be categorical when explicitly declared" ) assert "continuous_var" in model_default.numeric_targets, ( "continuous_var should be numeric" @@ -1313,7 +1314,8 @@ def test_qrf_not_numeric_categorical_override() -> None: X_train=donor_df, predictors=["predictor1", "predictor2"], imputed_variables=["discrete_var1", "discrete_var2", "continuous_var"], - not_numeric_categorical=["discrete_var1"], # Force discrete_var1 to be numeric + not_numeric_categorical=["discrete_var1"], # Legacy numeric override + target_types={"discrete_var2": "categorical"}, n_estimators=20, random_state=42, ) @@ -1697,3 +1699,43 @@ def test_seed_is_configurable_and_reproducible() -> None: second = QRF(log_level="WARNING", seed=1234).fit(train, ["x"], ["y"]).predict(test) np.testing.assert_allclose(first["y"], second["y"]) + + +def test_qrf_disjoint_target_filters_reject_missing_chained_predictors(): + """No joint observations cannot identify a conditional second-target model.""" + train = pd.DataFrame( + { + "x": [0.0, 1.0, 2.0, 3.0], + "y1": [10.0, 20.0, np.nan, np.nan], + "y2": [np.nan, np.nan, 30.0, 40.0], + } + ) + with pytest.raises(RuntimeError, match="NaN"): + QRF().fit( + train, + ["x"], + ["y1", "y2"], + target_filters={"y1": train.y1.notna(), "y2": train.y2.notna()}, + n_estimators=3, + ) + + +def test_qrf_independent_disjoint_target_filters_fit_marginals(): + """Separate donor subsets identify marginals without inventing joint data.""" + train = pd.DataFrame( + { + "x": [0.0, 1.0, 2.0, 3.0], + "y1": [10.0, 20.0, np.nan, np.nan], + "y2": [np.nan, np.nan, 30.0, 40.0], + } + ) + fitted = QRF(sequential=False).fit( + train, + ["x"], + ["y1", "y2"], + target_filters={"y1": train.y1.notna(), "y2": train.y2.notna()}, + n_estimators=3, + ) + prediction = fitted.predict(pd.DataFrame({"x": [1.5]}), quantiles=[0.5])[0.5] + assert 10 <= prediction.y1.iloc[0] <= 20 + assert 30 <= prediction.y2.iloc[0] <= 40 diff --git a/tests/test_models/test_qrf_distribution.py b/tests/test_models/test_qrf_distribution.py new file mode 100644 index 00000000..5ee80d25 --- /dev/null +++ b/tests/test_models/test_qrf_distribution.py @@ -0,0 +1,234 @@ +"""Publication regression checks for QRF distribution defaults and weights.""" + +import logging + +import numpy as np +import pandas as pd + +from microimpute.config import DEFAULT_MODEL_PARAMS +from microimpute.models.qrf import _QRFModel + + +def test_qrf_defaults_keep_multiple_observations_per_leaf(): + X = pd.DataFrame({"x": np.arange(100, dtype=float)}) + model = _QRFModel(4, logging.getLogger(__name__)) + model.fit(X, pd.Series(np.sin(X.x), name="y")) + assert ( + model.qrf.min_samples_leaf + == DEFAULT_MODEL_PARAMS["qrf"]["min_samples_leaf"] + >= 10 + ) + assert model.qrf.max_samples_leaf is None + overridden = _QRFModel(4, logging.getLogger(__name__)) + overridden.fit( + X, pd.Series(np.sin(X.x), name="y"), min_samples_leaf=3, n_estimators=7 + ) + assert overridden.qrf.min_samples_leaf == 3 + assert overridden.qrf.n_estimators == 7 + + +def test_qrf_weighted_conditional_cdf_without_predictor_signal(): + # No splits are possible, so the weighted median is analytically known. + X = pd.DataFrame({"x": np.zeros(100)}) + y = pd.Series(np.arange(100, dtype=float), name="y") + weights = np.r_[np.repeat(100.0, 10), np.ones(90)] + models = [] + for scale in [1, 50]: + model = _QRFModel(4, logging.getLogger(__name__)) + model.fit(X, y, sample_weight=weights * scale, bootstrap=False, n_estimators=5) + models.append(model) + query = pd.DataFrame({"x": [0.0]}) + for model in models: + assert model.predict(query, exact_quantile=0.5).iloc[0] == 5 + assert model.predict(query, exact_quantile=0.9).iloc[0] == 9 + + +def test_qrf_default_quantile_coverage_on_normal_noise(): + # Independent holdout from Y|X ~ N(2X,1): default tails must approximate the + # known .1/.9 probabilities. This checks this fixture, not universal calibration. + rng = np.random.default_rng(22) + X = pd.DataFrame({"x": rng.uniform(-2, 2, 3000)}) + y = pd.Series(2 * X.x + rng.normal(size=len(X)), name="y") + query = pd.DataFrame({"x": rng.uniform(-2, 2, 2000)}) + truth = 2 * query.x + rng.normal(size=len(query)) + model = _QRFModel(42, logging.getLogger(__name__)) + model.fit(X, y) + for q in [0.1, 0.5, 0.9]: + observed = np.mean(truth <= model.predict(query, exact_quantile=q)) + assert abs(observed - q) < 0.055, (q, observed) + + +def test_qrf_weighted_cdf_preserves_bootstrap_multiplicity(): + X = pd.DataFrame({"x": np.zeros(50)}) + y = pd.Series(np.arange(50, dtype=float), name="y") + weights = np.linspace(1, 10, 50) + model = _QRFModel(11, logging.getLogger(__name__)) + model.fit(X, y, sample_weight=weights, n_estimators=4, min_samples_leaf=5) + # One leaf per tree makes the bootstrap-weighted CDF directly calculable. + tree_mass = [] + for samples in model.qrf.estimators_samples_: + mass = np.bincount(samples, minlength=len(y)) * weights + tree_mass.append(mass / mass.sum()) + cdf = np.cumsum(np.mean(tree_mass, axis=0)) + quantiles = np.array([0.1, 0.5, 0.9]) + expected = np.searchsorted(cdf, quantiles) + observed = model.predict_quantiles_per_row( + pd.DataFrame({"x": np.zeros(3)}), quantiles + ) + np.testing.assert_array_equal(observed, expected) + + +def test_qrf_row_specific_quantiles_match_scalar_queries(): + rng = np.random.default_rng(19) + X = pd.DataFrame({"x": rng.normal(size=100)}) + y = pd.Series(rng.normal(size=100), name="y") + model = _QRFModel(4, logging.getLogger(__name__)) + model.fit(X, y, n_estimators=5) + query = X.iloc[:5] + quantiles = np.array([0, 0.21, 0.53, 0.97, 1]) + output = model.predict_quantiles_per_row(query, quantiles) + expected = [ + model.predict(query.iloc[[i]], exact_quantile=float(q)).iloc[0] + for i, q in enumerate(quantiles) + ] + np.testing.assert_allclose(output, expected) + + +def test_sequential_multi_target_quantiles_are_explicitly_unsupported(): + from microimpute.models import QRF + import pytest + + rng = np.random.default_rng(36) + data = pd.DataFrame( + {"x": rng.normal(size=80), "a": rng.normal(size=80), "b": rng.normal(size=80)} + ) + fitted = QRF().fit(data, ["x"], ["a", "b"], n_estimators=5) + with pytest.raises(NotImplementedError, match="sequential=False"): + fitted.predict(data[["x"]].iloc[:3], quantiles=[0.5]) + assert fitted.predict(data[["x"]].iloc[:3]).shape == (3, 2) + + +def test_independent_quantiles_are_invariant_to_target_order_and_batching(): + from microimpute.models import QRF + + rng = np.random.default_rng(6) + data = pd.DataFrame( + { + "x": rng.normal(size=200), + "a": rng.normal(size=200), + "b": rng.normal(size=200), + } + ) + query = data[["x"]].iloc[:15] + reference = None + for batch_size in [None, 1]: + for variables in [["a", "b"], ["b", "a"]]: + fitted = QRF(sequential=False, seed=9, batch_size=batch_size).fit( + data, ["x"], variables, n_estimators=7 + ) + for model in fitted.models.values(): + assert model.feature_columns == ["x"] + prediction = fitted.predict(query, quantiles=[0.1, 0.5, 0.9]) + if reference is None: + reference = prediction + for q in prediction: + pd.testing.assert_frame_equal( + reference[q].sort_index(axis=1), prediction[q].sort_index(axis=1) + ) + + +def test_independent_quantile_is_marginal_for_nonlinearly_related_targets(): + from microimpute.models import QRF + + # For A~Uniform(-2,2), median(A)=0 but median(A**2)=1. Plugging + # median(A) into a conditional second-stage model incorrectly gives zero. + rng = np.random.default_rng(95) + a = rng.uniform(-2, 2, 1200) + data = pd.DataFrame({"x": np.zeros(len(a)), "a": a, "b": a**2}) + fitted = QRF(sequential=False).fit(data, ["x"], ["a", "b"], n_estimators=15) + medians = fitted.predict(pd.DataFrame({"x": [0.0]}), quantiles=[0.5])[0.5] + assert abs(medians.a.iloc[0]) < 0.15 + assert abs(medians.b.iloc[0] - 1.0) < 0.2 + + +def test_sequential_stochastic_draws_retain_joint_dependence(): + from microimpute.models import QRF + + rng = np.random.default_rng(11) + a = rng.normal(size=1500) + data = pd.DataFrame( + { + "x": rng.normal(size=len(a)), + "a": a, + "b": -a + rng.normal(scale=0.1, size=len(a)), + } + ) + fitted = QRF().fit(data, ["x"], ["a", "b"], n_estimators=20) + draws = fitted.predict(data[["x"]].iloc[:600]) + assert draws.corr().loc["a", "b"] < -0.9 + + +def test_independent_mode_is_preserved_after_tuning(monkeypatch): + from microimpute.models import QRF + + rng = np.random.default_rng(6) + data = pd.DataFrame( + {"x": rng.normal(size=80), "a": rng.normal(size=80), "b": rng.normal(size=80)} + ) + model = QRF(sequential=False) + monkeypatch.setattr( + model, "_tune_hyperparameters", lambda **kwargs: {"n_estimators": 3} + ) + fitted, _ = model.fit(data, ["x"], ["a", "b"], tune_hyperparameters=True) + assert not fitted.sequential + assert fitted.predict(data[["x"]].iloc[:2], quantiles=[0.5])[0.5].shape == (2, 2) + + +def test_independent_inner_tuning_uses_original_predictors(monkeypatch): + from microimpute.models import QRF + from microimpute.models.qrf import _RandomForestClassifierModel + + rng = np.random.default_rng(19) + data = pd.DataFrame( + { + "x": rng.normal(size=80), + "a": rng.normal(size=80), + "b": rng.normal(size=80), + "c": np.tile(["A", "B"], 40), + } + ) + seen = [] + for model_class in [_QRFModel, _RandomForestClassifierModel]: + original_fit = model_class.fit + + def recording_fit(self, X, y, *args, _fit=original_fit, **kwargs): + seen.append((y.name, list(X.columns))) + return _fit(self, X, y, *args, **kwargs) + + monkeypatch.setattr(model_class, "fit", recording_fit) + model = QRF(sequential=False) + model.imputed_variables = ["a", "b", "c"] + model.categorical_targets = {"c": {"type": "categorical", "categories": ["A", "B"]}} + model._tune_qrf_hyperparameters(data, ["x"], ["a", "b"], n_cv_folds=2, n_trials=1) + model._tune_rfc_hyperparameters(data, ["x"], ["c"], n_cv_folds=2, n_trials=1) + assert {name for name, _ in seen} == {"a", "b", "c"} + assert all(columns == ["x"] for _, columns in seen) + + +def test_independent_batch_fit_supports_mixed_targets(): + from microimpute.models import QRF + + rng = np.random.default_rng(7) + data = pd.DataFrame( + { + "x": rng.normal(size=60), + "a": rng.normal(size=60), + "label": np.tile(["A", "B"], 30), + } + ) + fitted = QRF(sequential=False, batch_size=1).fit( + data, ["x"], ["a", "label"], n_estimators=3 + ) + result = fitted.predict(data[["x"]].iloc[:3], quantiles=[0.5], return_probs=True) + assert result[0.5].shape == (3, 2) + assert result["probabilities"]["label"]["probabilities"].shape == (3, 2) diff --git a/tests/test_models/test_regression_correctness.py b/tests/test_models/test_regression_correctness.py new file mode 100644 index 00000000..e4148d7e --- /dev/null +++ b/tests/test_models/test_regression_correctness.py @@ -0,0 +1,147 @@ +"""Regression API and stochastic correctness checks for submission issue #206/#208.""" + +import numpy as np +import pandas as pd +import pytest +import statsmodels.api as sm + +from microimpute.models.ols import OLS +from microimpute.models.quantreg import QuantReg + + +@pytest.fixture +def regression_data(): + rng = np.random.default_rng(42) + x = rng.normal(size=200) + return pd.DataFrame({"x": x, "y": 3 + 2 * x + rng.normal(size=200)}) + + +@pytest.mark.parametrize("model_class", [OLS, QuantReg]) +@pytest.mark.parametrize("x", [[2.0], [2.0, 2.0, 2.0]]) +def test_homogeneous_receivers_keep_intercept(model_class, x, regression_data): + fitted = model_class().fit(regression_data, ["x"], ["y"]) + receiver = pd.DataFrame({"x": x}, index=np.arange(len(x)) + 1000) + actual = fitted.predict(receiver, quantiles=[0.5])[0.5] + estimator = sm.OLS if model_class is OLS else sm.QuantReg + reference = estimator( + regression_data.y, sm.add_constant(regression_data[["x"]], has_constant="add") + ).fit() + expected = reference.predict(sm.add_constant(receiver, has_constant="add")) + np.testing.assert_allclose(actual.y, expected) + assert actual.index.equals(receiver.index) + + +@pytest.mark.parametrize("model_class", [OLS, QuantReg]) +def test_constant_training_predictor_keeps_stable_design(model_class): + data = pd.DataFrame({"x": np.ones(30), "y": np.arange(30.0)}) + fitted = model_class().fit(data, ["x"], ["y"]) + actual = fitted.predict(pd.DataFrame({"x": [1.0]}), quantiles=[0.5])[0.5] + np.testing.assert_allclose(actual.y, [14.5], atol=1e-5) + + +@pytest.mark.parametrize("model_class", [OLS, QuantReg]) +@pytest.mark.parametrize("column", ["x", "y"]) +@pytest.mark.parametrize("bad_value", [np.nan, np.inf, -np.inf]) +def test_nonfinite_training_values_fail( + model_class, column, bad_value, regression_data +): + regression_data.loc[0, column] = bad_value + with pytest.raises((ValueError, RuntimeError), match="finite|missing|NaN|inf"): + model_class().fit(regression_data, ["x"], ["y"]) + + +@pytest.mark.parametrize("model_class", [OLS, QuantReg]) +@pytest.mark.parametrize("bad_value", [np.nan, np.inf, -np.inf]) +def test_nonfinite_prediction_values_fail(model_class, bad_value, regression_data): + fitted = model_class().fit(regression_data, ["x"], ["y"]) + with pytest.raises((ValueError, RuntimeError), match="finite|missing|NaN|inf"): + fitted.predict(pd.DataFrame({"x": [bad_value, 1.0]})) + + +def test_quantreg_fits_new_prediction_quantiles_on_donor_data(regression_data): + fitted = QuantReg().fit(regression_data, ["x"], ["y"]) + receiver = pd.DataFrame({"x": [-1.0, 0.0, 1.0]}) + initial_default = fitted.predict(receiver) + grid = [0.13, 0.5, 0.87] + actual = fitted.predict(receiver, quantiles=grid) + for q in grid: + reference = sm.QuantReg( + regression_data.y, + sm.add_constant(regression_data[["x"]], has_constant="add"), + ).fit(q=q) + expected = reference.predict(sm.add_constant(receiver, has_constant="add")) + np.testing.assert_allclose(actual[q].y, expected) + pd.testing.assert_frame_equal(fitted.predict(receiver), initial_default) + + +@pytest.mark.parametrize("q", [0.0, 1.0]) +def test_quantreg_rejects_unsupported_endpoints(q, regression_data): + fitted = QuantReg().fit(regression_data, ["x"], ["y"]) + with pytest.raises(ValueError, match="strictly between"): + fitted.predict(pd.DataFrame({"x": [0.0]}), quantiles=[q]) + + +def test_ols_sampling_has_one_independent_shock_per_row_and_advances(regression_data): + # Identical target fits expose accidental re-use of the same residual shock. + regression_data["z"] = regression_data.y + fitted = OLS().fit(regression_data, ["x"], ["y", "z"]) + receiver = pd.DataFrame({"x": np.zeros(2000)}) + first = fitted.predict(receiver, random_quantile_sample=True) + second = fitted.predict(receiver, random_quantile_sample=True) + assert first.y.nunique() == len(receiver) + assert not np.array_equal(first.y, second.y) + assert abs(first.y.corr(first.z)) < 0.1 + fresh = OLS().fit(regression_data, ["x"], ["y", "z"]) + pd.testing.assert_frame_equal( + first, fresh.predict(receiver, random_quantile_sample=True) + ) + + +def test_quantreg_sampling_advances_independently_across_targets(regression_data): + regression_data["z"] = regression_data.y + fitted = QuantReg().fit( + regression_data, ["x"], ["y", "z"], quantiles=[0.1, 0.5, 0.9] + ) + receiver = pd.DataFrame({"x": np.zeros(1000)}) + first = fitted.predict(receiver, random_quantile_sample=True)[0.5] + second = fitted.predict(receiver, random_quantile_sample=True)[0.5] + assert not np.array_equal(first.y, second.y) + assert abs(first.y.corr(first.z)) < 0.1 + + +def test_ols_survey_weight_units_do_not_change_predictive_quantiles(regression_data): + # Survey weights express relative population mass, so changing their units + # must leave both the conditional mean and residual distribution invariant. + weights = np.linspace(0.5, 4.0, len(regression_data)) + receiver = pd.DataFrame({"x": [-1.0, 0.0, 4.0]}) + grid = [0.1, 0.5, 0.9] + first = OLS().fit(regression_data, ["x"], ["y"], weight_col=weights) + scaled = OLS().fit(regression_data, ["x"], ["y"], weight_col=100 * weights) + first_predictions = first.predict(receiver, quantiles=grid) + scaled_predictions = scaled.predict(receiver, quantiles=grid) + for q in grid: + np.testing.assert_allclose(first_predictions[q], scaled_predictions[q]) + + from scipy.stats import norm + + reference = sm.WLS( + regression_data.y, + sm.add_constant(regression_data[["x"]], has_constant="add"), + weights=weights / weights.mean(), + ).fit() + prediction = reference.get_prediction(sm.add_constant(receiver, has_constant="add")) + expected = prediction.predicted_mean + norm.ppf(0.9) * np.sqrt( + prediction.var_pred_mean + reference.scale + ) + np.testing.assert_allclose(first_predictions[0.9].y, expected) + + +def test_ols_explicit_quantiles_are_deterministic_even_if_sampling_flag_set( + regression_data, +): + fitted = OLS().fit(regression_data, ["x"], ["y"]) + receiver = pd.DataFrame({"x": [0.0, 1.0]}) + exact = fitted.predict(receiver, quantiles=[0.1, 0.9]) + actual = fitted.predict(receiver, quantiles=[0.1, 0.9], random_quantile_sample=True) + for q in exact: + pd.testing.assert_frame_equal(actual[q], exact[q]) diff --git a/tests/test_models/test_statmatch_bridge.py b/tests/test_models/test_statmatch_bridge.py new file mode 100644 index 00000000..eb79072e --- /dev/null +++ b/tests/test_models/test_statmatch_bridge.py @@ -0,0 +1,184 @@ +"""Validate Python/R argument contract against the documented StatMatch API. + +These unit tests use a fake R package; real integration tests remain optional. +""" + +import contextlib +import importlib.util +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import numpy as np +import pandas as pd +import pytest + + +@pytest.fixture +def bridge(monkeypatch): + calls = [] + + class Converter: + def __add__(self, other): + return self + + def run(name, **kwargs): + calls.append((name, kwargs)) + return SimpleNamespace(rx2=lambda key: SimpleNamespace(ncol=2)) + + def fused(**kwargs): + out = kwargs["data_rec"].copy() + out["y"] = kwargs["data_don"].y.iloc[0] + return out + + package = SimpleNamespace( + NND_hotdeck=lambda **kwargs: run("NND", **kwargs), + RANDwNND_hotdeck=lambda **kwargs: run("RAND", **kwargs), + create_fused=fused, + ) + robjects = ModuleType("rpy2.robjects") + robjects.StrVector = list + robjects.FloatVector = list + robjects.IntVector = list + robjects.default_converter = Converter() + robjects.numpy2ri = SimpleNamespace(converter=Converter()) + robjects.pandas2ri = SimpleNamespace(converter=Converter()) + conversion = ModuleType("rpy2.robjects.conversion") + conversion.py2rpy = lambda x: x + conversion.rpy2py = lambda x: x + conversion.localconverter = lambda _: contextlib.nullcontext() + robjects.conversion = conversion + packages = ModuleType("rpy2.robjects.packages") + packages.importr = lambda name: package + for name, module in { + "rpy2": ModuleType("rpy2"), + "rpy2.robjects": robjects, + "rpy2.robjects.conversion": conversion, + "rpy2.robjects.packages": packages, + }.items(): + monkeypatch.setitem(sys.modules, name, module) + module_path = Path(__file__).parents[2] / "microimpute/utils/statmatch_hotdeck.py" + spec = importlib.util.spec_from_file_location( + "statmatch_contract_test", module_path + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module, calls + + +def test_weighted_matching_uses_rand_and_a_donor_column_name(bridge): + # CRAN RANDwNND.hotdeck documents weight.don as a donor column NAME, + # and cut.don="min" retains all donors tied at the nearest distance. + module, calls = bridge + donor = pd.DataFrame({"x": [1.0, 1.0], "y": [10.0, 20.0]}) + original = donor.copy() + receiver = pd.DataFrame({"x": [1.0]}) + module.nnd_hotdeck_using_rpy2( + receiver, donor, ["x"], ["y"], donor_sample_weight=np.array([1.0, 100.0]) + ) + name, kwargs = calls[0] + assert name == "RAND" + assert kwargs["cut_don"] == "min" + weight_name = kwargs["weight_don"] + assert isinstance(weight_name, str) + np.testing.assert_array_equal(kwargs["data_don"][weight_name], [1.0, 100.0]) + pd.testing.assert_frame_equal(donor, original) + + +def test_unweighted_matching_keeps_nnd(bridge): + module, calls = bridge + donor = pd.DataFrame({"x": [1.0, 2.0], "y": [10.0, 20.0]}) + module.nnd_hotdeck_using_rpy2(donor[["x"]], donor, ["x"], ["y"]) + assert calls[0][0] == "NND" + assert "weight_don" not in calls[0][1] + + +@pytest.mark.parametrize( + "weights", [[1.0], [0.0, 0.0], [-1.0, 2.0], [1.0, np.nan], [np.inf, 1.0]] +) +def test_bridge_rejects_invalid_donor_weights(bridge, weights): + module, calls = bridge + donor = pd.DataFrame({"x": [1.0, 2.0], "y": [10.0, 20.0]}) + with pytest.raises(ValueError, match="weights"): + module.nnd_hotdeck_using_rpy2( + donor[["x"]], donor, ["x"], ["y"], donor_sample_weight=np.array(weights) + ) + assert not calls + + +def test_bridge_rejects_weighted_constrained_matching(bridge): + module, calls = bridge + donor = pd.DataFrame({"x": [1.0, 2.0], "y": [10.0, 20.0]}) + with pytest.raises(ValueError, match="constrained"): + module.nnd_hotdeck_using_rpy2( + donor[["x"]], + donor, + ["x"], + ["y"], + donor_sample_weight=np.ones(2), + constrained=True, + ) + assert not calls + + +def test_fallback_matching_pairs_use_r_column_major_order(bridge): + module, calls = bridge + donor = pd.DataFrame({"x": [1.0, 2.0, 3.0], "y": [10.0, 20.0, 30.0]}) + module._get_statmatch().NND_hotdeck = lambda **kwargs: SimpleNamespace( + rx2=lambda key: np.array([3, 1, 2]) + ) + matrices = [] + + def matrix(values, nrow, ncol): + result = np.asarray(values).reshape(nrow, ncol, order="F") + matrices.append(result) + return result + + module.ro.IntVector = list + module.ro.r = SimpleNamespace(matrix=matrix) + module.nnd_hotdeck_using_rpy2(donor[["x"]], donor, ["x"], ["y"]) + np.testing.assert_array_equal(matrices[0], [[1, 3], [2, 1], [3, 2]]) + + +def test_r_rng_state_is_restored_after_seeded_call(bridge): + module, calls = bridge + donor = pd.DataFrame({"x": [1.0, 2.0], "y": [10.0, 20.0]}) + original_state = np.array([1, 2, 3]) + module.ro.globalenv = {".Random.seed": original_state} + observed_seeds = [] + + def set_seed(seed): + observed_seeds.append(seed) + module.ro.globalenv[".Random.seed"] = np.array([seed]) + + module.ro.r = {"set.seed": set_seed} + module.nnd_hotdeck_using_rpy2( + donor[["x"]], + donor, + ["x"], + ["y"], + donor_sample_weight=np.ones(2), + random_state=73, + ) + assert observed_seeds == [73] + np.testing.assert_array_equal(module.ro.globalenv[".Random.seed"], original_state) + assert "random_state" not in calls[0][1] + + +def test_seeded_r_call_restores_absent_rng_state_even_on_error(bridge): + module, calls = bridge + donor = pd.DataFrame({"x": [1.0, 2.0], "y": [10.0, 20.0]}) + module.ro.globalenv = {} + module.ro.r = { + "set.seed": lambda seed: module.ro.globalenv.update({".Random.seed": [seed]}) + } + + def fail(**kwargs): + raise RuntimeError("injected R failure") + + module._get_statmatch().NND_hotdeck = fail + with pytest.raises(RuntimeError, match="injected R failure"): + module.nnd_hotdeck_using_rpy2( + donor[["x"]], donor, ["x"], ["y"], random_state=73 + ) + assert ".Random.seed" not in module.ro.globalenv diff --git a/tests/test_models/test_zero_inflated_quantiles.py b/tests/test_models/test_zero_inflated_quantiles.py new file mode 100644 index 00000000..d7bc8022 --- /dev/null +++ b/tests/test_models/test_zero_inflated_quantiles.py @@ -0,0 +1,190 @@ +"""Exact mixture inversions and survey weights for issues #203 and #206.""" + +import numpy as np +import pandas as pd +import pytest + +from microimpute.models.zero_inflated import ( + ZeroInflatedImputer, + ZeroInflatedImputerResults, +) + + +class _KnownComponent: + def __init__(self, lower, upper): + self.lower, self.upper = lower, upper + + def predict(self, X, quantiles=None, **kwargs): + if quantiles is None: + return pd.DataFrame({"y": np.repeat(self.lower, len(X))}, index=X.index) + return { + q: pd.DataFrame( + {"y": np.repeat(self.lower + (self.upper - self.lower) * q, len(X))}, + index=X.index, + ) + for q in quantiles + } + + +class _KnownGate: + def __init__(self, classes, probabilities): + self.classes_ = np.asarray(classes) + self.probabilities = np.asarray(probabilities) + + def predict_proba(self, X): + return self.probabilities[np.asarray(X[:, 0], dtype=int)] + + +def _mixture(kind, classes, probabilities): + return ZeroInflatedImputerResults( + predictors=["x"], + imputed_variables=["y"], + seed=42, + regimes={"y": kind}, + per_variable={ + "y": { + "kind": kind, + "classifier": _KnownGate(classes, probabilities), + "positive_base": _KnownComponent(2, 6), + "negative_base": _KnownComponent(-4, -2), + } + }, + ) + + +@pytest.mark.parametrize( + "kind,classes,probabilities,expected", + [ + ("zi_positive", [0, 1], [[0.6, 0.4]], [0, 0, 0, 3, 5]), + ("zi_negative", [0, 1], [[0.6, 0.4]], [-3.5, -2.5, 0, 0, 0]), + ("sign_only", [0, 1], [[0.4, 0.6]], [-3.5, -2.5, 2 + 4 / 6, 4, 2 + 10 / 3]), + ("three_sign", [0, 1, 2], [[0.2, 0.4, 0.4]], [-3, 0, 0, 3, 5]), + ], +) +def test_quantiles_invert_ordered_mixture(kind, classes, probabilities, expected): + fitted = _mixture(kind, classes, probabilities) + query = pd.DataFrame({"x": [0] * 12}, index=np.arange(12) + 100) + quantiles = [0.1, 0.3, 0.5, 0.7, 0.9] + first = fitted.predict(query, quantiles=quantiles) + second = fitted.predict(query, quantiles=quantiles) + for q, value in zip(quantiles, expected): + np.testing.assert_allclose(first[q]["y"], value) + pd.testing.assert_frame_equal(first[q], second[q]) + assert ( + np.diff(np.column_stack([first[q]["y"] for q in quantiles]), axis=1) >= 0 + ).all() + + +def test_mixture_rescaling_varies_by_row_and_handles_zero_mass(): + fitted = _mixture( + "three_sign", [2, 0, 1], [[0.5, 0.25, 0.25], [0, 1, 0], [0, 0, 1], [1, 0, 0]] + ) + query = pd.DataFrame({"x": [0, 1, 2, 3]}, index=[9, 7, 3, 1]) + output = fitted.predict(query, quantiles=[0, 0.25, 0.5, 1]) + np.testing.assert_allclose(output[0]["y"], [-4, -4, 0, 2]) + np.testing.assert_allclose(output[0.25]["y"], [-2, -3.5, 0, 3]) + np.testing.assert_allclose(output[0.5]["y"], [0, -3, 0, 4]) + np.testing.assert_allclose(output[1]["y"], [6, -2, 0, 6]) + + +@pytest.mark.parametrize("weights_kind", ["column", "array", "series"]) +def test_numeric_weights_reach_gate_and_component(weights_kind): + # With no predictor signal, the gate must estimate the weighted prevalence; + # the positive component's median must reflect its own conditional weights. + values = np.r_[np.zeros(40), np.arange(1, 61, dtype=float)] + weights = np.r_[np.repeat(30.0, 40), np.repeat(25.0, 10), np.ones(50)] + data = pd.DataFrame( + {"x": np.zeros(100), "y": values, "w": weights}, index=np.arange(100) * 2 + ) + weight_arg = {"column": "w", "array": weights, "series": data["w"].iloc[::-1]}[ + weights_kind + ] + fitted = ZeroInflatedImputer().fit(data, ["x"], ["y"], weight_col=weight_arg) + gate = fitted._per_variable["y"]["classifier"] + np.testing.assert_allclose(gate.predict_proba([[0]])[0], [0.8, 0.2], atol=1e-8) + # q=.9 is the positive component's median because p0=.8. + output = fitted.predict(pd.DataFrame({"x": [0.0]}), quantiles=[0.5, 0.9]) + assert output[0.5]["y"].iloc[0] == 0 + assert 1 <= output[0.9]["y"].iloc[0] <= 10 + + +@pytest.mark.parametrize("bad", [np.nan, np.inf, -1.0, 0.0]) +def test_numeric_weights_are_validated(bad): + data = pd.DataFrame({"x": [0.0] * 30, "y": np.arange(30, dtype=float)}) + weights = np.ones(30) + weights[3] = bad + with pytest.raises(ValueError, match="[Ww]eight"): + ZeroInflatedImputer().fit(data, ["x"], ["y"], weight_col=weights) + + +def test_sequential_multi_target_quantiles_fail_explicitly(): + data = pd.DataFrame( + { + "x": np.arange(60, dtype=float), + "a": np.arange(60, dtype=float) + 1, + "b": np.arange(60, dtype=float) + 2, + } + ) + fitted = ZeroInflatedImputer().fit(data, ["x"], ["a", "b"], n_estimators=3) + with pytest.raises(NotImplementedError, match="sequential=False"): + fitted.predict(data[["x"]], quantiles=[0.5]) + assert fitted.predict(data[["x"]]).shape == (60, 2) + + +def test_components_with_incompatible_sign_support_fail_explicitly(): + fitted = _mixture("zi_positive", [0, 1], [[0.1, 0.9]]) + fitted._per_variable["y"]["positive_base"] = _KnownComponent(-5, 5) + with pytest.raises(ValueError, match="outside its sign support"): + fitted.predict(pd.DataFrame({"x": [0]}), quantiles=[0.2]) + + +def test_explicit_quantiles_do_not_advance_stochastic_rng(): + fitted = _mixture("zi_positive", [0, 1], [[0.6, 0.4]]) + reference = _mixture("zi_positive", [0, 1], [[0.6, 0.4]]) + query = pd.DataFrame({"x": [0] * 50}) + fitted.predict(query, quantiles=[0.1, 0.9]) + pd.testing.assert_frame_equal(fitted.predict(query), reference.predict(query)) + + +@pytest.mark.parametrize("quantiles", [[], [-0.1], [1.1], [np.nan]]) +def test_invalid_quantile_requests_fail_explicitly(quantiles): + fitted = _mixture("zi_positive", [0, 1], [[0.6, 0.4]]) + with pytest.raises(ValueError, match="[Qq]uantile"): + fitted.predict(pd.DataFrame({"x": [0]}), quantiles=quantiles) + + +def test_independent_numeric_components_have_distinct_reproducible_draws(): + rng = np.random.default_rng(4) + data = pd.DataFrame( + {"x": np.zeros(300), "a": rng.uniform(1, 3, 300), "b": rng.uniform(1, 3, 300)} + ) + query = pd.DataFrame({"x": np.zeros(500)}) + first = ( + ZeroInflatedImputer(sequential=False, seed=17) + .fit(data, ["x"], ["a", "b"], n_estimators=10) + .predict(query) + ) + second = ( + ZeroInflatedImputer(sequential=False, seed=17) + .fit(data, ["x"], ["a", "b"], n_estimators=10) + .predict(query) + ) + pd.testing.assert_frame_equal(first, second) + assert abs(first.corr().loc["a", "b"]) < 0.15 + + +def test_explicit_categorical_target_routes_to_auxiliary_model(): + data = pd.DataFrame({"x": np.arange(30, dtype=float), "y": np.tile([0, 1, 2], 10)}) + fitted = ZeroInflatedImputer().fit( + data, ["x"], ["y"], target_types={"y": "categorical"}, n_estimators=3 + ) + assert "y" not in fitted._regimes + assert fitted.predict(data[["x"]])["y"].isin([0, 1, 2]).all() + + +def test_quantile_endpoints_keep_tiny_nonzero_components(): + fitted = _mixture("three_sign", [0, 1, 2], [[1e-20, 1.0, 1e-20]]) + predicted = fitted.predict(pd.DataFrame({"x": [0]}), quantiles=[0, 0.5, 1]) + assert predicted[0]["y"].iloc[0] == -4 + assert predicted[0.5]["y"].iloc[0] == 0 + assert predicted[1]["y"].iloc[0] == 6 diff --git a/tests/test_predictor_analysis.py b/tests/test_predictor_analysis.py index 9fd32ce8..b418ba2d 100644 --- a/tests/test_predictor_analysis.py +++ b/tests/test_predictor_analysis.py @@ -12,6 +12,8 @@ from sklearn.datasets import make_classification, make_regression from microimpute.evaluations.predictor_analysis import ( + _compute_losses_from_predictions, + _evaluate_model_performance, compute_predictor_correlations, leave_one_out_analysis, progressive_predictor_inclusion, @@ -19,6 +21,98 @@ from microimpute.models import OLS, QRF, QuantReg +@pytest.mark.parametrize("truth", [["a", "a"], ["a", "b"]]) +def test_predictor_analysis_constant_categorical_forecast(truth): + """A constant training target has an exact point-mass probability forecast.""" + from sklearn.metrics import log_loss + + training = pd.DataFrame({"x": np.arange(10.0), "label": ["a"] * 10}) + testing = pd.DataFrame({"x": [1.0, 2.0], "label": truth}) + result = _evaluate_model_performance( + training, testing, ["x"], ["label"], QRF, None, [0.5], 42 + ) + expected = log_loss(truth, np.array([[1.0, 0.0], [1.0, 0.0]]), labels=["a", "b"]) + assert result["log_loss"] == pytest.approx(expected, abs=1e-12) + + +def test_predictor_analysis_scores_real_probabilities(): + """Equal-probability binary forecasts score log(2), regardless of hard labels.""" + predictions = { + 0.5: pd.DataFrame({"choice": ["a", "a", "a", "a"]}), + "probabilities": { + "choice": { + "probabilities": np.full((4, 2), 0.5), + "classes": np.array(["b", "a"]), + } + }, + } + actual = _compute_losses_from_predictions( + predictions, pd.DataFrame({"choice": ["a", "b", "a", "b"]}), ["choice"], [0.5] + ) + assert actual["log_loss"] == pytest.approx(np.log(2)) + + +def test_predictor_analysis_refuses_missing_probabilities(): + with pytest.raises(ValueError, match="probabilit"): + _compute_losses_from_predictions( + {0.5: pd.DataFrame({"choice": ["a", "a"]})}, + pd.DataFrame({"choice": ["a", "b"]}), + ["choice"], + [0.5], + ) + + +@pytest.mark.parametrize("n", [200, 1000, 5000]) +def test_normalized_mi_identical_continuous_variables(n): + """Identical discretizations have MI equal to their entropy, at any size.""" + x = np.random.default_rng(3).normal(size=n) + result = compute_predictor_correlations( + pd.DataFrame({"a": x, "b": x, "target": x}), + ["a", "b"], + ["target"], + method="mutual_info", + ) + assert result["mutual_info"].loc["a", "b"] == pytest.approx(1.0) + assert result["predictor_target_mi"].loc["a", "target"] == pytest.approx(1.0) + + +def test_normalized_mi_categorical_independence_and_constant(): + """A balanced Cartesian product is independent; constants convey no information.""" + data = pd.DataFrame( + { + "a": ["a", "a", "b", "b"] * 30, + "b": ["x", "y", "x", "y"] * 30, + "copy": ["A", "A", "B", "B"] * 30, + "constant": [1.0] * 120, + } + ) + matrix = compute_predictor_correlations(data, list(data), method="mutual_info")[ + "mutual_info" + ] + assert matrix.loc["a", "copy"] == pytest.approx(1.0) + assert matrix.loc["a", "b"] == pytest.approx(0.0, abs=1e-12) + assert (matrix.loc["constant"] == 0.0).all() + np.testing.assert_allclose(matrix, matrix.T) + + +def test_normalized_mi_pairwise_missing_and_column_order(): + data = pd.DataFrame( + { + "x": [1.0, 2.0, 3.0, 4.0, np.nan], + "y": [1.0, 2.0, 3.0, 4.0, 90.0], + "group": ["a", "a", "b", "b", None], + } + ) + first = compute_predictor_correlations(data, list(data), method="mutual_info")[ + "mutual_info" + ] + second = compute_predictor_correlations( + data, list(reversed(data.columns)), method="mutual_info" + )["mutual_info"] + assert first.loc["x", "y"] == pytest.approx(1.0) + np.testing.assert_allclose(first, second.loc[first.index, first.columns]) + + @pytest.fixture def sample_regression_data(): """Create sample regression data for testing.""" diff --git a/tests/test_publication_pipeline.py b/tests/test_publication_pipeline.py new file mode 100644 index 00000000..925511fb --- /dev/null +++ b/tests/test_publication_pipeline.py @@ -0,0 +1,316 @@ +"""Behavioral regressions for publication audit issues #202, #206, #209, #213.""" + +import numpy as np +import pandas as pd +import pytest +from sklearn.metrics import log_loss as sklearn_log_loss +from sklearn.model_selection import train_test_split + +from microimpute.comparisons.autoimpute_helpers import prepare_data_for_imputation +from microimpute.comparisons.metrics import compare_metrics, compute_loss +from microimpute.comparisons.validation import validate_imputation_inputs +from microimpute.models import OLS +from microimpute.utils.data import preprocess_data + + +def test_receiver_uses_donor_normalization(): + donor = pd.DataFrame({"x": np.arange(20.0), "y": np.arange(20.0) * 2}) + receiver = pd.DataFrame({"x": [30.0]}) + train, test, _ = prepare_data_for_imputation( + donor, receiver, ["x"], ["y"], None, 1.0, 0.0, preprocessing={"x": "normalize"} + ) + assert test.x.iloc[0] == pytest.approx((30 - donor.x.mean()) / donor.x.std()) + assert train.x.mean() == pytest.approx(0) + + +def test_preprocess_split_fits_training_statistics_only(): + data = pd.DataFrame({"x": np.arange(30.0) ** 2}) + train, test, params = preprocess_data(data, normalize=["x"], random_state=3) + raw_train, raw_test = train_test_split( + data, train_size=0.8, test_size=0.2, random_state=3 + ) + assert params["normalization"]["x"]["mean"] == pytest.approx(raw_train.x.mean()) + np.testing.assert_allclose( + test.x, (raw_test.x - raw_train.x.mean()) / raw_train.x.std() + ) + + +def test_count_target_stays_numeric_and_explicit_override_is_available(): + donor = pd.DataFrame({"x": np.arange(60.0), "y": np.tile(np.arange(6), 10)}) + model = OLS() + model.fit(donor, ["x"], ["y"]) + assert model.numeric_targets == ["y"] + model.fit(donor, ["x"], ["y"], target_types={"y": "categorical"}) + assert "y" in model.categorical_targets + assert model.numeric_targets == [] + + +def test_logloss_requires_real_probabilities(): + with pytest.raises(ValueError, match="probabilit"): + compute_loss(np.array([0, 1]), np.array([0, 1]), "log_loss") + actual = compute_loss(np.array([0, 1]), np.array([0.25, 0.75]), "log_loss")[1] + assert actual == pytest.approx(-np.log(0.75)) + + +def test_compare_metrics_uses_probabilities_and_custom_quantiles(): + y = pd.DataFrame({"label": ["a", "b"], "count": [2, 4]}) + probabilities = np.array([[0.8, 0.2], [0.3, 0.7]]) + predictions = { + "OLS": { + 0.25: pd.DataFrame({"label": ["a", "b"], "count": [1, 3]}), + "probabilities": { + "label": { + "probabilities": probabilities, + "classes": np.array(["a", "b"]), + } + }, + } + } + result = compare_metrics(y, predictions, ["label", "count"]) + assert result.loc[result["Imputed Variable"] == "label", "Loss"].iloc[ + 0 + ] == pytest.approx(sklearn_log_loss(y.label, probabilities)) + assert ( + result.loc[result["Imputed Variable"] == "count", "Percentile"].iloc[0] == 0.25 + ) + + +@pytest.mark.parametrize( + "predictors,targets,receiver,message", + [ + (["x", "x"], ["y"], pd.DataFrame({"x": [1]}), "Duplicate"), + (["x"], ["x"], pd.DataFrame({"x": [1]}), "overlap"), + (["x"], ["y"], pd.DataFrame({"x": ["1"]}), "dtype"), + ], +) +def test_invalid_imputation_inputs_are_explicit(predictors, targets, receiver, message): + with pytest.raises(ValueError, match=message): + validate_imputation_inputs( + pd.DataFrame({"x": [1.0, 2.0], "y": [3, 4]}), receiver, predictors, targets + ) + + +def test_cv_preprocessing_uses_each_fold_and_scores_original_scale(): + from microimpute.evaluations.cross_validation import cross_validate_model + + data = pd.DataFrame({"x": np.arange(60.0), "y": 10 + np.arange(60.0) * 3}) + result = cross_validate_model( + OLS, + data, + ["x"], + ["y"], + n_splits=3, + quantiles=[0.5], + preprocessing={"x": "normalize", "y": "normalize"}, + ) + assert result["quantile_loss"]["mean_test"] < 1e-10 + + +def test_autoimpute_train_fraction_seed_custom_quantiles_and_mixed_transform(): + from microimpute import autoimpute + + data = pd.DataFrame({"x": np.arange(40.0), "y": 4 + np.arange(40.0) * 2}) + result = autoimpute( + data, + pd.DataFrame({"x": [50.0]}), + ["x"], + ["y"], + models=[OLS], + train_size=0.5, + random_state=19, + k_folds=2, + imputation_quantiles=[0.25, 0.75], + preprocessing={"x": "normalize", "y": "normalize"}, + ) + fitted = result.fitted_models["best_method"] + assert fitted.models["y"].model.nobs == 20 + assert fitted.seed == 19 + assert {0.25, 0.5, 0.75} <= set(result.imputations["best_method"]) + assert result.receiver_data.y.iloc[0] == pytest.approx(104) + + +def test_cv_refits_tuning_on_all_rows_instead_of_selecting_test_fold(monkeypatch): + import importlib + import joblib + + cv = importlib.import_module("microimpute.evaluations.cross_validation") + calls = [] + + class ProbeResults: + def predict(self, data, quantiles, **kwargs): + return { + q: pd.DataFrame({"y": np.zeros(len(data))}, index=data.index) + for q in quantiles + } + + def fake_fit(model, model_class, data, *args): + calls.append(data.index.to_list()) + return ProbeResults(), {"training_rows": len(data)} + + monkeypatch.setattr(cv, "_fit_model_for_fold", fake_fit) + data = pd.DataFrame({"x": np.arange(30.0), "y": np.arange(30.0)}) + with joblib.parallel_backend("threading"): + _, params = cv.cross_validate_model( + OLS, + data, + ["x"], + ["y"], + n_splits=3, + quantiles=[0.25], + tune_hyperparameters=True, + ) + assert params == {"training_rows": 30} + assert sorted(map(len, calls)) == [20, 20, 20, 30] + + +def test_explicit_categorical_target_is_stable_through_autoimpute(): + from microimpute import autoimpute + + data = pd.DataFrame( + {"x": np.tile([0.2, 1.2, 2.2], 20), "y": np.tile([0, 1, 2], 20)} + ) + result = autoimpute( + data, + data[["x"]].iloc[:3], + ["x"], + ["y"], + models=[OLS], + k_folds=2, + target_types={"y": "categorical"}, + train_size=1, + ) + assert result.cv_results["OLS"]["log_loss"]["variables"] == ["y"] + assert result.cv_results["OLS"]["quantile_loss"]["variables"] == [] + + +def test_predict_accepts_int_float_equivalence_and_rejects_strings(): + data = pd.DataFrame({"x": np.arange(20.0), "y": np.arange(20.0) * 2}) + fitted = OLS().fit(data, ["x"], ["y"]) + assert fitted.predict(pd.DataFrame({"x": [3]}), [0.5])[0.5].y.iloc[ + 0 + ] == pytest.approx(6) + with pytest.raises(ValueError, match="dtype"): + fitted.predict(pd.DataFrame({"x": ["3"]}), [0.5]) + + +def test_logloss_scores_constant_and_unseen_classes_without_fabricating_probs(): + probabilities = np.ones((2, 1)) + assert ( + compute_loss( + np.array(["a", "a"]), probabilities, "log_loss", labels=np.array(["a"]) + )[1] + == 0 + ) + actual = compute_loss( + np.array(["a", "b"]), probabilities, "log_loss", labels=np.array(["a"]) + )[1] + expected = sklearn_log_loss( + ["a", "b"], np.array([[1.0, 0.0], [1.0, 0.0]]), labels=["a", "b"] + ) + assert actual == pytest.approx(expected) + + +def test_returned_fitted_model_replays_donor_preprocessing_on_raw_receivers(): + from microimpute import autoimpute + + donor = pd.DataFrame({"x": np.arange(40.0), "y": 10 + np.arange(40.0) * 2}) + receiver = pd.DataFrame({"x": [50.0, 70.0]}) + result = autoimpute( + donor, + receiver, + ["x"], + ["y"], + models=[OLS], + train_size=1, + k_folds=2, + preprocessing={"x": "normalize", "y": "normalize"}, + ) + replay = result.fitted_models["best_method"].predict(receiver, quantiles=[0.5])[0.5] + pd.testing.assert_frame_equal(replay, result.imputations["best_method"]) + np.testing.assert_allclose(replay.y, [110, 150]) + + +def test_constant_categorical_get_imputations_has_exact_probabilities(): + from microimpute.comparisons import get_imputations + + donor = pd.DataFrame({"x": np.arange(20.0), "label": ["a"] * 20}) + receiver = pd.DataFrame({"x": [21.0, 22.0], "label": ["a", "b"]}) + predictions = get_imputations([OLS], donor, receiver, ["x"], ["label"], [0.5]) + assert predictions["OLS"]["probabilities"]["label"]["probabilities"].shape == (2, 1) + result = compare_metrics(receiver[["label"]], predictions, ["label"]) + expected = sklearn_log_loss(["a", "b"], [[1.0, 0.0], [1.0, 0.0]], labels=["a", "b"]) + assert result.loc[result["Imputed Variable"] == "label", "Loss"].iloc[ + 0 + ] == pytest.approx(expected) + + +def test_returned_model_preserves_probabilities_and_partial_target_transforms(): + from microimpute import autoimpute + + donor = pd.DataFrame( + { + "x": np.arange(40.0), + "y": 10 + np.arange(40.0) * 2, + "label": np.tile(["a", "b"], 20), + } + ) + receiver = pd.DataFrame({"x": [50.0, 70.0]}) + result = autoimpute( + donor, + receiver, + ["x"], + ["y", "label"], + models=[OLS], + train_size=1, + k_folds=2, + preprocessing={"x": "normalize", "y": "normalize"}, + imputation_quantiles=[0.5], + ) + replay = result.fitted_models["best_method"].predict( + receiver, quantiles=[0.5], return_probs=True + ) + pd.testing.assert_frame_equal(replay[0.5], result.imputations["best_method"][0.5]) + np.testing.assert_allclose( + replay["probabilities"]["label"]["probabilities"], + result.imputations["best_method"]["probabilities"]["label"]["probabilities"], + ) + + +def test_qrf_comparison_quantiles_equal_independent_single_target_forecasts(): + from microimpute.comparisons import get_imputations + from microimpute.models import QRF + + rng = np.random.default_rng(31) + donor = pd.DataFrame({"x": rng.normal(size=160), "first": rng.normal(size=160)}) + donor["second"] = donor["first"] * 4 + rng.normal(size=160) + receiver = pd.DataFrame({"x": np.linspace(-1, 1, 5)}) + quantiles = [0.1, 0.9] + together = get_imputations( + [QRF], donor, receiver, ["x"], ["first", "second"], quantiles + )["QRF"] + separate = get_imputations([QRF], donor, receiver, ["x"], ["second"], quantiles)[ + "QRF" + ] + for q in quantiles: + np.testing.assert_allclose(together[q]["second"], separate[q]["second"]) + + +def test_preprocessing_replay_ignores_receiver_target_placeholders(): + from microimpute import autoimpute + + donor = pd.DataFrame({"x": np.linspace(0, 2, 40)}) + donor["y"] = np.exp(1 + donor.x) + receiver = pd.DataFrame({"x": [1.5], "y": [0.0]}) + result = autoimpute( + donor, + receiver, + ["x"], + ["y"], + models=[OLS], + train_size=1, + k_folds=2, + preprocessing={"y": "log"}, + ) + replay = result.fitted_models["best_method"].predict(receiver, quantiles=[0.5])[0.5] + assert replay.y.iloc[0] == pytest.approx(np.exp(2.5)) + assert receiver.y.iloc[0] == 0 diff --git a/tests/test_quantile_comparison.py b/tests/test_quantile_comparison.py index 9eb93199..71542410 100644 --- a/tests/test_quantile_comparison.py +++ b/tests/test_quantile_comparison.py @@ -170,8 +170,6 @@ def test_multiple_imputed_variables(split_data: tuple) -> None: Y_test = X_test[imputed_variables] model_classes = [OLS, QRF, QuantReg] - if HAS_MATCHING: - model_classes.append(Matching) method_imputations = get_imputations( model_classes, X_train, X_test, predictors, imputed_variables ) @@ -367,9 +365,9 @@ def test_log_loss_for_categorical_variables() -> None: "x1": np.random.randn(n_train), "x2": np.random.randn(n_train), # Binary categorical variable - "binary_cat": np.random.choice([0, 1], size=n_train), + "binary_cat": np.random.choice([False, True], size=n_train), # Multi-class categorical variable (3 classes) - "multi_cat": np.random.choice([0, 1, 2], size=n_train), + "multi_cat": pd.Categorical(np.random.choice([0, 1, 2], size=n_train)), # String categorical variable "string_cat": np.random.choice(["A", "B", "C"], size=n_train), # Numerical variable for comparison @@ -382,8 +380,8 @@ def test_log_loss_for_categorical_variables() -> None: { "x1": np.random.randn(n_test), "x2": np.random.randn(n_test), - "binary_cat": np.random.choice([0, 1], size=n_test), - "multi_cat": np.random.choice([0, 1, 2], size=n_test), + "binary_cat": np.random.choice([False, True], size=n_test), + "multi_cat": pd.Categorical(np.random.choice([0, 1, 2], size=n_test)), "string_cat": np.random.choice(["A", "B", "C"], size=n_test), "numerical": np.random.randn(n_test), } @@ -471,7 +469,7 @@ def test_mixed_variable_types() -> None: "predictor1": np.random.randn(n_train), "predictor2": np.random.choice(["X", "Y", "Z"], size=n_train), "numerical_target": np.random.randn(n_train), - "categorical_target": np.random.choice([0, 1], size=n_train), + "categorical_target": np.random.choice([False, True], size=n_train), } ) @@ -480,7 +478,7 @@ def test_mixed_variable_types() -> None: "predictor1": np.random.randn(n_test), "predictor2": np.random.choice(["X", "Y", "Z"], size=n_test), "numerical_target": np.random.randn(n_test), - "categorical_target": np.random.choice([0, 1], size=n_test), + "categorical_target": np.random.choice([False, True], size=n_test), } ) diff --git a/tests/test_type_handling.py b/tests/test_type_handling.py index bb4b1410..b6fb6568 100644 --- a/tests/test_type_handling.py +++ b/tests/test_type_handling.py @@ -21,9 +21,9 @@ def test_is_boolean_variable_true_for_bool_dtype() -> None: assert VariableTypeDetector.is_boolean_variable(s) is True -def test_is_boolean_variable_true_for_int_0_1() -> None: +def test_is_boolean_variable_false_for_int_0_1() -> None: s = pd.Series([0, 1, 0, 1], dtype=int) - assert VariableTypeDetector.is_boolean_variable(s) is True + assert VariableTypeDetector.is_boolean_variable(s) is False def test_is_boolean_variable_false_for_float_0_1() -> None: From b2c7ce3d8336669bc1ee1c3dfbb3a75569c6db67 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 10:38:06 +0100 Subject: [PATCH 5/5] Fix issues from review: preserve numeric and weight semantics --- changelog.d/219-adversarial.fixed.md | 1 + microimpute/comparisons/autoimpute.py | 4 +- microimpute/comparisons/autoimpute_helpers.py | 6 +- microimpute/comparisons/imputations.py | 13 +- microimpute/comparisons/metrics.py | 4 +- microimpute/evaluations/cross_validation.py | 9 +- microimpute/models/matching.py | 4 +- microimpute/models/ols.py | 2 +- microimpute/utils/type_handling.py | 4 + tests/test_adversarial_regressions.py | 387 ++++++++++++++++++ 10 files changed, 422 insertions(+), 12 deletions(-) create mode 100644 changelog.d/219-adversarial.fixed.md create mode 100644 tests/test_adversarial_regressions.py diff --git a/changelog.d/219-adversarial.fixed.md b/changelog.d/219-adversarial.fixed.md new file mode 100644 index 00000000..66ac97cb --- /dev/null +++ b/changelog.d/219-adversarial.fixed.md @@ -0,0 +1 @@ +Prevent integer overflow in quantile loss and Matching tuning; consistently honor numeric boolean targets, preserve survey weights when transforming a weight predictor, retain default constant-category probabilities, and support nullable numeric OLS predictors. diff --git a/microimpute/comparisons/autoimpute.py b/microimpute/comparisons/autoimpute.py index e5d5ff17..c8878f6d 100644 --- a/microimpute/comparisons/autoimpute.py +++ b/microimpute/comparisons/autoimpute.py @@ -338,7 +338,7 @@ def _generate_imputations_for_all_models( imputing_data, predictors, imputed_variables, - weight_col, + donor_data[weight_col].copy() if weight_col is not None else None, imputation_q, model_hyperparams, log_level, @@ -621,7 +621,7 @@ def autoimpute( imputing_data, predictors, imputed_variables, - weight_col, + donor_data[weight_col].copy() if weight_col is not None else None, imputation_q, model_hyperparams, log_level, diff --git a/microimpute/comparisons/autoimpute_helpers.py b/microimpute/comparisons/autoimpute_helpers.py index 37379c39..1e63089f 100644 --- a/microimpute/comparisons/autoimpute_helpers.py +++ b/microimpute/comparisons/autoimpute_helpers.py @@ -14,7 +14,7 @@ """ import logging -from typing import Any, Dict, List, Optional, Tuple, Type +from typing import Any, Dict, List, Optional, Tuple, Type, Union import numpy as np import pandas as pd @@ -165,7 +165,7 @@ def prepare_data_for_imputation( ) training_data = transformed_training - if weight_col: + if weight_col and weight_col not in all_training_cols: training_data[weight_col] = donor_data[weight_col] imputing_data = transformed_imputing @@ -283,7 +283,7 @@ def fit_and_predict_model( imputing_data: pd.DataFrame, predictors: List[str], imputed_variables: List[str], - weight_col: Optional[str], + weight_col: Optional[Union[str, np.ndarray, pd.Series]], quantile: float, hyperparams: Optional[Dict[str, Any]] = None, log_level: str = "WARNING", diff --git a/microimpute/comparisons/imputations.py b/microimpute/comparisons/imputations.py index defdb44e..c9b35117 100644 --- a/microimpute/comparisons/imputations.py +++ b/microimpute/comparisons/imputations.py @@ -98,10 +98,16 @@ def get_imputations( predictors, imputed_variables, quantiles=quantiles, + target_types=target_types, ) else: log.info(f"Fitting {model_name}") - fitted_model = model.fit(X_train, predictors, imputed_variables) + fitted_model = model.fit( + X_train, + predictors, + imputed_variables, + target_types=target_types, + ) # Get predictions log.info(f"Generating predictions with {model_name}") @@ -118,6 +124,11 @@ def get_imputations( get_metric_for_variable_type(X_train[variable], variable) == "log_loss" ): + # Default point predictions may be returned as a frame. + # Probability metadata uses the same median-keyed + # container as nonconstant categorical predictions. + if isinstance(imputations, pd.DataFrame): + imputations = {0.5: imputations} imputations.setdefault("probabilities", {})[variable] = { "probabilities": np.ones((len(X_test), 1)), "classes": np.asarray([info["value"]]), diff --git a/microimpute/comparisons/metrics.py b/microimpute/comparisons/metrics.py index d8f0c34d..cde92519 100644 --- a/microimpute/comparisons/metrics.py +++ b/microimpute/comparisons/metrics.py @@ -64,7 +64,9 @@ def quantile_loss(q: float, y: np.ndarray, f: np.ndarray) -> np.ndarray: Returns: Array of quantile losses. """ - e = y - f + # Integer counts may be unsigned or narrower than their residuals. Promote + # both operands before subtraction so neither wrapping nor overflow occurs. + e = np.asarray(y, dtype=float) - np.asarray(f, dtype=float) return np.maximum(q * e, (q - 1) * e) diff --git a/microimpute/evaluations/cross_validation.py b/microimpute/evaluations/cross_validation.py index c36cffff..acf7e02c 100644 --- a/microimpute/evaluations/cross_validation.py +++ b/microimpute/evaluations/cross_validation.py @@ -68,6 +68,9 @@ def _process_single_fold( # Split data for this fold train_data = data.iloc[train_idx] test_data = data.iloc[test_idx] + # The sampling mass must retain its original units even when the same + # column is transformed as a predictor. Its index also survives fit filters. + fit_weights = train_data[weight_col].copy() if weight_col is not None else None # Store actual values for this fold organized by variable train_y = {var: train_data[var].values for var in imputed_variables} @@ -90,7 +93,7 @@ def _process_single_fold( train_data, predictors, imputed_variables, - weight_col, + fit_weights, quantiles, model_hyperparams, tune_hyperparameters, @@ -158,7 +161,7 @@ def _fit_model_for_fold( train_data: pd.DataFrame, predictors: List[str], imputed_variables: List[str], - weight_col: Optional[str], + weight_col: Optional[Union[str, np.ndarray, pd.Series]], quantiles: List[float], model_hyperparams: Optional[dict], tune_hyperparameters: bool, @@ -682,7 +685,7 @@ def cross_validate_model( tuning_data, predictors, imputed_variables, - weight_col, + data[weight_col].copy() if weight_col is not None else None, quantiles, model_hyperparams, True, diff --git a/microimpute/models/matching.py b/microimpute/models/matching.py index 07fb1117..e8522bc7 100644 --- a/microimpute/models/matching.py +++ b/microimpute/models/matching.py @@ -513,7 +513,9 @@ def objective(trial: optuna.Trial) -> float: if var in discrete_targets: errors.append(float(np.mean(actual != estimate))) else: - if not np.isfinite(estimate.astype(float)).all(): + actual = actual.astype(float) + estimate = estimate.astype(float) + if not np.isfinite(estimate).all(): raise optuna.TrialPruned( "Matching returned nonfinite numeric predictions" ) diff --git a/microimpute/models/ols.py b/microimpute/models/ols.py index df6f0bf4..91b7e891 100644 --- a/microimpute/models/ols.py +++ b/microimpute/models/ols.py @@ -256,7 +256,7 @@ def _predict_variable( # rows far from the training centroid; at extreme quantiles # (0.01, 0.99) the under-dispersion is material. X_test_with_const = sm.add_constant( - X_test[self.predictors], has_constant="add" + X_test[self.predictors].astype(float), has_constant="add" ) prediction = model.model.get_prediction(X_test_with_const) # var_pred_mean is the leverage term (x' (X'X)^-1 x) * scale; diff --git a/microimpute/utils/type_handling.py b/microimpute/utils/type_handling.py index f3a1f556..308f4c1a 100644 --- a/microimpute/utils/type_handling.py +++ b/microimpute/utils/type_handling.py @@ -445,6 +445,10 @@ def declare_target_types( f"Boolean target '{variable}' must contain only 0 and 1" ) result[variable] = result[variable].astype(bool) + elif pd.api.types.is_bool_dtype(result[variable]): + # A numeric declaration must survive dtype-based routing in every + # public caller, including preprocessing and metric selection. + result[variable] = result[variable].astype(float) elif not pd.api.types.is_numeric_dtype(result[variable]): if isinstance( result[variable].dtype, pd.CategoricalDtype diff --git a/tests/test_adversarial_regressions.py b/tests/test_adversarial_regressions.py new file mode 100644 index 00000000..0aec965c --- /dev/null +++ b/tests/test_adversarial_regressions.py @@ -0,0 +1,387 @@ +"""Independent dtype, weight and prediction-contract regressions for PR 219.""" + +import importlib + +import joblib +import numpy as np +import optuna +import pandas as pd +import pytest +from scipy.spatial.distance import cdist +from sklearn.model_selection import KFold + +from microimpute import ( + autoimpute, + compare_metrics, + cross_validate_model, + get_imputations, +) +from microimpute.comparisons.metrics import quantile_loss +from microimpute.models import OLS, QRF, QuantReg +from microimpute.models.matching import Matching + + +@pytest.fixture(autouse=True) +def threaded_cv(): + with joblib.parallel_backend("threading", n_jobs=2): + yield + + +@pytest.mark.parametrize("dtype", ["uint8", "uint16", "int8", "int16"]) +def test_integer_pinball_loss_matches_real_arithmetic(dtype): + limits = np.iinfo(dtype) + truth = np.array([limits.min, limits.max], dtype=dtype) + prediction = truth[::-1].copy() + # At q=.25, errors of +/- the full integer range have these exact losses. + width = int(limits.max) - int(limits.min) + np.testing.assert_allclose( + quantile_loss(0.25, truth, prediction), [0.75 * width, 0.25 * width] + ) + + +@pytest.mark.parametrize("model_class", [OLS, QRF, QuantReg]) +def test_unsigned_constant_comparison_uses_numeric_loss(model_class): + donor = pd.DataFrame({"x": np.arange(20.0), "y": np.full(20, 3, dtype="uint8")}) + truth = pd.DataFrame({"x": [21.0, 22.0], "y": np.array([1, 5], dtype="uint8")}) + predictions = get_imputations([model_class], donor, truth, ["x"], ["y"], [0.5]) + scores = compare_metrics(truth, predictions, ["y"]) + assert set(scores.Metric) == {"quantile_loss"} + # Both errors are two, so the median pinball loss is one. + np.testing.assert_allclose(scores.Loss, 1.0) + + +def test_unsigned_cv_matches_float_control(): + donor = pd.DataFrame( + {"x": np.arange(8.0), "y": np.array([3] * 7 + [1], dtype="uint8")} + ) + unsigned = cross_validate_model( + OLS, donor, ["x"], ["y"], n_splits=2, quantiles=[0.5] + ) + floating = cross_validate_model( + OLS, donor.astype({"y": float}), ["x"], ["y"], n_splits=2, quantiles=[0.5] + ) + pd.testing.assert_frame_equal( + unsigned["quantile_loss"]["results"], floating["quantile_loss"]["results"] + ) + + +def _nearest_donor( + receiver, donor, matching_variables, z_variables, dist_fun="Manhattan", **kwargs +): + train = donor[matching_variables].to_numpy(dtype=float) + test = receiver[matching_variables].to_numpy(dtype=float) + if dist_fun == "Gower": + scale = np.ptp(train, axis=0) + scale[scale == 0] = 1 + distances = cdist(test / scale, train / scale, metric="cityblock") + else: + metric = { + "Manhattan": "cityblock", + "Euclidean": "euclidean", + "minimax": "chebyshev", + "Mahalanobis": "mahalanobis", + }[dist_fun] + options = ( + {"VI": np.linalg.inv(np.cov(train, rowvar=False))} + if metric == "mahalanobis" + else {} + ) + distances = cdist(test, train, metric=metric, **options) + nearest = np.argmin(distances, axis=1) + result = receiver.copy() + for variable in z_variables: + result[variable] = donor[variable].to_numpy()[nearest] + return result, result + + +def test_matching_tunes_with_numeric_donor_error(monkeypatch): + rng = np.random.default_rng(0) + features = rng.normal(size=(45, 2)) * [1, 5] + donor = pd.DataFrame( + { + "x": features[:, 0], + "z": features[:, 1], + "y": ((features[:, 0] + rng.normal(size=45)) > -0.8).astype("uint8"), + } + ) + studies = [] + create_study = optuna.create_study + + def capture_study(*args, **kwargs): + study = create_study(*args, **kwargs) + studies.append(study) + return study + + monkeypatch.setattr(optuna, "create_study", capture_study) + _, params = Matching(matching_hotdeck=_nearest_donor, seed=42).fit( + donor, ["x", "z"], ["y"], tune_hyperparameters=True + ) + oracle = {} + for distance in {trial.params["dist_fun"] for trial in studies[0].trials}: + fold_losses = [] + for train, test in KFold(n_splits=3, shuffle=True, random_state=42).split( + donor + ): + training, held_out = donor.iloc[train], donor.iloc[test] + predicted, _ = _nearest_donor( + held_out.drop(columns="y"), + training, + ["x", "z"], + ["y"], + dist_fun=distance, + ) + errors = [ + abs(int(actual) - int(estimate)) + for actual, estimate in zip(held_out.y, predicted.y) + ] + fold_losses.append(np.mean(errors) / training.y.std(ddof=0)) + oracle[distance] = np.mean(fold_losses) + assert params["dist_fun"] == min(oracle, key=oracle.get) + for trial in studies[0].trials: + assert trial.value == pytest.approx(oracle[trial.params["dist_fun"]]) + + +@pytest.fixture +def boolean_donor(): + rng = np.random.default_rng(57) + return pd.DataFrame({"x": np.linspace(-2, 2, 90), "y": rng.uniform(size=90) < 0.4}) + + +@pytest.mark.parametrize("dtype", ["bool", "boolean"]) +def test_explicit_numeric_boolean_agrees_across_public_paths(boolean_donor, dtype): + donor = boolean_donor.astype({"y": dtype}) + floating = donor.astype({"y": float}) + receiver = donor.iloc[[1, 17, 61]] + quantiles = [0.1, 0.5, 0.9] + direct = ( + OLS() + .fit(donor, ["x"], ["y"], target_types={"y": "numeric"}) + .predict(receiver, quantiles) + ) + generated = get_imputations( + [OLS], donor, receiver, ["x"], ["y"], quantiles, target_types={"y": "numeric"} + ) + control = get_imputations([OLS], floating, receiver, ["x"], ["y"], quantiles) + for q in quantiles: + pd.testing.assert_frame_equal(direct[q], control["OLS"][q]) + pd.testing.assert_frame_equal(generated["OLS"][q], control["OLS"][q]) + scores = compare_metrics(receiver, generated, ["y"], target_types={"y": "numeric"}) + expected = compare_metrics(receiver.astype({"y": float}), control, ["y"]) + pd.testing.assert_frame_equal(scores, expected) + assert set(scores.Metric) == {"quantile_loss"} + + +@pytest.mark.parametrize("dtype", ["bool", "boolean"]) +def test_explicit_numeric_boolean_cv_matches_float(boolean_donor, dtype): + donor = boolean_donor.astype({"y": dtype}) + result = cross_validate_model( + OLS, + donor, + ["x"], + ["y"], + n_splits=2, + quantiles=[0.5], + target_types={"y": "numeric"}, + ) + control = cross_validate_model( + OLS, donor.astype({"y": float}), ["x"], ["y"], n_splits=2, quantiles=[0.5] + ) + assert result["quantile_loss"]["variables"] == ["y"] + assert result["log_loss"]["variables"] == [] + pd.testing.assert_frame_equal( + result["quantile_loss"]["results"], control["quantile_loss"]["results"] + ) + + +@pytest.mark.parametrize("declared_type", [None, "categorical", "bool"]) +def test_boolean_classification_and_probabilities_remain_available( + boolean_donor, declared_type +): + declaration = None if declared_type is None else {"y": declared_type} + truth = boolean_donor.iloc[:6] + predictions = get_imputations( + [OLS], boolean_donor, truth, ["x"], ["y"], [0.5], target_types=declaration + ) + probabilities = predictions["OLS"]["probabilities"]["y"] + assert probabilities["probabilities"].shape == (6, 2) + np.testing.assert_allclose(probabilities["probabilities"].sum(axis=1), 1) + scores = compare_metrics(truth, predictions, ["y"], target_types=declaration) + assert set(scores.Metric) == {"log_loss"} + cv = cross_validate_model( + OLS, + boolean_donor, + ["x"], + ["y"], + n_splits=2, + quantiles=[0.5], + target_types=declaration, + ) + assert cv["log_loss"]["variables"] == ["y"] + assert cv["quantile_loss"]["variables"] == [] + + +def test_weight_predictor_normalization_preserves_exact_linear_oracle(): + donor = pd.DataFrame({"w": np.arange(1.0, 61.0), "y": 1 + 3 * np.arange(1.0, 61.0)}) + receiver = pd.DataFrame({"w": [61.0, 65.0]}, index=[101, 107]) + result = autoimpute( + donor, + receiver, + ["w"], + ["y"], + weight_col="w", + models=[OLS], + preprocessing={"w": "normalize"}, + train_size=1, + k_folds=2, + ) + np.testing.assert_allclose(result.receiver_data.y, [184, 196]) + assert result.cv_results["OLS"]["quantile_loss"]["mean_test"] < 1e-10 + replay = result.fitted_models["best_method"].predict(receiver, [0.5])[0.5] + pd.testing.assert_frame_equal(replay, result.imputations["best_method"]) + + +@pytest.mark.parametrize("transform", ["normalize", "log", "asinh"]) +def test_transformed_weight_predictor_matches_separate_weights_after_sampling_and_filtering( + transform, +): + rng = np.random.default_rng(123) + donor = pd.DataFrame( + {"w": np.arange(1.0, 61.0), "y": rng.normal(size=60) + np.arange(60.0) ** 0.5}, + index=np.arange(60) * 7 + 3, + ) + receiver = pd.DataFrame({"w": [3.0, 25.0, 65.0]}, index=[901, 902, 903]) + settings = dict( + predictors=["w"], + imputed_variables=["y"], + models=[OLS], + preprocessing={"w": transform}, + train_size=0.8, + random_state=11, + k_folds=2, + imputation_quantiles=[0.1, 0.5, 0.9], + hyperparameters={ + "OLS": {"row_filter": pd.Series(donor.index % 3 != 0, index=donor.index)} + }, + ) + result = autoimpute(donor, receiver, weight_col="w", **settings) + control = autoimpute( + donor.assign(weight=donor.w), receiver, weight_col="weight", **settings + ) + for q in [0.1, 0.5, 0.9]: + pd.testing.assert_frame_equal( + result.imputations["best_method"][q], control.imputations["best_method"][q] + ) + pd.testing.assert_frame_equal( + result.cv_results["OLS"]["quantile_loss"]["results"], + control.cv_results["OLS"]["quantile_loss"]["results"], + ) + sampled = donor.sample(frac=0.8, random_state=11) + selected = sampled.loc[sampled.index % 3 != 0] + np.testing.assert_allclose( + result.fitted_models["best_method"].models["y"].model.model.weights, + selected.w / selected.w.mean(), + ) + + +def test_cv_tuning_refit_preserves_raw_weight_predictor(monkeypatch): + cv_module = importlib.import_module("microimpute.evaluations.cross_validation") + original_fit = cv_module._fit_model_for_fold + seen = [] + donor = pd.DataFrame( + {"w": np.arange(1.0, 31.0), "y": np.sin(np.arange(30.0))}, + index=np.arange(30) * 7, + ) + + def capture_fit(model, model_class, data, predictors, targets, weight_col, *args): + weights = data[weight_col] if isinstance(weight_col, str) else weight_col + np.testing.assert_array_equal(weights, donor.loc[data.index, "w"]) + assert abs(data.w.mean()) < 1e-12 + seen.append(len(data)) + return original_fit( + model, model_class, data, predictors, targets, weight_col, *args + ) + + monkeypatch.setattr(cv_module, "_fit_model_for_fold", capture_fit) + # Bound the unrelated parameter search while retaining real QRF weighted fits. + monkeypatch.setattr( + QRF, + "_tune_hyperparameters", + lambda self, **kwargs: {"n_estimators": 5, "min_samples_leaf": 2}, + ) + _, params = cross_validate_model( + QRF, + donor, + ["w"], + ["y"], + weight_col="w", + preprocessing={"w": "normalize"}, + n_splits=2, + quantiles=[0.5], + tune_hyperparameters=True, + ) + assert params == {"n_estimators": 5, "min_samples_leaf": 2} + assert sorted(seen) == [15, 15, 30] + + +@pytest.mark.parametrize("model_class", [OLS, QRF]) +@pytest.mark.parametrize("value", ["a", True]) +def test_default_constant_category_has_default_frame_and_point_mass(model_class, value): + donor = pd.DataFrame({"x": np.arange(20.0), "y": [value] * 20}) + receiver = pd.DataFrame({"x": [21.0, 22.0]}, index=[101, 109]) + predictions = get_imputations( + [model_class], donor, receiver, ["x"], ["y"], quantiles=None + )[model_class.__name__] + assert set(predictions) == {0.5, "probabilities"} + pd.testing.assert_frame_equal( + predictions[0.5], pd.DataFrame({"y": [value, value]}, index=receiver.index) + ) + info = predictions["probabilities"]["y"] + np.testing.assert_array_equal(info["classes"], [value]) + np.testing.assert_array_equal(info["probabilities"], np.ones((2, 1))) + + +@pytest.mark.parametrize("model_class", [OLS, QRF, QuantReg]) +def test_default_numeric_constant_return_contract_is_unchanged(model_class): + donor = pd.DataFrame({"x": np.arange(20.0), "y": [3.0] * 20}) + receiver = pd.DataFrame({"x": [21.0, 22.0]}) + predictions = get_imputations( + [model_class], donor, receiver, ["x"], ["y"], quantiles=None + )[model_class.__name__] + assert isinstance(predictions, pd.DataFrame) + np.testing.assert_array_equal(predictions.y, [3, 3]) + + +@pytest.mark.parametrize("dtype", ["Int64", "UInt8", "Float64"]) +def test_ols_nullable_numeric_receiver_matches_native_float(dtype): + rng = np.random.default_rng(456) + donor = pd.DataFrame( + { + "children": pd.Series(np.tile(np.arange(5), 12), dtype=dtype), + "y": rng.normal(size=60) + np.tile(np.arange(5), 12) * 3, + } + ) + receiver = donor[["children"]].iloc[[1, 7, 19]].copy() + receiver.index = [101, 103, 109] + fitted = OLS().fit(donor, ["children"], ["y"]) + actual = fitted.predict(receiver, [0.1, 0.5, 0.9]) + control = ( + OLS() + .fit(donor.astype({"children": float}), ["children"], ["y"]) + .predict(receiver.astype(float), [0.1, 0.5, 0.9]) + ) + for q in actual: + pd.testing.assert_frame_equal(actual[q], control[q]) + pd.testing.assert_index_equal(actual[q].index, receiver.index) + # A single/homogeneous receiver must still get the fitted intercept. + one = fitted.predict(receiver.iloc[[0]], [0.5])[0.5] + pd.testing.assert_frame_equal(one, actual[0.5].iloc[[0]]) + + +@pytest.mark.parametrize("bad", [pd.NA, np.inf]) +def test_ols_nullable_numeric_receiver_rejects_nonfinite(bad): + donor = pd.DataFrame( + {"x": np.arange(20.0), "y": np.arange(20.0) + np.sin(np.arange(20.0))} + ) + fitted = OLS().fit(donor, ["x"], ["y"]) + with pytest.raises(ValueError, match="finite|missing|NaN"): + fitted.predict(pd.DataFrame({"x": pd.Series([bad], dtype="Float64")}), [0.5])