Add some tests on transform() with xCEBRA

Author: CeliaBenquet

tests/test_integration_xcebra.py

@@ -1,5 +1,6 @@
 import pickle
 
+import numpy as np
 import pytest
 import torch
 
@@ -150,3 +151,20 @@ def test_synthetic_data_training(synthetic_data, device):
     assert Z2_hat.shape == Z2.shape, f"Incorrect Z2 embedding dimension: {Z2_hat.shape}"
     assert not torch.isnan(Z1_hat).any(), "NaN values in Z1 embedding"
     assert not torch.isnan(Z2_hat).any(), "NaN values in Z2 embedding"
+
+    # Test the transform
+    solver.model.split_outputs = False
+    transform_embedding = solver.transform(data.neural.to(device))
+    assert transform_embedding.shape[
+        1] == n_latents, "Incorrect embedding dimension"
+    assert not torch.isnan(transform_embedding).any(), "NaN values in embedding"
+    assert np.allclose(embedding, transform_embedding, rtol=1e-02)
+
+    # Test the transform with batching
+    batched_embedding = solver.transform(data.neural.to(device), batch_size=512)
+    assert batched_embedding.shape[
+        1] == n_latents, "Incorrect embedding dimension"
+    assert not torch.isnan(batched_embedding).any(), "NaN values in embedding"
+    assert np.allclose(embedding, batched_embedding, rtol=1e-02)
+
+    assert np.allclose(transform_embedding, batched_embedding, rtol=1e-02)