microimpute 3.1.2__tar.gz → 3.1.4__tar.gz

{microimpute-3.1.2 → microimpute-3.1.4}/PKG-INFO

@@ -1,8 +1,17 @@
 Metadata-Version: 2.4
 Name: microimpute
-Version: 3.1.2
+Version: 3.1.4
 Summary: Benchmarking imputation methods for microdata
-Author-email: María Juaristi <juaristi@uni.minerva.edu>, Nikhil Woodruff <nikhil.woodruff@outlook.com>
+Author-email: Vahid Ahmadi <vahid@policyengine.org>, Max Ghenis <max@policyengine.org>, María Juaristi <juaristi@uni.minerva.edu>, Nikhil Woodruff <nikhil.woodruff@outlook.com>
+Project-URL: Repository, https://github.com/PolicyEngine/microimpute
+Project-URL: Documentation, https://policyengine.github.io/microimpute
+Project-URL: Issues, https://github.com/PolicyEngine/microimpute/issues
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Topic :: Scientific/Engineering :: Information Analysis
 Requires-Python: <3.15,>=3.12
 Description-Content-Type: text/markdown
 Requires-Dist: numpy<3.0.0,>=2.0.0

{microimpute-3.1.2 → microimpute-3.1.4}/microimpute/__init__.py

@@ -62,7 +62,14 @@ from microimpute.evaluations.predictor_analysis import (
 )
 
 # Import main models and utilities
-from microimpute.models import OLS, QRF, Imputer, ImputerResults, QuantReg
+from microimpute.models import (
+    OLS,
+    QRF,
+    Imputer,
+    ImputerResults,
+    QuantReg,
+    ZeroInflatedImputer,
+)
 
 # Import data handling functions
 from microimpute.utils.data import preprocess_data, unnormalize_predictions

{microimpute-3.1.2 → microimpute-3.1.4}/microimpute/config.py

@@ -1,8 +1,8 @@
 """
 Configuration module for MicroImpute.
 
-This module centralizes all constants and configuration parameters used across
-the package.
+Shared validation, analysis and plotting settings, plus compatibility constants
+for existing callers. Model implementations define their own runtime defaults.
 """
 
 from typing import Any, Dict, List
@@ -14,7 +14,8 @@ from pydantic import ConfigDict
 # arbitrary types like pd.DataFrame
 VALIDATE_CONFIG = ConfigDict(arbitrary_types_allowed=True)
 
-# Data configuration
+# Historical SCF years retained for imports in existing notebooks and callers.
+# This compatibility list does not restrict datasets accepted by the imputers.
 VALID_YEARS: List[int] = [
     1989,
     1992,
@@ -30,6 +31,7 @@ VALID_YEARS: List[int] = [
     2022,
 ]
 
+# Data configuration
 TRAIN_SIZE: float = 0.8
 TEST_SIZE: float = 0.2
 
@@ -39,40 +41,31 @@ QUANTILES: List[float] = [round(q, 2) for q in np.arange(0.05, 1.00, 0.05)]
 # Random state for reproducibility
 RANDOM_STATE: int = 42
 
-# Model parameters (passed via **kwargs to fit() or as __init__ params)
+# Historical parameter mapping retained for import compatibility. This mapping
+# does not configure the learners; their implementations own runtime defaults.
 DEFAULT_MODEL_PARAMS: Dict[str, Dict[str, Any]] = {
     "qrf": {
-        # RandomForestQuantileRegressor parameters
         "n_estimators": 100,
         "max_depth": None,
         "min_samples_split": 2,
         "min_samples_leaf": 1,
         "max_features": 1.0,
     },
-    "quantreg": {
-        # statsmodels QuantReg uses default parameters
-    },
+    "quantreg": {},
     "ols": {
-        # statsmodels OLS uses default parameters
-        # LogisticRegression params for categorical targets:
         "l1_ratio": 0,
         "C": 1.0,
         "max_iter": 1000,
     },
-    "matching": {
-        # StatMatch NND hotdeck default parameters
-    },
+    "matching": {},
     "mdn": {
-        # Backbone network parameters
         "layers": "128-64-32",
         "activation": "ReLU",
         "dropout": 0.0,
         "use_batch_norm": False,
-        # MDN head parameters
         "num_gaussian": 5,
         "softmax_temperature": 1.0,
         "n_samples": 100,
-        # Training parameters
         "learning_rate": 1e-3,
         "max_epochs": 100,
         "early_stopping_patience": 10,

{microimpute-3.1.2 → microimpute-3.1.4}/microimpute/models/__init__.py

@@ -11,6 +11,9 @@ Available models:
     - Matching: statistical matching/hot-deck imputation (optional, requires rpy2)
     - MDN: Mixture Density Network for probabilistic imputation
         (optional, requires pytorch-tabular)
+    - ZeroInflatedImputer: wrapper composing a model for the probability of a
+        zero with a model for the positive part, for variables such as asset
+        holdings where a large share of the population is at zero
 
 Base classes:
     - Imputer: abstract base class for all imputation models
@@ -34,3 +37,4 @@ except ImportError:
 from microimpute.models.ols import OLS
 from microimpute.models.qrf import QRF
 from microimpute.models.quantreg import QuantReg
+from microimpute.models.zero_inflated import ZeroInflatedImputer

{microimpute-3.1.2 → microimpute-3.1.4}/microimpute/models/imputer.py

@@ -261,7 +261,7 @@ class Imputer(ABC):
             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 of the column or column array/series containing sampling weights. When provided, the weights are passed to the underlying learner's own weighted-fit interface as sample weights. QRF, OLS and Matching support this; QuantReg and MDN raise `NotImplementedError` rather than silently returning an unweighted 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

{microimpute-3.1.2 → microimpute-3.1.4}/microimpute/models/matching.py

@@ -74,6 +74,7 @@ class MatchingResults(ImputerResults):
         self.categorical_targets = categorical_targets or {}
         self.boolean_targets = boolean_targets or {}
         self.dummy_processor = dummy_processor
+        self.n_failed_records = 0
 
     @validate_call(config=VALIDATE_CONFIG)
     def _predict(
@@ -93,11 +94,24 @@ class MatchingResults(ImputerResults):
             Dictionary mapping quantiles to imputed values.
             If return_probs=True, includes 'probabilities' key with one-hot encodings.
 
+        Side effects:
+            Sets ``self.n_failed_records`` to the number of recipient records
+            that could not be matched and are NaN in the result, and mirrors it
+            on ``result.attrs["n_failed_records"]``. It is reset to 0 on entry,
+            so it always describes the most recent call. Matching runs
+            single-threaded (``autoimpute`` forces ``n_jobs=1`` when a Matching
+            model is present), so concurrent calls on one fitted object would
+            race on it.
+
         Raises:
             ValueError: If model is not properly set up or
                 input data is invalid.
             RuntimeError: If matching or prediction fails.
         """
+        # Reset before any work: a prediction that raises must not leave the
+        # previous successful call's count readable as if it described this one.
+        self.n_failed_records = 0
+
         try:
             self.logger.info(f"Performing matching for {len(X_test)} recipient records")
 
@@ -325,6 +339,18 @@ class MatchingResults(ImputerResults):
             )
             raise RuntimeError("Failed to process matching results") from convert_error
 
+        # Both single-call and chunked predictions replace the previous count.
+        # Only missing target values represent unmatched output records.
+        self.n_failed_records = int(
+            fused0[self.imputed_variables].isna().any(axis=1).sum()
+        )
+        if self.n_failed_records:
+            self.logger.warning(
+                f"{self.n_failed_records} of {len(fused0)} records "
+                f"({self.n_failed_records / len(fused0):.1%}) could not be "
+                "matched and are NaN in the result."
+            )
+
         # Create output dictionary with results
         imputations: Dict[float, pd.DataFrame] = {}
         prob_results = {} if return_probs else None
@@ -363,6 +389,12 @@ class MatchingResults(ImputerResults):
                 if return_probs and prob_results:
                     imputations["probabilities"] = prob_results
 
+                # Mirror the unmatched count onto each frame, so a caller can
+                # see it without reaching into the fitted model.
+                for frame in imputations.values():
+                    if isinstance(frame, pd.DataFrame):
+                        frame.attrs["n_failed_records"] = self.n_failed_records
+
                 return imputations
             else:
                 # If no quantiles specified, use a default one
@@ -388,6 +420,7 @@ class MatchingResults(ImputerResults):
                             prob_results[variable] = prob_df
 
                 imputations[q_default] = imputed_df
+                imputed_df.attrs["n_failed_records"] = self.n_failed_records
 
                 # Add probabilities to results if requested
                 if return_probs and prob_results:
@@ -572,7 +605,13 @@ class Matching(Imputer):
             f"Tuning Matching hyperparameters with {n_cv_folds}-fold CV and {n_trials} trials"
         )
 
+        # Keeps the most recent underlying failure so an all-pruned study can
+        # report why, rather than only that nothing succeeded.
+        last_trial_error: Optional[BaseException] = None
+
         def objective(trial: optuna.Trial) -> float:
+            nonlocal last_trial_error
+
             params = {
                 "dist_fun": trial.suggest_categorical(
                     "dist_fun",
@@ -639,11 +678,19 @@ class Matching(Imputer):
                                 )
                                 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."
+                                )
+                                last_trial_error = e
+                                raise optuna.TrialPruned() from e
 
                         # Combine chunk results
                         y_pred = np.concatenate(y_pred_chunks)
@@ -660,11 +707,15 @@ class Matching(Imputer):
                             )
                             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."
+                            )
+                            last_trial_error = e
+                            raise optuna.TrialPruned() from e
 
                     # Use appropriate metric based on variable type
                     metric = variable_metrics[var]
@@ -709,6 +760,16 @@ class Matching(Imputer):
 
         study.optimize(objective, n_trials=n_trials)
 
+        if not any(
+            trial.state == optuna.trial.TrialState.COMPLETE for trial in study.trials
+        ):
+            raise ValueError(
+                "No matching hyperparameter trial succeeded. Last error: "
+                f"{last_trial_error}"
+                if last_trial_error
+                else "No matching hyperparameter trial succeeded"
+            )
+
         best_value = study.best_value
         self.logger.info(
             f"Matching - Lowest average normalized quantile loss ({n_cv_folds}-fold CV): {best_value}"

{microimpute-3.1.2 → microimpute-3.1.4}/microimpute/models/qrf.py

@@ -10,7 +10,7 @@ from pydantic import validate_call
 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 @@ class QRF(Imputer):
         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,19 @@ class QRF(Imputer):
             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)
+        if seed is not None and (
+            isinstance(seed, bool)
+            or not isinstance(seed, (int, np.integer))
+            or not 0 <= seed < 2**32
+        ):
+            raise ValueError(
+                f"seed must be an integer from 0 to 2**32 - 1, or None, got {seed!r}"
+            )
+        super().__init__(log_level=log_level, seed=seed)
         self.models = {}
         self.log_level = log_level
         self.memory_efficient = memory_efficient
@@ -695,20 +707,45 @@ class QRF(Imputer):
 
         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 is shared with target-specific subsampling and wraps within
+        sklearn's uint32 seed range.
+        """
+        if self.seed is None:
+            return None
+        try:
+            variable_offset = (self.imputed_variables or []).index(variable)
+        except ValueError as error:
+            # Falling back to offset 0 would hand this variable the same draws
+            # as the first target, which is the comonotonicity this method
+            # exists to prevent - and it would do so silently.
+            raise ValueError(
+                f"Cannot derive a seed for {variable!r}: it is not among the "
+                f"imputed variables {list(self.imputed_variables or [])}."
+            ) from error
+        return (int(self.seed) + variable_offset) % 2**32
+
     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,
@@ -800,12 +837,7 @@ class QRF(Imputer):
             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
             )
@@ -1465,7 +1497,9 @@ class QRF(Imputer):
                         y_val = X_val_fold[var]
 
                         # Create and fit QRF model with trial parameters
-                        model = _QRFModel(seed=self.seed, logger=self.logger)
+                        model = _QRFModel(
+                            seed=self._seed_for_variable(var), logger=self.logger
+                        )
                         model.fit(
                             X_train_augmented[encoded_predictors],
                             X_train_fold[var],
@@ -1623,7 +1657,7 @@ class QRF(Imputer):
 
                         # Create and fit RFC model with trial parameters
                         model = _RandomForestClassifierModel(
-                            seed=self.seed, logger=self.logger
+                            seed=self._seed_for_variable(var), logger=self.logger
                         )
 
                         # Determine variable type and fit appropriately

{microimpute-3.1.2 → microimpute-3.1.4}/microimpute.egg-info/PKG-INFO

@@ -1,8 +1,17 @@
 Metadata-Version: 2.4
 Name: microimpute
-Version: 3.1.2
+Version: 3.1.4
 Summary: Benchmarking imputation methods for microdata
-Author-email: María Juaristi <juaristi@uni.minerva.edu>, Nikhil Woodruff <nikhil.woodruff@outlook.com>
+Author-email: Vahid Ahmadi <vahid@policyengine.org>, Max Ghenis <max@policyengine.org>, María Juaristi <juaristi@uni.minerva.edu>, Nikhil Woodruff <nikhil.woodruff@outlook.com>
+Project-URL: Repository, https://github.com/PolicyEngine/microimpute
+Project-URL: Documentation, https://policyengine.github.io/microimpute
+Project-URL: Issues, https://github.com/PolicyEngine/microimpute/issues
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Topic :: Scientific/Engineering :: Information Analysis
 Requires-Python: <3.15,>=3.12
 Description-Content-Type: text/markdown
 Requires-Dist: numpy<3.0.0,>=2.0.0

{microimpute-3.1.2 → microimpute-3.1.4}/microimpute.egg-info/SOURCES.txt

@@ -2,6 +2,7 @@ README.md
 pyproject.toml
 microimpute/__init__.py
 microimpute/config.py
+microimpute/py.typed
 microimpute.egg-info/PKG-INFO
 microimpute.egg-info/SOURCES.txt
 microimpute.egg-info/dependency_links.txt
@@ -38,6 +39,7 @@ tests/test_dashboard_formatter.py
 tests/test_data_preprocessing.py
 tests/test_metrics.py
 tests/test_predictor_analysis.py
+tests/test_public_api.py
 tests/test_quantile_comparison.py
 tests/test_smoke_qrf.py
 tests/test_type_handling.py

{microimpute-3.1.2 → microimpute-3.1.4}/pyproject.toml

@@ -7,13 +7,23 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "microimpute"
-version = "3.1.2"
+version = "3.1.4"
 description = "Benchmarking imputation methods for microdata"
 readme = "README.md"
 authors = [
+    { name = "Vahid Ahmadi", email = "vahid@policyengine.org" },
+    { name = "Max Ghenis", email = "max@policyengine.org" },
     { name = "María Juaristi", email = "juaristi@uni.minerva.edu" },
     { name = "Nikhil Woodruff", email = "nikhil.woodruff@outlook.com" }
 ]
+classifiers = [
+    "Intended Audience :: Science/Research",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: 3.14",
+    "Topic :: Scientific/Engineering :: Information Analysis",
+]
 requires-python = ">=3.12,<3.15"
 dependencies = [
     "numpy>=2.0.0,<3.0.0",
@@ -62,12 +72,17 @@ images = [
     "kaleido>=0.2.1,<0.3.0",  # For exporting plots as PNG/JPG
 ]
 
-[tool.setuptools]
-packages = ["microimpute"]
-include-package-data = true
+[project.urls]
+Repository = "https://github.com/PolicyEngine/microimpute"
+Documentation = "https://policyengine.github.io/microimpute"
+Issues = "https://github.com/PolicyEngine/microimpute/issues"
+
+[tool.setuptools.packages.find]
+include = ["microimpute*"]
+exclude = ["tests*", "docs*", "examples*", "experiments*"]
 
 [tool.setuptools.package-data]
-"microimpute" = ["**/*"]
+"microimpute" = ["py.typed"]
 
 [tool.towncrier]
 package = "microimpute"

{microimpute-3.1.2 → microimpute-3.1.4}/tests/test_autoimpute.py

@@ -7,11 +7,10 @@ from sklearn.datasets import load_diabetes
 
 from microimpute.comparisons.autoimpute import autoimpute, AutoImputeResult
 from microimpute.visualizations import *
-from microimpute.models import QRF, QuantReg, OLS
 
 # Check if Matching is available
 try:
-    from microimpute.models import Matching
+    from microimpute.models import Matching  # noqa: F401  (import is the probe)
 
     HAS_MATCHING = True
 except ImportError:
@@ -19,12 +18,27 @@ except ImportError:
 
 # Check if MDN is available
 try:
-    from microimpute.models import MDN
+    from microimpute.models import MDN  # noqa: F401  (import is the probe)
 
     HAS_MDN = True
 except ImportError:
     HAS_MDN = False
 
+
+def available_models():
+    """The models autoimpute should use in this environment.
+
+    Matching needs rpy2 and R's StatMatch, and MDN needs pytorch-tabular.
+    Naming a model that did not import raises NameError before the call under
+    test runs, which is how a missing optional dependency turned into eight
+    test errors rather than a skip.
+    """
+    # autoimpute's own default is already dependency-aware, building
+    # [QRF, OLS, QuantReg] plus Matching and MDN when they import. Passing None
+    # uses it, so the installed set is exercised whatever is present.
+    return None
+
+
 # === Fixtures ===
 
 
@@ -87,7 +101,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=available_models(),
         hyperparameters={
             "QRF": {"n_estimators": 50},
             "Matching": {"constrained": True},
@@ -137,7 +151,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=available_models(),
         impute_all=True,  # Return results for all models
         log_level="WARNING",
     )
@@ -202,7 +216,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=available_models(),
         hyperparameters=hyperparameters,
         log_level="WARNING",
     )
@@ -224,7 +238,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=available_models(),
         log_level="WARNING",
     )
 
@@ -246,7 +260,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=available_models(),
         log_level="WARNING",
     )
 
@@ -267,7 +281,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=available_models(),
         log_level="WARNING",
     )
 
@@ -326,7 +340,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=available_models(),
         log_level="WARNING",
     )
 
@@ -372,13 +386,13 @@ def test_autoimpute_missing_predictors() -> None:
         }
     )
 
-    with pytest.raises(Exception):
+    with pytest.raises(Exception, match=r"Missing columns in receiver data"):
         autoimpute(
             donor_data=donor,
             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=available_models(),
             log_level="WARNING",
         )
 
@@ -392,7 +406,9 @@ def test_autoimpute_invalid_model_specification() -> None:
     receiver = pd.DataFrame({"x": np.random.randn(10)})
 
     # Invalid model type
-    with pytest.raises(Exception):
+    with pytest.raises(
+        Exception, match=r"(?s)validation error.*Input should be a type"
+    ):
         autoimpute(
             donor_data=donor,
             receiver_data=receiver,
@@ -461,7 +477,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=available_models(),
         log_level="WARNING",
     )
 
@@ -470,7 +486,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=available_models(),
         log_level="WARNING",
     )
 
@@ -480,5 +496,6 @@ def test_autoimpute_consistency(simple_data: tuple) -> None:
         if model_name in results2.cv_results:
             loss1 = results1.cv_results[model_name]["quantile_loss"]["mean_test"]
             loss2 = results2.cv_results[model_name]["quantile_loss"]["mean_test"]
-            if not np.isnan(loss1) and not np.isnan(loss2):
-                np.testing.assert_allclose(loss1, loss2, rtol=0.10)
+            assert not np.isnan(loss1), f"{model_name} produced a NaN loss"
+            assert not np.isnan(loss2), f"{model_name} produced a NaN loss"
+            np.testing.assert_allclose(loss1, loss2, rtol=0.10)

microimpute-3.1.4/tests/test_public_api.py

@@ -0,0 +1,22 @@
+"""Compatibility checks for public imports used by published examples."""
+
+
+def test_legacy_scf_years_remain_importable() -> None:
+    """Existing SCF examples retain their historical supported-year list."""
+    from microimpute.config import VALID_YEARS
+
+    assert VALID_YEARS == list(range(1989, 2023, 3))
+
+
+def test_legacy_model_parameters_remain_importable() -> None:
+    """The mapping stays importable for callers that still read it.
+
+    Asserting its full contents would freeze values nothing in the package
+    reads: `DEFAULT_MODEL_PARAMS` has no remaining callers inside microimpute,
+    and the real defaults live in each model. Check only what a downstream
+    caller relies on - that the import works and the expected keys are there.
+    """
+    from microimpute.config import DEFAULT_MODEL_PARAMS
+
+    assert set(DEFAULT_MODEL_PARAMS) == {"qrf", "quantreg", "ols", "matching", "mdn"}
+    assert all(isinstance(v, dict) for v in DEFAULT_MODEL_PARAMS.values())