Add tests to solver

Author: CeliaBenquet

cebra/data/multi_session.py

@@ -116,6 +116,18 @@ def configure_for(self, model: "cebra.models.Model"):
         for i, session in enumerate(self.iter_sessions()):
             session.configure_for(model[i])
 
+    def configure_for(self, model: "cebra.models.Model"):
+        """Configure the dataset offset for the provided model.
+
+        Call this function before indexing the dataset. This sets the
+        :py:attr:`offset` attribute of the dataset.
+
+        Args:
+            model: The model to configure the dataset for.
+        """
+        for i, session in enumerate(self.iter_sessions()):
+            session.configure_for(model[i])
+
 
 @dataclasses.dataclass
 class MultiSessionLoader(cebra_data.Loader):

cebra/data/single_session.py

@@ -73,7 +73,7 @@ def configure_for(self, model: "cebra.models.Model"):
         """Configure the dataset offset for the provided model.
 
         Call this function before indexing the dataset. This sets the
-        `offset` attribute of the dataset.
+        :py:attr:`offset` attribute of the dataset.
 
         Args:
             model: The model to configure the dataset for.

cebra/solver/base.py

@@ -38,6 +38,7 @@
 import literate_dataclasses as dataclasses
 import numpy.typing as npt
 import numpy as np
+import numpy.typing as npt
 import torch
 import torch.nn.functional as F
 import tqdm
@@ -569,8 +570,12 @@ def _select_model(
         raise NotImplementedError
 
     @property
-    def is_fitted(self):
-        return hasattr(self, "n_features")
+    def _check_is_fitted(self):
+        #NOTE(celia): instead of hasattr(model, "n_features_"), double check this!
+        if not (hasattr(self, "history") and len(self.history) > 0):
+            raise ValueError(
+                f"This {type(self).__name__} instance is not fitted yet. Call 'fit' with "
+                "appropriate arguments before using this estimator.")
 
     @torch.no_grad()
     def transform(self,

tests/test_solver.py

@@ -59,11 +59,13 @@
      cebra.data.ContinuousMultiSessionDataLoader, "offset1-model"),
     ("demo-continuous-multisession",
      cebra.data.ContinuousMultiSessionDataLoader, "offset10-model"),
+    ("demo-discrete-multisession",
+              cebra.data.DiscreteMultiSessionDataLoader, "offset1-model"),
+    ("demo-discrete-multisession",
+     cebra.data.DiscreteMultiSessionDataLoader, "offset10-model"),
 ]:
     multi_session_tests.append((*args, cebra.solver.MultiSessionSolver))
 
-# multi_session_tests.append((*args, cebra.solver.MultiSessionAuxVariableSolver))
-
 
 def _get_loader(data, loader_initfunc):
     kwargs = dict(num_steps=5, batch_size=32)
@@ -165,6 +167,28 @@ def test_single_session(data_name, loader_initfunc, model_architecture,
 
     assert solver.num_sessions == None
     assert solver.n_features == X.shape[1]
+    
+    embedding = solver.transform(X)
+    assert isinstance(embedding, torch.Tensor)
+    assert embedding.shape == (X.shape[0], OUTPUT_DIMENSION)
+    embedding = solver.transform(torch.Tensor(X))
+    assert isinstance(embedding, torch.Tensor)
+    assert embedding.shape == (X.shape[0], OUTPUT_DIMENSION)
+    embedding = solver.transform(X, session_id=0)
+    assert isinstance(embedding, torch.Tensor)
+    assert embedding.shape == (X.shape[0], OUTPUT_DIMENSION)
+    embedding = solver.transform(X, pad_before_transform=False)
+    assert isinstance(embedding, torch.Tensor)
+    assert embedding.shape == (X.shape[0] - len(offset) + 1, OUTPUT_DIMENSION)
+
+    with pytest.raises(ValueError, match="torch.Tensor"):
+        solver.transform(X.numpy())
+    with pytest.raises(RuntimeError, match="Invalid.*session_id"):
+        embedding = solver.transform(X, session_id=2)
+
+    for param in solver.parameters():
+        assert isinstance(param, torch.Tensor)
+
 
     embedding = solver.transform(X)
     assert isinstance(embedding, torch.Tensor)
@@ -320,6 +344,183 @@ def test_multi_session(data_name, loader_initfunc, model_architecture,
     assert solver.num_sessions == 3
     assert solver.n_features == [X[i].shape[1] for i in range(len(X))]
 
+    embedding = solver.transform(X[0], session_id=0)
+    assert isinstance(embedding, torch.Tensor)
+    assert embedding.shape == (X[0].shape[0], OUTPUT_DIMENSION)
+    embedding = solver.transform(X[1], session_id=1)
+    assert isinstance(embedding, torch.Tensor)
+    assert embedding.shape == (X[1].shape[0], OUTPUT_DIMENSION)
+    embedding = solver.transform(X[0], session_id=0, pad_before_transform=False)
+    assert isinstance(embedding, torch.Tensor)
+    assert embedding.shape == (X[0].shape[0] -
+                               len(solver.model[0].get_offset()) + 1,
+                               OUTPUT_DIMENSION)
+
+    with pytest.raises(ValueError, match="torch.Tensor"):
+        embedding = solver.transform(X[0].numpy(), session_id=0)
+
+    with pytest.raises(ValueError, match="shape"):
+        embedding = solver.transform(X[1], session_id=0)
+    with pytest.raises(ValueError, match="shape"):
+        embedding = solver.transform(X[0], session_id=1)
+
+    with pytest.raises(RuntimeError, match="No.*session_id"):
+        embedding = solver.transform(X[0])
+    with pytest.raises(RuntimeError, match="single.*session"):
+        embedding = solver.transform(X)
+    with pytest.raises(RuntimeError, match="Invalid.*session_id"):
+        embedding = solver.transform(X[0], session_id=5)
+    with pytest.raises(RuntimeError, match="Invalid.*session_id"):
+        embedding = solver.transform(X[0], session_id=-1)
+
+    for param in solver.parameters(session_id=0):
+        assert isinstance(param, torch.Tensor)
+
+    with pytest.raises(RuntimeError, match="No.*session_id"):
+        for param in solver.parameters():
+            assert isinstance(param, torch.Tensor)
+
+
+@pytest.mark.parametrize(
+    "inputs, add_padding, offset, start_batch_idx, end_batch_idx, expected_output",
+    [
+        # Test case 1: No padding
+        (torch.tensor([[1, 2], [3, 4], [5, 6]]), False, cebra.data.Offset(
+            0, 1), 0, 2, torch.tensor([[1, 2], [3, 4]])),  # first batch
+        (torch.tensor([[1, 2], [3, 4], [5, 6]]), False, cebra.data.Offset(
+            0, 1), 1, 3, torch.tensor([[3, 4], [5, 6]])),  # last batch
+        (torch.tensor(
+            [[1, 2], [3, 4], [5, 6], [7, 8]]), False, cebra.data.Offset(
+                0, 1), 1, 3, torch.tensor([[3, 4], [5, 6]])),  # middle batch
+
+        # Test case 2: First batch with padding
+        (
+            torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]),
+            True,
+            cebra.data.Offset(0, 1),
+            0,
+            2,
+            torch.tensor([[1, 2, 3], [4, 5, 6]]),
+        ),
+        (
+            torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]),
+            True,
+            cebra.data.Offset(1, 1),
+            0,
+            3,
+            torch.tensor([[1, 2, 3], [1, 2, 3], [4, 5, 6], [7, 8, 9]]),
+        ),
+
+        # Test case 3: Last batch with padding
+        (
+            torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]),
+            True,
+            cebra.data.Offset(0, 1),
+            1,
+            3,
+            torch.tensor([[4, 5, 6], [7, 8, 9]]),
+        ),
+        (
+            torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12],
+                          [13, 14, 15]]),
+            True,
+            cebra.data.Offset(1, 2),
+            1,
+            3,
+            torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]]),
+        ),
+
+        # Test case 4: Middle batch with padding
+        (
+            torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]]),
+            True,
+            cebra.data.Offset(0, 1),
+            1,
+            3,
+            torch.tensor([[4, 5, 6], [7, 8, 9]]),
+        ),
+        (
+            torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]]),
+            True,
+            cebra.data.Offset(1, 1),
+            1,
+            3,
+            torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]),
+        ),
+        (
+            torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12],
+                          [13, 14, 15]]),
+            True,
+            cebra.data.Offset(0, 1),
+            2,
+            4,
+            torch.tensor([[7, 8, 9], [10, 11, 12]]),
+        ),
+        (
+            torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]]),
+            True,
+            cebra.data.Offset(0, 1),
+            0,
+            3,
+            torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]),
+        ),
+
+        # Examples that throw an error:
+
+        # Padding without offset (should raise an error)
+        (torch.tensor([[1, 2]]), True, None, 0, 2, ValueError),
+        # Negative start_batch_idx or end_batch_idx (should raise an error)
+        (torch.tensor([[1, 2]]), False, cebra.data.Offset(
+            0, 1), -1, 2, ValueError),
+        # out of bound indices because offset is too large
+        (torch.tensor([[1, 2], [3, 4]]), True, cebra.data.Offset(
+            5, 5), 1, 2, ValueError),
+        # Batch length is smaller than offset.
+        (torch.tensor([[1, 2], [3, 4]]), False, cebra.data.Offset(
+            0, 1), 0, 1, ValueError),  # first batch
+    ],
+)
+def test_get_batch(inputs, add_padding, offset, start_batch_idx, end_batch_idx,
+                   expected_output):
+    if expected_output == ValueError:
+        with pytest.raises(ValueError):
+            cebra.solver.base._get_batch(inputs, offset, start_batch_idx,
+                                         end_batch_idx, add_padding)
+    else:
+        result = cebra.solver.base._get_batch(inputs, offset, start_batch_idx,
+                                              end_batch_idx, add_padding)
+        assert torch.equal(result, expected_output)
+
+
+@pytest.mark.parametrize("data_name, loader_initfunc, solver_initfunc",
+                         multi_session_tests)
+def test_multi_session_2(data_name, loader_initfunc, solver_initfunc):
+    loader = _get_loader(data_name, loader_initfunc)
+    criterion = cebra.models.InfoNCE()
+    model = nn.ModuleList(
+        [_make_model(dataset) for dataset in loader.dataset.iter_sessions()])
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+
+    solver = solver_initfunc(model=model,
+                             criterion=criterion,
+                             optimizer=optimizer,
+                             tqdm_on=True)
+
+    batch = next(iter(loader))
+    for session_id, dataset in enumerate(loader.dataset.iter_sessions()):
+        assert batch[session_id].reference.shape == (32,
+                                                     dataset.input_dimension,
+                                                     10)
+        assert batch[session_id].index is not None
+
+    log = solver.step(batch)
+    assert isinstance(log, dict)
+
+    solver.fit(loader)
+
+    assert solver.num_sessions == 3
+    assert solver.n_features == [X[i].shape[1] for i in range(len(X))]
+
     embedding = solver.transform(X[0], session_id=0)
     assert isinstance(embedding, torch.Tensor)
     assert embedding.shape == (X[0].shape[0], OUTPUT_DIMENSION)
@@ -504,8 +705,8 @@ def create_model(model_name, input_dimension):
 
 @pytest.mark.parametrize(
     "data_name, model_name ,session_id, loader_initfunc, solver_initfunc",
-    single_session_tests_select_model + single_session_hybrid_tests_select_model
-)
+    single_session_tests_select_model +
+    single_session_hybrid_tests_select_model)
 def test_select_model_single_session(data_name, model_name, session_id,
                                      loader_initfunc, solver_initfunc):
     dataset = cebra.datasets.init(data_name)
@@ -576,7 +777,7 @@ def test_select_model_multi_session(data_name, model_name, session_id,
     "offset40-model-4x-subsample",
     "offset1-model",
     "offset10-model",
-]
+]  #NOTE(rodrigo): there is an issue with "offset4-model-2x-subsample" because it's not a convolutional model.
 batch_size_inference = [40_000, 99_990, 99_999]
 
 single_session_tests_transform = []