Fix tests that pass for the wrong reason, and export ZeroInflatedImputer

changelog.d/213.fixed.md

@@ -0,0 +1 @@
+The test suite now runs cleanly without `rpy2` or `pytorch-tabular` instead of erroring, and five tests that could pass without asserting anything now assert. `ZeroInflatedImputer` is exported from `microimpute` and `microimpute.models`. Keeps the legacy `VALID_YEARS` and `DEFAULT_MODEL_PARAMS` imports, and corrects the `Imputer.fit` weight docstring.

microimpute/__init__.py

@@ -62,7 +62,14 @@
 )
 
 # 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/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 @@
 # 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 @@
     2022,
 ]
 
+# Data configuration
 TRAIN_SIZE: float = 0.8
 TEST_SIZE: float = 0.2
 
@@ -39,40 +41,31 @@
 # 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/models/__init__.py

@@ -11,6 +11,9 @@
     - 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 @@
 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/models/imputer.py

@@ -261,7 +261,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 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

tests/test_autoimpute.py

@@ -7,24 +7,38 @@
 
 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:
     HAS_MATCHING = False
 
 # 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)

tests/test_models/test_imputers.py

@@ -719,9 +719,12 @@ def test_missing_predictors_in_test(model_class: Type[Imputer]) -> None:
     else:
         fitted = model.fit(train_data, ["x1", "x2"], ["y"])
 
-    # Should raise an error when predictor is missing
-    with pytest.raises(Exception):
-        predictions = fitted.predict(test_data, quantiles=[0.5])
+    # Should raise an error when predictor is missing. Matching surfaces this
+    # from R rather than from pandas, so its message differs.
+    with pytest.raises(
+        Exception, match=r"not in index|Missing columns|matching failed"
+    ):
+        fitted.predict(test_data, quantiles=[0.5])
 
 
 # === Reproducibility Tests ===

tests/test_models/test_qrf.py

@@ -747,7 +747,7 @@ def test_qrf_all_variables_missing() -> None:
 def test_qrf_error_handling() -> None:
     """Test error handling in QRF model."""
     # Test with empty data
-    with pytest.raises(Exception):
+    with pytest.raises(Exception, match=r"must not be None or empty"):
         model = QRF()
         model.fit(pd.DataFrame(), predictors=[], imputed_variables=[])
 
@@ -760,10 +760,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(Exception, match=r"(?i)are in the"):
+        fitted_model.predict(test_data)
 
 
 # === Internal Model Tests ===

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