PolicyEngine
diff --git a/‎changelog.d/ols-prediction-se-and-logistic-penalty.fixed.md‎
Lines changed: 1 addition & 0 deletions b/‎changelog.d/ols-prediction-se-and-logistic-penalty.fixed.md‎
Lines changed: 1 addition & 0 deletions
diff --git a/‎microimpute/models/ols.py‎
Lines changed: 61 additions & 23 deletions b/‎microimpute/models/ols.py‎
Lines changed: 61 additions & 23 deletions
diff --git a/‎tests/test_models/test_ols.py‎
Lines changed: 131 additions & 0 deletions b/‎tests/test_models/test_ols.py‎
Lines changed: 131 additions & 0 deletions
@@ -0,0 +1 @@
+Fixed three OLS imputation bugs: (#6) `_OLSModel.predict` previously used `se = sqrt(model.scale)` (residual std only) to scale normal quantiles, which under-dispersed imputations for test rows far from the training centroid; now uses `statsmodels`' `model.get_prediction(X).var_pred_mean + model.scale` to include the leverage term, producing a per-row prediction SE. The quantile is also clipped to `(1e-6, 1-1e-6)` so `q=0` / `q=1` no longer produce ±inf from `norm.ppf`. (#8) `_LogisticRegressionModel.fit` previously passed `l1_ratio` to `LogisticRegression` with the default `penalty="l2"`, which silently ignored the parameter; now sets `penalty="elasticnet"` and `solver="saga"` when a non-zero `l1_ratio` is supplied so the parameter actually takes effect. (followup) `OLSResults._predict_quantile` now returns a `pd.Series` indexed to `mean_preds.index` rather than a bare `ndarray`. When a mixed-type imputation built the output `DataFrame` by assigning a numeric column (ndarray) before a categorical column (indexed Series), pandas anchored the DataFrame to a default `RangeIndex` and the subsequent categorical assignment failed to align — silently producing all-NaN categorical columns, which later broke downstream `sklearn.log_loss`.
@@ -59,17 +59,29 @@ def fit(
                 )
                 y_encoded = y_encoded.fillna(0)  # Default to first category
 
-        # Extract relevant LR parameters from kwargs
-        # Use l1_ratio instead of penalty (deprecated in sklearn 1.8)
+        # Extract relevant LR parameters from kwargs.
+        # sklearn's LogisticRegression ignores l1_ratio unless
+        # penalty="elasticnet" (and solver="saga"). Previously we passed
+        # l1_ratio through with the default penalty="l2", so tuning
+        # l1_ratio had no effect — silent misconfiguration.
+        l1_ratio = lr_kwargs.get("l1_ratio", None)
         classifier_params = {
-            "l1_ratio": lr_kwargs.get("l1_ratio", 0),
             "C": lr_kwargs.get("C", 1.0),
             "max_iter": lr_kwargs.get("max_iter", 1000),
-            "solver": lr_kwargs.get(
-                "solver", "lbfgs" if len(self.categories) <= 2 else "saga"
-            ),
             "random_state": self.seed,
         }
+        if l1_ratio is not None and l1_ratio != 0:
+            # Caller explicitly supplied a non-zero l1_ratio: wire up the
+            # elasticnet penalty and saga solver so it actually applies.
+            classifier_params["penalty"] = "elasticnet"
+            classifier_params["l1_ratio"] = float(l1_ratio)
+            classifier_params["solver"] = lr_kwargs.get("solver", "saga")
+        else:
+            # No elasticnet requested — use the default L2 penalty and a
+            # solver that supports it.
+            classifier_params["solver"] = lr_kwargs.get(
+                "solver", "lbfgs" if len(self.categories) <= 2 else "saga"
+            )
 
         self.classifier = LogisticRegression(**classifier_params)
         fit_kwargs = {}
@@ -216,12 +228,22 @@ def _predict_variable(
                 X_test[self.predictors], return_probs=False, quantile=quantile
             )
         else:
-            # Regression for numeric targets
+            # Regression for numeric targets.
+            # Use the full prediction SE (leverage + residual) rather than
+            # just sqrt(model.scale). Previously se = sqrt(scale) used only
+            # 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])
-            mean_preds = model.predict(X_test_with_const)
-            se = np.sqrt(model.scale)
+            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
+            # variance for a new observation.
+            pred_var = np.asarray(prediction.var_pred_mean) + model.scale
+            mean_preds = np.asarray(prediction.predicted_mean)
+            se = np.sqrt(np.maximum(pred_var, 0.0))
             imputed_values = self._predict_quantile(
-                mean_preds=mean_preds,
+                mean_preds=pd.Series(mean_preds, index=X_test.index, name=variable),
                 se=se,
                 mean_quantile=quantile,
                 random_sample=random_sample,
@@ -362,37 +384,50 @@ def _predict(
     def _predict_quantile(
         self,
         mean_preds: pd.Series,
-        se: float,
+        se: Any,
         mean_quantile: float,
         random_sample: bool,
         count_samples: int = 10,
-    ) -> np.ndarray:
+    ) -> pd.Series:
         """Predict values at a specified quantile.
 
         Args:
             mean_preds: Mean predictions from the model.
-            se: Standard error of the predictions.
+            se: Standard error of the predictions. May be a scalar (legacy,
+                residual std) or a per-row array (prediction SE including
+                leverage).
             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:
-            Array of predicted values at the specified quantile.
+            Series of predicted values at the specified quantile, indexed to
+            match ``mean_preds``. Returning a Series (rather than a bare
+            ndarray) preserves the test-row index so downstream
+            ``DataFrame[col] = series`` assignments align correctly when a
+            numeric column is set before a categorical column whose
+            predictions come back indexed.
 
         Raises:
             RuntimeError: If prediction fails.
         """
+        # Clip q away from 0 and 1 to avoid ±inf from norm.ppf (and the
+        # degenerate a=0/a=inf case in the beta-alpha formula).
+        # Previously mean_quantile could be 0.0 or 1.0 with no guard.
+        q_eps = 1e-6
+        q_clipped = float(np.clip(mean_quantile, q_eps, 1.0 - q_eps))
         try:
-            if random_sample == True:
+            if random_sample:
                 self.logger.info(
-                    f"Predicting at random quantiles sampled from a beta distribution with mean quantile {mean_quantile}"
+                    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
-                a = mean_quantile / (1 - mean_quantile)
+                # Calculate alpha parameter for beta distribution (q is
+                # safely in (0,1) after clipping).
+                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)
@@ -406,12 +441,15 @@ def _predict_quantile(
                 )
                 selected_quantiles = normal_quantiles[sampled_indices]
 
-                # Adjust each mean prediction by corresponding sampled quantile times standard error
-                return mean_preds + selected_quantiles * se
+                # Adjust each mean prediction by the sampled quantile
+                # times its per-row SE (or scalar SE if se is a float).
+                values = mean_preds.values + selected_quantiles * np.asarray(se)
             else:
-                self.logger.info(f"Predicting at specified quantile {mean_quantile}")
-                specified_quantile = norm.ppf(mean_quantile)
-                return mean_preds + specified_quantile * se
+                self.logger.info(f"Predicting at specified quantile {q_clipped}")
+                specified_quantile = norm.ppf(q_clipped)
+                values = mean_preds.values + specified_quantile * np.asarray(se)
+
+            return pd.Series(values, index=mean_preds.index, name=mean_preds.name)
 
         except Exception as e:
             if isinstance(e, ValueError):
 
@@ -258,3 +258,134 @@ def test_ols_prediction_quality(diabetes_data: pd.DataFrame) -> None:
 
     # Check prediction variance is reasonable
     assert np.var(pred_values) > 0, "OLS predictions have no variance"
+
+
+def test_ols_quantile_uses_full_prediction_se() -> None:
+    """Regression test for #6: OLS quantile prediction must use the full
+    prediction SE (leverage + residual) rather than sqrt(scale). The
+    prediction SE for a new row is strictly greater than the residual
+    std; the gap grows for rows far from the training centroid.
+
+    We verify the fix by checking that (a) at a fixed quantile, the
+    prediction interval at an extreme x is WIDER than at the centroid
+    (which the old implementation could not produce — both used the
+    same residual std), and (b) at q=0.99 the quantile prediction for
+    an extreme row is larger than the residual-std-only formulation
+    would have given.
+    """
+    rng = np.random.default_rng(0)
+    n = 200
+    x = rng.normal(size=n)
+    y = 2.0 * x + rng.normal(size=n) * 0.3
+    train = pd.DataFrame({"x": x, "y": y})
+
+    model = OLS()
+    fitted = model.fit(train, ["x"], ["y"])
+
+    # One test row at the centroid, one far outside the support.
+    x_test = pd.DataFrame({"x": [0.0, 10.0]})
+    upper = fitted.predict(x_test, quantiles=[0.99])[0.99]["y"].values
+    lower = fitted.predict(x_test, quantiles=[0.01])[0.01]["y"].values
+    widths = upper - lower
+
+    # The extrapolated point must have a wider prediction interval than
+    # the centroid (leverage effect). With the old se = sqrt(scale),
+    # both rows had identical widths.
+    assert widths[1] > widths[0], (
+        "Prediction SE must grow with leverage; widths were "
+        f"{widths}, indicating residual-std-only (pre-fix) behaviour"
+    )
+
+
+def test_ols_quantile_clips_q_away_from_zero_and_one() -> None:
+    """Regression test for #6: q=0 and q=1 previously produced ±inf via
+    norm.ppf; the clipped implementation should return finite values."""
+    rng = np.random.default_rng(0)
+    n = 100
+    x = rng.normal(size=n)
+    y = 2.0 * x + rng.normal(size=n) * 0.3
+    train = pd.DataFrame({"x": x, "y": y})
+
+    model = OLS()
+    fitted = model.fit(train, ["x"], ["y"])
+
+    x_test = pd.DataFrame({"x": [0.0]})
+    preds_0 = fitted.predict(x_test, quantiles=[0.0])[0.0]["y"].values
+    preds_1 = fitted.predict(x_test, quantiles=[1.0])[1.0]["y"].values
+
+    assert np.all(np.isfinite(preds_0)), "q=0 produced non-finite predictions"
+    assert np.all(np.isfinite(preds_1)), "q=1 produced non-finite predictions"
+
+
+def test_ols_mixed_targets_preserve_test_index() -> None:
+    """Regression test: when OLS imputes a DataFrame containing both a
+    numeric target (returned via the OLS path) and a categorical/boolean
+    target (returned via the logistic-regression path), the predictions
+    assembled into the output DataFrame must align by the X_test index.
+
+    The bug: the numeric path previously returned a bare ndarray. When
+    that ndarray was assigned as the first column of a fresh empty
+    DataFrame, the DataFrame took on a default RangeIndex (0..N). The
+    subsequent categorical prediction (a pd.Series with the real
+    X_test index, e.g. 160..199) then failed to align on assignment
+    and the whole column came back as NaN — which later produced
+    ``Input contains NaN`` in sklearn's log_loss.
+    """
+    rng = np.random.default_rng(42)
+    n = 200
+    df = pd.DataFrame(
+        {
+            "num_pred1": rng.normal(size=n),
+            "num_pred2": rng.normal(size=n) * 2 + 1,
+            "num_target": rng.normal(size=n) * 3,
+            "binary_target": rng.choice([0, 1], size=n),
+        }
+    )
+    # Split so the test slice has a non-zero-based index.
+    train_data = df.iloc[:160].copy()
+    test_data = df.iloc[160:].copy()
+
+    model = OLS()
+    fitted = model.fit(
+        train_data,
+        predictors=["num_pred1", "num_pred2"],
+        imputed_variables=["num_target", "binary_target"],
+    )
+    predictions = fitted.predict(test_data, quantiles=[0.5])
+    out = predictions[0.5]
+
+    assert not out["num_target"].isna().any(), "num_target should not contain NaN"
+    assert not out["binary_target"].isna().any(), (
+        "binary_target should not be NaN; the output DataFrame index must "
+        "align with the X_test index so the categorical Series assignment "
+        "lines up with the numeric column"
+    )
+    # Index must match the test slice, not a fresh RangeIndex.
+    assert list(out.index) == list(test_data.index)
+
+
+def test_logistic_l1_ratio_activates_elasticnet() -> None:
+    """Regression test for #8: passing l1_ratio must activate the
+    elasticnet penalty (and saga solver). Previously l1_ratio was
+    passed through with the default L2 penalty and was silently ignored.
+    """
+    from microimpute.models.ols import _LogisticRegressionModel
+    import logging
+
+    rng = np.random.default_rng(0)
+    n = 100
+    X = pd.DataFrame(rng.normal(size=(n, 3)), columns=["a", "b", "c"])
+    y = pd.Series((X["a"] + rng.normal(size=n) > 0).astype(int), name="y")
+
+    model = _LogisticRegressionModel(seed=0, logger=logging.getLogger("test"))
+    model.fit(X, y, var_type="boolean", l1_ratio=0.5)
+
+    # When l1_ratio=0.5, penalty must be "elasticnet" and solver "saga"
+    # so l1_ratio actually has an effect.
+    assert model.classifier.penalty == "elasticnet", (
+        f"Expected penalty='elasticnet' with l1_ratio=0.5, got "
+        f"penalty={model.classifier.penalty!r} (l1_ratio silently ignored)"
+    )
+    assert model.classifier.solver == "saga", (
+        f"Expected solver='saga' for elasticnet, got {model.classifier.solver!r}"
+    )
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	+Fixed three OLS imputation bugs: (#6) `_OLSModel.predict` previously used `se = sqrt(model.scale)` (residual std only) to scale normal quantiles, which under-dispersed imputations for test rows far from the training centroid; now uses `statsmodels`' `model.get_prediction(X).var_pred_mean + model.scale` to include the leverage term, producing a per-row prediction SE. The quantile is also clipped to `(1e-6, 1-1e-6)` so `q=0` / `q=1` no longer produce ±inf from `norm.ppf`. (#8) `_LogisticRegressionModel.fit` previously passed `l1_ratio` to `LogisticRegression` with the default `penalty="l2"`, which silently ignored the parameter; now sets `penalty="elasticnet"` and `solver="saga"` when a non-zero `l1_ratio` is supplied so the parameter actually takes effect. (followup) `OLSResults._predict_quantile` now returns a `pd.Series` indexed to `mean_preds.index` rather than a bare `ndarray`. When a mixed-type imputation built the output `DataFrame` by assigning a numeric column (ndarray) before a categorical column (indexed Series), pandas anchored the DataFrame to a default `RangeIndex` and the subsequent categorical assignment failed to align — silently producing all-NaN categorical columns, which later broke downstream `sklearn.log_loss`.