Median-impute NaN in calc_cross_validation_score

openavmkit/utilities/stats.py

@@ -1634,6 +1634,45 @@ def calc_cross_validation_score(
     float
         The mean cross-validated mean squared error.
     """
+    # Impute NaN with column medians before fitting.
+    # LinearRegression (sklearn) does not accept NaN natively; median imputation is a
+    # neutral choice for this variable-selection pre-pass -- LightGBM training still sees
+    # real NaN. Without this, cross_val_score raises inside individual folds, and the
+    # `nan` this function would otherwise return silently disables the caller's
+    # cross-validation refinement (`nan < best_score` is always False).
+    if isinstance(X, pd.DataFrame):
+        if X.isnull().values.any():
+            import warnings
+            warnings.warn(
+                f"calc_cross_validation_score: NaN detected in "
+                f"{list(X.columns[X.isnull().any()])}. "
+                "Imputing with column medians for the cross-validation step only.",
+                UserWarning,
+            )
+            X = X.fillna(X.median(numeric_only=True))
+    else:
+        X = np.asarray(X, dtype=float)
+        if np.isnan(X).any():
+            import warnings
+            warnings.warn(
+                "calc_cross_validation_score: NaN detected in input array. "
+                "Imputing with column medians for the cross-validation step only.",
+                UserWarning,
+            )
+            X = np.where(np.isnan(X), np.nanmedian(X, axis=0), X)
+
+    # Any column that was entirely NaN has no median to impute from; drop it rather than
+    # letting a residual NaN fail the fit.
+    if isinstance(X, pd.DataFrame):
+        all_nan = X.columns[X.isnull().all()]
+        if len(all_nan) > 0:
+            X = X.drop(columns=list(all_nan))
+    else:
+        X = np.nan_to_num(X, nan=0.0)
+
+    if isinstance(X, pd.DataFrame) and X.shape[1] == 0:
+        return float("nan")
+
     model = LinearRegression()
     # Use negative MSE and negate it to return positive MSE
     try:

tests/test_stats.py

@@ -37,4 +37,45 @@ def test_cod_bootstrap():
   print(expected)
   print("***")
 
-  assert objects_are_equal(results, expected)
+  assert objects_are_equal(results, expected)
+
+def test_cross_validation_score_with_nan():
+  """A feature column with genuine missing values must still yield a finite score.
+
+  Without imputation, sklearn raises inside individual CV folds and the mean of the
+  surviving scores is nan -- which silently disables the caller's `cv_score <
+  best_score` refinement, since `nan < x` is always False.
+  """
+
+  import pandas as pd
+  from openavmkit.utilities.stats import calc_cross_validation_score
+
+  rng = np.random.default_rng(0)
+  n = 2000
+  X = pd.DataFrame({
+    "bldg_area_finished_sqft": rng.normal(1500, 400, n),
+    "bldg_rooms_bed": rng.integers(1, 6, n).astype(float),
+    "latitude": rng.normal(39.95, 0.05, n),
+  })
+  y = pd.Series(
+    2.0 * X["bldg_area_finished_sqft"]
+    + 5000 * X["bldg_rooms_bed"]
+    + rng.normal(0, 1000, n)
+  )
+
+  score_clean = calc_cross_validation_score(X, y)
+  assert np.isfinite(score_clean)
+
+  # Genuine missing building attributes, spread across folds.
+  X_nan = X.copy()
+  X_nan.loc[rng.choice(n, 200, replace=False), "bldg_rooms_bed"] = np.nan
+  score_nan = calc_cross_validation_score(X_nan, y)
+  assert np.isfinite(score_nan), "NaN in a feature column must not produce a nan score"
+
+  # An all-NaN column has no median to impute from; it must be dropped, not propagated.
+  X_dead = X.copy()
+  X_dead["never_recorded"] = np.nan
+  assert np.isfinite(calc_cross_validation_score(X_dead, y))
+
+  # The numpy path is supported by the signature too.
+  assert np.isfinite(calc_cross_validation_score(X_nan.to_numpy(), y.to_numpy()))