Repository: Parietal-INRIA/fmralign

Issue:
Using surface data
Currently, the code takes Niimg-like objects. PairwiseAlignment also takes as input a mask (also Niimg-like object). How can we straightforwardly use surface space data with fmralign?

One idea: Make the input images numpy arrays (ntimepoints/ncontrasts x nvertices). 'clustering' could also become an array (nvertices,) This puts the onus on the user to convert any data format to numpy array before using fmralign. Ther user would also have to generate 'clustering' and mask the data before using fmralign. It improves usability without restricting to particular inputs data formats. One disadvantage of this is that passing full arrays to fmralign may consume too much memory, especially when doing many subjects in parallel. 

Discussion:

Patch:
diff --git a/fmralign/_utils.py b/fmralign/_utils.py
index 7f8e4ac..6fc8111 100644
--- a/fmralign/_utils.py
+++ b/fmralign/_utils.py
@@ -1,5 +1,4 @@
 # -*- coding: utf-8 -*-
-import os
 import warnings
 
 import nibabel as nib
@@ -8,6 +7,7 @@ from nilearn._utils.niimg_conversions import check_same_fov
 from nilearn.image import new_img_like, smooth_img
 from nilearn.masking import apply_mask_fmri, intersect_masks
 from nilearn.regions.parcellations import Parcellations
+from nilearn.surface import SurfaceImage
 
 
 class ParceledData:
@@ -219,12 +219,13 @@ def _make_parcellation(
         Parcellation of features in clusters
     """
     # check if clustering is provided
-    if isinstance(clustering, nib.nifti1.Nifti1Image) or os.path.isfile(
-        clustering
-    ):
+    if isinstance(clustering, nib.nifti1.Nifti1Image):
         check_same_fov(masker.mask_img_, clustering)
         labels = apply_mask_fmri(clustering, masker.mask_img_).astype(int)
 
+    elif isinstance(clustering, SurfaceImage):
+        labels = masker.transform(clustering)[0].astype(int)
+
     # otherwise check it's needed, if not return 1 everywhere
     elif n_pieces == 1:
         labels = np.ones(
@@ -256,9 +257,7 @@ def _make_parcellation(
             )
             err.args += (errmsg,)
             raise err
-        labels = apply_mask_fmri(
-            parcellation.labels_img_, masker.mask_img_
-        ).astype(int)
+        labels = masker.transform(parcellation.labels_img_)[0].astype(int)
 
     if verbose > 0:
         unique_labels, counts = np.unique(labels, return_counts=True)
diff --git a/fmralign/preprocessing.py b/fmralign/preprocessing.py
index d271b51..369be67 100644
--- a/fmralign/preprocessing.py
+++ b/fmralign/preprocessing.py
@@ -8,6 +8,8 @@ from joblib import Memory, Parallel, delayed
 from nibabel.nifti1 import Nifti1Image
 from nilearn._utils.masker_validation import check_embedded_masker
 from nilearn.image import concat_imgs
+from nilearn.maskers._utils import concatenate_surface_images
+from nilearn.surface import SurfaceImage
 from sklearn.base import BaseEstimator, TransformerMixin
 
 from fmralign._utils import (
@@ -118,10 +120,7 @@ class ParcellationMasker(BaseEstimator, TransformerMixin):
 
     def _fit_masker(self, imgs):
         """Fit the masker on a single or multiple images."""
-        self.masker_ = check_embedded_masker(self)
-        self.masker_.n_jobs = self.n_jobs
-
-        if isinstance(imgs, Nifti1Image):
+        if isinstance(imgs, (Nifti1Image, SurfaceImage)):
             imgs = [imgs]
         # If images are 3D, add a fourth dimension
         for i, img in enumerate(imgs):
@@ -136,32 +135,46 @@ class ParcellationMasker(BaseEstimator, TransformerMixin):
             raise NotImplementedError(
                 "fmralign does not support images of different shapes."
             )
-        if self.masker_.mask_img is None:
-            self.masker_.fit(imgs)
-        else:
-            self.masker_.fit()
-
-        if isinstance(self.clustering, Nifti1Image) or os.path.isfile(
-            self.clustering
-        ):
-            # check that clustering provided fills the mask, if not, reduce the mask
-            if 0 in self.masker_.transform(self.clustering):
-                reduced_mask = _intersect_clustering_mask(
-                    self.clustering, self.masker_.mask_img
-                )
-                self.mask = reduced_mask
-                self.masker_ = check_embedded_masker(self)
-                self.masker_.n_jobs = self.n_jobs
+
+        masker_type = (
+            "surface" if isinstance(imgs[0], SurfaceImage) else "multi_nii"
+        )
+        self.masker_ = check_embedded_masker(self, masker_type=masker_type)
+        self.masker_.n_jobs = self.n_jobs
+
+        # Fit the masker for volume data
+        if masker_type == "multi_nii":
+            if self.masker_.mask_img is None:
+                self.masker_.fit(imgs)
+            else:
                 self.masker_.fit()
-                warnings.warn(
-                    "Mask used was bigger than clustering provided. "
-                    + "Its intersection with the clustering was used instead."
-                )
+
+            if isinstance(self.clustering, Nifti1Image) or os.path.isfile(
+                self.clustering
+            ):
+                # check that clustering provided fills the mask, if not, reduce the mask
+                if 0 in self.masker_.transform(self.clustering):
+                    reduced_mask = _intersect_clustering_mask(
+                        self.clustering, self.masker_.mask_img
+                    )
+                    self.mask = reduced_mask
+                    self.masker_ = check_embedded_masker(self)
+                    self.masker_.n_jobs = self.n_jobs
+                    self.masker_.fit()
+                    warnings.warn(
+                        "Mask used was bigger than clustering provided. "
+                        + "Its intersection with the clustering was used instead."
+                    )
+        else:
+            self.masker_.fit(imgs)
 
     def _one_parcellation(self, imgs):
         """Compute one parcellation for all images."""
         if isinstance(imgs, list):
-            imgs = concat_imgs(imgs)
+            if isinstance(imgs[0], (Nifti1Image)):
+                imgs = concat_imgs(imgs)
+            else:
+                imgs = concatenate_surface_images(imgs)
         self.labels = _make_parcellation(
             imgs,
             self.clustering,
@@ -234,7 +247,7 @@ class ParcellationMasker(BaseEstimator, TransformerMixin):
             List of ParceledData objects containing the data and parcelation
             information for each image.
         """
-        if isinstance(imgs, Nifti1Image):
+        if isinstance(imgs, (Nifti1Image, SurfaceImage)):
             imgs = [imgs]
 
         parceled_data = Parallel(n_jobs=self.n_jobs)(


Test patch:
diff --git a/fmralign/tests/test_pairwise_alignment.py b/fmralign/tests/test_pairwise_alignment.py
index cd4a375..f7ed206 100644
--- a/fmralign/tests/test_pairwise_alignment.py
+++ b/fmralign/tests/test_pairwise_alignment.py
@@ -6,11 +6,13 @@ import pytest
 from nibabel import Nifti1Image
 from nilearn.image import new_img_like
 from nilearn.maskers import NiftiMasker
+from nilearn.surface import SurfaceImage
 
 from fmralign.pairwise_alignment import PairwiseAlignment
 from fmralign.tests.utils import (
     assert_algo_transform_almost_exactly,
     random_niimg,
+    surf_img,
     zero_mean_coefficient_determination,
 )
 
@@ -130,3 +132,26 @@ def test_parcellation_before_fit():
         AttributeError, match="Parcellation has not been computed yet"
     ):
         alignment.get_parcellation()
+
+
+def test_surface_alignment():
+    """Test compatibility with `SurfaceImage`"""
+    alignment = PairwiseAlignment()
+    n_pieces = 3
+    img1 = surf_img(20)
+    img2 = surf_img(20)
+    alignment = PairwiseAlignment(n_pieces=n_pieces)
+
+    # Test fitting
+    alignment.fit(img1, img2)
+
+    # Test transformation
+    img_transformed = alignment.transform(img1)
+    assert img_transformed.shape == img1.shape
+    assert isinstance(img_transformed, SurfaceImage)
+
+    # Test parcellation retrieval
+    labels, parcellation_image = alignment.get_parcellation()
+    assert isinstance(labels, np.ndarray)
+    assert len(np.unique(labels)) == n_pieces
+    assert isinstance(parcellation_image, SurfaceImage)
diff --git a/fmralign/tests/test_preprocessing.py b/fmralign/tests/test_preprocessing.py
index 0efa96f..e04c203 100644
--- a/fmralign/tests/test_preprocessing.py
+++ b/fmralign/tests/test_preprocessing.py
@@ -5,7 +5,7 @@ from nibabel.nifti1 import Nifti1Image
 
 from fmralign._utils import ParceledData
 from fmralign.preprocessing import ParcellationMasker
-from fmralign.tests.utils import random_niimg
+from fmralign.tests.utils import random_niimg, surf_img
 
 
 def test_init_default_params():
@@ -36,22 +36,22 @@ def test_fit_single_image():
     """Test that ParcellationMasker fits a single image"""
     img, _ = random_niimg((8, 7, 6))
     parcel_masker = ParcellationMasker(n_pieces=2)
-    fitted_parcel_masker = parcel_masker.fit(img)
+    parcel_masker.fit(img)
 
-    assert hasattr(fitted_parcel_masker, "masker_")
-    assert fitted_parcel_masker.labels is not None
-    assert isinstance(fitted_parcel_masker.labels, np.ndarray)
-    assert len(np.unique(fitted_parcel_masker.labels)) == 2  # n_pieces=2
+    assert hasattr(parcel_masker, "masker_")
+    assert parcel_masker.labels is not None
+    assert isinstance(parcel_masker.labels, np.ndarray)
+    assert len(np.unique(parcel_masker.labels)) == 2  # n_pieces=2
 
 
 def test_fit_multiple_images():
     """Test that ParcellationMasker fits multiple images"""
     imgs = [random_niimg((8, 7, 6))[0]] * 3
     parcel_masker = ParcellationMasker(n_pieces=2)
-    fitted_parcel_masker = parcel_masker.fit(imgs)
+    parcel_masker = parcel_masker.fit(imgs)
 
-    assert hasattr(fitted_parcel_masker, "masker_")
-    assert fitted_parcel_masker.labels is not None
+    assert hasattr(parcel_masker, "masker_")
+    assert parcel_masker.labels is not None
 
 
 def test_transform_single_image():
@@ -181,13 +181,43 @@ def test_standardization():
     assert np.abs(np.std(data_array) - 1.0) < 1e-5
 
 
+def test_one_surface_image():
+    """Test that ParcellationMasker can handle surface images"""
+    img = surf_img(20)
+    n_pieces = 2
+    n_vertices_total = img.shape[0]
+    parcel_masker = ParcellationMasker(n_pieces=n_pieces)
+    parcel_masker.fit(img)
+
+    assert hasattr(parcel_masker, "masker_")
+    assert parcel_masker.labels is not None
+    assert isinstance(parcel_masker.labels, np.ndarray)
+    assert len(np.unique(parcel_masker.labels)) == n_pieces
+    assert len(parcel_masker.labels) == n_vertices_total
+
+
+def test_multiple_surface_images():
+    """Test that ParcellationMasker can handle multiple surface images"""
+    imgs = [surf_img(20)] * 3
+    n_pieces = 2
+    n_vertices_total = imgs[0].shape[0]
+    parcel_masker = ParcellationMasker(n_pieces=n_pieces)
+    parcel_masker = parcel_masker.fit(imgs)
+
+    assert hasattr(parcel_masker, "masker_")
+    assert parcel_masker.labels is not None
+    assert isinstance(parcel_masker.labels, np.ndarray)
+    assert len(np.unique(parcel_masker.labels)) == n_pieces
+    assert len(parcel_masker.labels) == n_vertices_total
+
+
 def test_one_contrast():
     """Test that ParcellationMasker handles both 3D and\n
     4D images in the case of one contrast"""
     img1, _ = random_niimg((8, 7, 6))
     img2, _ = random_niimg((8, 7, 6, 1))
-    pmasker = ParcellationMasker()
-    pmasker.fit([img1, img2])
+    parcel_masker = ParcellationMasker()
+    parcel_masker.fit([img1, img2])
 
 
 def test_get_parcellation_img():
diff --git a/fmralign/tests/test_template_alignment.py b/fmralign/tests/test_template_alignment.py
index 7b914ec..60606ea 100644
--- a/fmralign/tests/test_template_alignment.py
+++ b/fmralign/tests/test_template_alignment.py
@@ -3,6 +3,7 @@ import pytest
 from nibabel import Nifti1Image
 from nilearn.image import concat_imgs, math_img
 from nilearn.maskers import NiftiMasker
+from nilearn.surface import SurfaceImage
 from numpy.testing import assert_array_almost_equal
 
 from fmralign._utils import ParceledData
@@ -19,6 +20,7 @@ from fmralign.tests.utils import (
     random_niimg,
     sample_parceled_data,
     sample_subjects_data,
+    surf_img,
     zero_mean_coefficient_determination,
 )
 
@@ -258,5 +260,34 @@ def test_parcellation_before_fit():
         alignment.get_parcellation()
 
 
-if __name__ == "__main__":
-    test_parcellation_retrieval()
+def test_surface_template():
+    """Test compatibility with `SurfaceImage`"""
+    n_pieces = 3
+    img1 = surf_img(20)
+    img2 = surf_img(20)
+    img3 = surf_img(20)
+    alignment = TemplateAlignment(n_pieces=n_pieces)
+
+    # Test fitting
+    alignment.fit([img1, img2, img3])
+    assert isinstance(alignment.template, SurfaceImage)
+
+    # Test transformation from new subject
+    img_transformed = alignment.transform(surf_img(20))
+    assert isinstance(img_transformed, SurfaceImage)
+    assert img_transformed.shape == surf_img(20).shape
+
+    # Test transformation on real subject
+    img_transformed = alignment.transform(img1, subject_index=0)
+    assert isinstance(img_transformed, SurfaceImage)
+    for hemi in ["left", "right"]:
+        assert np.allclose(
+            img_transformed.data.parts[hemi],
+            img1.data.parts[hemi],
+        )
+
+    # Test parcellation retrieval
+    labels, parcellation_image = alignment.get_parcellation()
+    assert isinstance(labels, np.ndarray)
+    assert len(np.unique(labels)) == n_pieces
+    assert isinstance(parcellation_image, SurfaceImage)
diff --git a/fmralign/tests/utils.py b/fmralign/tests/utils.py
index 61af947..d943ee1 100644
--- a/fmralign/tests/utils.py
+++ b/fmralign/tests/utils.py
@@ -1,6 +1,7 @@
 import nibabel
 import numpy as np
 from nilearn.maskers import NiftiMasker
+from nilearn.surface import InMemoryMesh, PolyMesh, SurfaceImage
 from numpy.random import default_rng
 from numpy.testing import assert_array_almost_equal
 
@@ -43,9 +44,7 @@ def zero_mean_coefficient_determination(
     nonzero_numerator = numerator != 0
     valid_score = nonzero_denominator & nonzero_numerator
     output_scores = np.ones([y_true.shape[1]])
-    output_scores[valid_score] = 1 - (
-        numerator[valid_score] / denominator[valid_score]
-    )
+    output_scores[valid_score] = 1 - (numerator[valid_score] / denominator[valid_score])
     output_scores[nonzero_numerator & ~nonzero_denominator] = 0
 
     if multioutput == "raw_values":
@@ -56,8 +55,7 @@ def zero_mean_coefficient_determination(
         avg_weights = None
     elif multioutput == "variance_weighted":
         avg_weights = (
-            weight
-            * (y_true - np.average(y_true, axis=0, weights=sample_weight)) ** 2
+            weight * (y_true - np.average(y_true, axis=0, weights=sample_weight)) ** 2
         ).sum(axis=0, dtype=np.float64)
         # avoid fail on constant y or one-element arrays
         if not np.any(nonzero_denominator):
@@ -134,7 +132,54 @@ def sample_parceled_data(n_pieces=1):
     return data, masker, labels
 
 
+def _make_mesh():
+    """Create a sample mesh with two parts: left and right, and total of
+    9 vertices and 10 faces.
+
+    The left part is a tetrahedron with four vertices and four faces.
+    The right part is a pyramid with five vertices and six faces.
+    """
+    left_coords = np.asarray([[0.0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]])
+    left_faces = np.asarray([[1, 0, 2], [0, 1, 3], [0, 3, 2], [1, 2, 3]])
+    right_coords = (
+        np.asarray([[0.0, 0, 0], [1, 0, 0], [1, 1, 0], [0, 1, 0], [0, 0, 1]]) + 2.0
+    )
+    right_faces = np.asarray(
+        [
+            [0, 1, 4],
+            [0, 3, 1],
+            [1, 3, 2],
+            [1, 2, 4],
+            [2, 3, 4],
+            [0, 4, 3],
+        ]
+    )
+    return PolyMesh(
+        left=InMemoryMesh(left_coords, left_faces),
+        right=InMemoryMesh(right_coords, right_faces),
+    )
+
+
+def surf_img(n_samples=1):
+    """Create a sample surface image using the sample mesh. # noqa: D202
+    This will add some random data to the vertices of the mesh.
+    The shape of the data will be (n_vertices, n_samples).
+    n_samples by default is 1.
+    """
+
+    mesh = _make_mesh()
+    data = {}
+    for i, (key, val) in enumerate(mesh.parts.items()):
+        data_shape = (val.n_vertices, n_samples)
+        data_part = (
+            np.arange(np.prod(data_shape)).reshape(data_shape[::-1]) + 1.0
+        ) * 10**i
+        data[key] = data_part.T
+    return SurfaceImage(mesh, data)
+
+
 def sample_subjects_data(n_subjects=3):
     """Sample data in one parcel for n_subjects"""
     subjects_data = [np.random.rand(10, 20) for _ in range(n_subjects)]
     return subjects_data
+