diff --git a/tabfm/src/classifier_and_regressor.py b/tabfm/src/classifier_and_regressor.py index 4a0ebee..6923b6e 100644 --- a/tabfm/src/classifier_and_regressor.py +++ b/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 diff --git a/tabfm/src/classifier_and_regressor_test.py b/tabfm/src/classifier_and_regressor_test.py index 546fbee..bdedca6 100644 --- a/tabfm/src/classifier_and_regressor_test.py +++ b/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):