Blend out-of-fold predictions only over members that predicted each row

tabfm/src/classifier_and_regressor.py

@@ -2508,15 +2508,19 @@ def fit(self, X: Any, y: Any) -> "TabFMClassifier":
         and self.active_calibration_method_ != "none"
     ):
       oof_probs = self.predict_oof_proba(cv=self.num_folds_for_cv)
-      val_idx = getattr(self, "oof_val_indices_", None)
-      if val_idx is not None:
-        oof_probs_fit = oof_probs[:, val_idx, :]
-        y_orig_fit = y_orig[val_idx]
-        y_fit = y[val_idx]
-      else:
-        oof_probs_fit = oof_probs
-        y_orig_fit = y_orig
-        y_fit = y
+      valid_mask = getattr(self, "oof_valid_mask_", None)
+      if valid_mask is None:
+        valid_mask = np.ones(oof_probs.shape[:2], dtype=bool)
+      # Keep only rows some member predicted out of fold, and carry which
+      # members predicted each kept row so the blending below averages over
+      # those members alone. Without subsampling every member predicts every
+      # row, so this keeps all rows and the mask is all True (see
+      # predict_oof_proba).
+      kept = np.flatnonzero(valid_mask.any(axis=0))
+      oof_probs_fit = oof_probs[:, kept, :]
+      mask_fit = valid_mask[:, kept]
+      y_orig_fit = y_orig[kept]
+      y_fit = y[kept]
 
     if self.enable_nnls and oof_probs_fit is not None:
       n_classes = self.n_classes_
@@ -2548,10 +2552,15 @@ def fit(self, X: Any, y: Any) -> "TabFMClassifier":
         and self.active_calibration_method_ != "none"
         and oof_probs_fit is not None
     ):
+      # enable_nnls forbids max_num_rows (see __init__), so the weighted blend
+      # never sees subsampling. The plain mean can, and must average each row
+      # over the members that predicted it, not the zero-filled gaps (which
+      # would otherwise make the row sum to k / n_estimators).
       if self.enable_nnls:
         P = np.tensordot(self.ensemble_weights_, oof_probs_fit, axes=(0, 0))
       else:
-        P = np.mean(oof_probs_fit, axis=0)
+        counts = mask_fit.sum(axis=0)
+        P = oof_probs_fit.sum(axis=0) / counts[:, None]
       assert P.shape == (len(y_fit), self.n_classes_), (
           f"Expected calibration input shape {(len(y_fit), self.n_classes_)},"
           f" got {P.shape}"
@@ -2957,14 +2966,16 @@ def predict_oof_proba(self, cv: int = 5) -> jt.Float[Array | np.ndarray, "E N K"
       folds_to_run = folds_base
 
     outputs_oof = np.zeros((n_estimators, N, n_classes))
-
-    self.oof_val_indices_ = None
-    for fold_idx, (train_fold, val_fold) in enumerate(folds_to_run):
+    # Which rows each member actually predicted out of fold. Under row
+    # subsampling (max_num_rows) members validate on different absolute rows,
+    # so outputs_oof stays zero everywhere else; fit() uses this mask to blend
+    # each row over the members that predicted it instead of averaging in the
+    # zero-filled gaps.
+    self.oof_valid_mask_ = np.zeros((n_estimators, N), dtype=bool)
+    for train_fold, val_fold in folds_to_run:
       data_fold, val_indices_list = self.ensemble_generator_.transform_fold(
           train_fold, val_fold
       )
-      if fold_idx == 0 and len(folds_to_run) == 1:
-        self.oof_val_indices_ = val_indices_list[0]
       (
           Xs_batch,
           ys_batch,
@@ -2988,6 +2999,7 @@ def predict_oof_proba(self, cv: int = 5) -> jt.Float[Array | np.ndarray, "E N K"
             out_i, axis=-1, temperature=self.softmax_temperature
         )
         outputs_oof[i, val_indices_list[i]] = out_i
+        self.oof_valid_mask_[i, val_indices_list[i]] = True
 
     return outputs_oof
 

tabfm/src/classifier_and_regressor_test.py

@@ -796,6 +796,87 @@ def test_calibration_multiclass(self):
         probs = classifier.predict_proba(X)
         self.assertEqual(probs.shape, (150, 3))
 
+  def test_predict_oof_proba_sets_valid_mask_under_subsampling(self):
+    # Row subsampling makes each member validate on a different set of absolute
+    # rows; oof_valid_mask_ must mark exactly the rows a member predicted so
+    # fit() can blend each row over those members alone (issue #55).
+    classifier = TabFMClassifier(
+        model=self.model,
+        n_estimators=3,
+        batch_size=2,
+        max_num_rows=8,
+    )
+    X = np.random.rand(12, 3)
+    y = np.array([0, 1] * 6)
+    classifier.fit(X, y)
+
+    oof = classifier.predict_oof_proba(cv=2)
+    mask = classifier.oof_valid_mask_
+
+    self.assertEqual(mask.shape, (3, 12))
+    self.assertEqual(mask.dtype, np.bool_)
+    for i in range(3):
+      # A member's OOF row is a softmax (sums to 1) exactly where it predicted,
+      # and left at zero elsewhere.
+      predicted = oof[i].sum(axis=1) > 0
+      np.testing.assert_array_equal(predicted, mask[i])
+      np.testing.assert_allclose(oof[i, mask[i]].sum(axis=1), 1.0, atol=1e-6)
+
+  def test_calibration_blends_only_predicted_members(self):
+    # Issue #55: with per-member row subsampling the calibration input must
+    # average only the members that predicted each row, not the zero-filled
+    # gaps that would otherwise make rows sum to k / n_estimators.
+    mask = np.array([
+        [True, True, True, False, False, False],
+        [False, True, True, True, False, False],
+        [False, False, True, True, True, False],
+    ])
+    rng = np.random.RandomState(0)
+    oof = np.zeros((3, 6, 2))
+    for i in range(3):
+      for r in np.flatnonzero(mask[i]):
+        oof[i, r] = rng.dirichlet([1.0, 1.0])
+    X = rng.rand(6, 3)
+    y = np.array([0, 1, 0, 1, 0, 1])
+
+    captured = {}
+    real_fit_calibration = TabFMClassifier._fit_calibration
+
+    def spy(inner_self, P, targets):
+      captured["P"] = np.asarray(P).copy()
+      captured["y"] = np.asarray(targets).copy()
+      return real_fit_calibration(inner_self, P, targets)
+
+    classifier = TabFMClassifier(
+        model=self.model,
+        n_estimators=3,
+        binary_calibration_method="platt",
+    )
+
+    def fake_oof(cv=5):
+      classifier.oof_valid_mask_ = mask
+      return oof
+
+    patched_oof = mock.patch.object(
+        classifier, "predict_oof_proba", side_effect=fake_oof
+    )
+    patched_cal = mock.patch.object(TabFMClassifier, "_fit_calibration", spy)
+    with patched_oof, patched_cal:
+      classifier.fit(X, y)
+
+    P = captured["P"]
+    # Row 5 was predicted by nobody, so it is dropped; the rest stay in order.
+    self.assertEqual(P.shape, (5, 2))
+    np.testing.assert_array_equal(captured["y"], np.array([0, 1, 0, 1, 0]))
+    np.testing.assert_allclose(P.sum(axis=1), 1.0, atol=1e-6)
+    np.testing.assert_allclose(P[0], oof[0, 0], atol=1e-6)
+    np.testing.assert_allclose(P[1], (oof[0, 1] + oof[1, 1]) / 2, atol=1e-6)
+    np.testing.assert_allclose(
+        P[2], (oof[0, 2] + oof[1, 2] + oof[2, 2]) / 3, atol=1e-6
+    )
+    np.testing.assert_allclose(P[3], (oof[1, 3] + oof[2, 3]) / 2, atol=1e-6)
+    np.testing.assert_allclose(P[4], oof[2, 4], atol=1e-6)
+
 
 @unittest.skipUnless(HAS_JAX, "JAX is required")
 class StackingTest(absltest.TestCase):