fix: calls to lovo_split

Author: oesteban

src/nifreeze/estimator.py

@@ -141,13 +141,13 @@ def estimate(
                         pbar.set_description_str(
                             f"Pass {i_iter + 1}/{n_iter} | Fit and predict b-index <{i}>"
                         )
-                        data_train, data_test = lovo_split(data, i, with_b0=True)
-                        grad_str = f"{i}, {data_test[1][:3]}, b={int(data_test[1][3])}"
+                        data_train, data_test = lovo_split(data, i)
+                        grad_str = f"{i}, {data_test[-1][:3]}, b={int(data_test[-1][3])}"
                         pbar.set_description_str(f"[{grad_str}], {n_jobs} jobs")
 
                         if not single_model:  # A true LOGO estimator
                             if hasattr(data, "gradients"):
-                                kwargs["gtab"] = data_train[1]
+                                kwargs["gtab"] = data_train[-1]
                             # Factory creates the appropriate model and pipes arguments
                             dwmodel = ModelFactory.init(
                                 model=model,
@@ -162,7 +162,7 @@ def estimate(
                             )
 
                         # generate a synthetic dw volume for the test gradient
-                        predicted = dwmodel.predict(data_test[1])
+                        predicted = dwmodel.predict(data_test[-1])
 
                         # prepare data for running ANTs
                         fixed, moving = _prepare_registration_data(
@@ -180,7 +180,7 @@ def estimate(
                             data.motion_affines,
                             data.affine,
                             data.dataobj.shape[:3],
-                            data_test[1][3],
+                            data_test[-1][3],
                             i_iter,
                             i,
                             ptmp_dir,

test/test_model.py

@@ -154,19 +154,19 @@ def test_two_initialisations(datadir):
 
     # Direct initialisation
     model1 = model.AverageDWIModel(
-        gtab=data_train[1],
+        gtab=data_train[-1],
         S0=dmri_dataset.bzero,
         th_low=100,
         th_high=1000,
         bias=False,
         stat="mean",
     )
-    model1.fit(data_train[0], gtab=data_train[1])
-    predicted1 = model1.predict(data_test[1])
+    model1.fit(data_train[0], gtab=data_train[-1])
+    predicted1 = model1.predict(data_test[-1])
 
     # Initialisation via ModelFactory
     model2 = model.ModelFactory.init(
-        gtab=data_train[1],
+        gtab=data_train[-1],
         model="avgdwi",
         S0=dmri_dataset.bzero,
         th_low=100,
@@ -176,9 +176,9 @@ def test_two_initialisations(datadir):
     )
 
     with pytest.raises(ModelNotFittedError):
-        model2.predict(data_test[1])
+        model2.predict(data_test[-1])
 
-    model2.fit(data_train[0], gtab=data_train[1])
-    predicted2 = model2.predict(data_test[1])
+    model2.fit(data_train[0], gtab=data_train[-1])
+    predicted2 = model2.predict(data_test[-1])
 
     assert np.all(predicted1 == predicted2)

test/test_splitting.py

@@ -37,6 +37,7 @@ def test_lovo_split(datadir):
 
     Returns:
     None
+
     """
     data = DWI.from_filename(datadir / "dwi.h5")
 
@@ -52,11 +53,11 @@ def test_lovo_split(datadir):
     data.gradients[..., index] = 1
 
     # Apply the lovo_split function at the specified index
-    (train_data, train_gradients), (test_data, test_gradients) = lovo_split(data, index)
+    train_data, test_data = lovo_split(data, index)
 
     # Check if the test data contains only 1s
     # and the train data contains only 0s after the split
-    assert np.all(test_data == 1)
-    assert np.all(train_data == 0)
-    assert np.all(test_gradients == 1)
-    assert np.all(train_gradients == 0)
+    assert np.all(test_data[0] == 1)
+    assert np.all(train_data[0] == 0)
+    assert np.all(test_data[-1] == 1)
+    assert np.all(train_data[-1] == 0)