Update GoF computation

Author: stes

cebra/integrations/sklearn/cebra.py

@@ -599,6 +599,13 @@ def num_sessions(self) -> Optional[int]:
         """
         return self.num_sessions_
 
+    @property
+    def num_negatives_(self) -> int:
+        """The number of negative examples."""
+        if self.num_negatives is None:
+            return self.batch_size
+        return self.num_negatives
+
     @property
     def state_dict_(self) -> dict:
         return self.solver_.state_dict()

cebra/integrations/sklearn/metrics.py

@@ -100,12 +100,12 @@ def infonce_loss(
     solver.to(cebra_model.device_)
     avg_loss = solver.validation(loader=loader, session_id=session_id)
     if correct_by_batchsize:
-        if cebra_model.batch_size is None:
+        if cebra_model.num_negatives_ is None:
             raise ValueError(
                 "Batch size is None, please provide a model with a batch size to correct the InfoNCE."
             )
         else:
-            avg_loss = avg_loss - np.log(cebra_model.batch_size)
+            avg_loss = avg_loss - np.log(cebra_model.num_negatives_)
     return avg_loss
 
 
@@ -211,7 +211,7 @@ def infonce_to_goodness_of_fit(
     Args:
         infonce: The InfoNCE loss, either a single value or an iterable of values.
         model: The trained CEBRA model.
-        batch_size: The batch size used to train the model.
+        batch_size: The batch size (or number of negatives, if different from the batch size) used to train the model.
         num_sessions: The number of sessions used to train the model.
 
     Returns:
@@ -228,19 +228,15 @@ def infonce_to_goodness_of_fit(
             )
         if not hasattr(model, "state_dict_"):
             raise RuntimeError("Fit the CEBRA model first.")
-        if model.batch_size is None:
+        if model.num_negatives_ is None:
             raise ValueError(
                 "Computing the goodness of fit is not yet supported for "
                 "models trained on the full dataset (batchsize = None). ")
-        batch_size = model.batch_size
+        batch_size = model.num_negatives_
         num_sessions = model.num_sessions_
         if num_sessions is None:
             num_sessions = 1
 
-        if model.batch_size is None:
-            raise ValueError(
-                "Computing the goodness of fit is not yet supported for "
-                "models trained on the full dataset (batchsize = None). ")
     else:
         if batch_size is None or num_sessions is None:
             raise ValueError(

tests/test_sklearn_metrics.py

@@ -482,14 +482,22 @@ def _fit_and_get_history(X, y):
 
 
 @pytest.mark.parametrize("seed", [42, 24, 10])
-def test_infonce_to_goodness_of_fit(seed):
+@pytest.mark.parametrize("batch_size", [100, 200])
+@pytest.mark.parametrize("num_negatives", [None, 100, 200])
+def test_infonce_to_goodness_of_fit(seed, batch_size, num_negatives):
     """Test the conversion from InfoNCE loss to goodness of fit metric."""
+    nats_to_bits = np.log2(np.e)
+
     # Test with model
     cebra_model = cebra_sklearn_cebra.CEBRA(
         model_architecture="offset10-model",
         max_iterations=5,
-        batch_size=128,
+        batch_size=batch_size,
+        num_negatives=num_negatives,
     )
+    if num_negatives is None:
+        num_negatives = batch_size
+
     generator = torch.Generator().manual_seed(seed)
     X = torch.rand(1000, 50, dtype=torch.float32, generator=generator)
     cebra_model.fit(X)
@@ -498,19 +506,22 @@ def test_infonce_to_goodness_of_fit(seed):
     gof = cebra_sklearn_metrics.infonce_to_goodness_of_fit(1.0,
                                                            model=cebra_model)
     assert isinstance(gof, float)
+    assert np.isclose(gof, (np.log(num_negatives) - 1.0) * nats_to_bits)
 
     # Test array of values
     infonce_values = np.array([1.0, 2.0, 3.0])
     gof_array = cebra_sklearn_metrics.infonce_to_goodness_of_fit(
         infonce_values, model=cebra_model)
     assert isinstance(gof_array, np.ndarray)
     assert gof_array.shape == infonce_values.shape
+    assert np.allclose(gof_array,
+                       (np.log(num_negatives) - infonce_values) * nats_to_bits)
 
     # Test with explicit batch_size and num_sessions
-    gof = cebra_sklearn_metrics.infonce_to_goodness_of_fit(1.0,
-                                                           batch_size=128,
-                                                           num_sessions=1)
+    gof = cebra_sklearn_metrics.infonce_to_goodness_of_fit(
+        1.0, batch_size=batch_size, num_sessions=1)
     assert isinstance(gof, float)
+    assert np.isclose(gof, (np.log(batch_size) - 1.0) * nats_to_bits)
 
     # Test error cases
     with pytest.raises(ValueError, match="batch_size.*should not be provided"):