From 3042dd7c5f2a1ac7eb6b5ecd60196c8b8a420d17 Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Thu, 7 Dec 2023 15:34:24 +0100
Subject: [PATCH 01/69] add: Feilong's Hyperalignment method in
 TemplateAlignment

---
 fmralign/template_alignment.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/fmralign/template_alignment.py b/fmralign/template_alignment.py
index 08a6ba9..c0121d3 100644
--- a/fmralign/template_alignment.py
+++ b/fmralign/template_alignment.py
@@ -13,6 +13,7 @@
 from sklearn.base import BaseEstimator, TransformerMixin
 
 from fmralign.pairwise_alignment import PairwiseAlignment
+from hyperalignment.hyperalignment import HyperAlignment
 
 
 def _rescaled_euclidean_mean(imgs, masker, scale_average=False):
@@ -376,6 +377,15 @@ def fit(self, imgs):
             Length : n_samples
 
         """
+        if self.alignment_method == "hyperalignment":
+            self.model = HyperAlignment(
+                n_jobs=self.n_jobs,
+            )
+
+            return self.model.fit(
+                imgs, self.masker_, self.masker_.mask_img_, verbose=self.verbose
+            )
+
         # Check if the input is a list, if list of lists, concatenate each subjects
         # data into one unique image.
         if not isinstance(imgs, (list, np.ndarray)) or len(imgs) < 2:
@@ -441,6 +451,9 @@ def transform(self, imgs, train_index, test_index):
             Each Niimg has the same length as the list test_index
 
         """
+        if self.alignment_method == "hyperalignment":
+            return self.model.transform(imgs, verbose=self.verbose)
+
         if not isinstance(imgs, (list, np.ndarray)):
             raise ValueError(
                 "The method TemplateAlignment.transform() need a list input. "

From b538a935267e2eceb34924f8c159f05a96f89561 Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Thu, 21 Dec 2023 12:04:14 +0100
Subject: [PATCH 02/69] adding pcha and slha

---
 fmralign/template_alignment.py | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/fmralign/template_alignment.py b/fmralign/template_alignment.py
index c0121d3..95f0473 100644
--- a/fmralign/template_alignment.py
+++ b/fmralign/template_alignment.py
@@ -266,7 +266,7 @@ def __init__(
         alignment_method: string
             Algorithm used to perform alignment between X_i and Y_i :
             * either 'identity', 'scaled_orthogonal', 'optimal_transport',
-            'ridge_cv', 'permutation', 'diagonal'
+            'ridge_cv', 'permutation', 'diagonal', 'pcha', 'slha'
             * or an instance of one of alignment classes (imported from
             functional_alignment.alignment_methods)
         n_pieces: int, optional (default = 1)
@@ -377,11 +377,22 @@ def fit(self, imgs):
             Length : n_samples
 
         """
-        if self.alignment_method == "hyperalignment":
+
+        # Alignment method is hyperalignment (too different from other methods)
+        if self.alignment_method == "slha":
             self.model = HyperAlignment(
+                method="searchlight",
                 n_jobs=self.n_jobs,
             )
+            return self.model.fit(
+                imgs, self.masker_, self.masker_.mask_img_, verbose=self.verbose
+            )
 
+        elif self.alignment_method == "pcha":
+            self.model = HyperAlignment(
+                method="parcels",
+                n_jobs=self.n_jobs,
+            )
             return self.model.fit(
                 imgs, self.masker_, self.masker_.mask_img_, verbose=self.verbose
             )

From baa899553c6bf441a316e9170fb1f4f4856c628a Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Tue, 2 Jan 2024 14:11:45 +0100
Subject: [PATCH 03/69] edit pyproject.toml

---
 .gitignore     | 2 --
 pyproject.toml | 3 ++-
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/.gitignore b/.gitignore
index 5c4fd70..e479904 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,5 +1,3 @@
-.DS_Store
-
 # hatchling
 fmralign/_version.py
 fmarlign/__pycache__
diff --git a/pyproject.toml b/pyproject.toml
index ec15836..289c829 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -25,7 +25,8 @@ dependencies = [
     "scipy",
     "nibabel",
     "nilearn@git+https://github.com/nilearn/nilearn#egg=1b11edf5789813d20995e7744bc58f3a2b6bc4ea",
-    "POT"
+    "POT",
+    "hyperalignment"
 ]
 dynamic = ["version"]
 

From c0995b2d849591f5370ecb871487b4a80df94d9e Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Wed, 3 Jan 2024 11:57:09 +0100
Subject: [PATCH 04/69] rework hyperalignment incorporation

---
 .flake8                                      |   2 +-
 fmralign/alignment_methods.py                | 289 +++++++++++
 fmralign/hyperalignment/__init__.py          |   0
 fmralign/hyperalignment/hyperalignment.py    | 106 ++++
 fmralign/hyperalignment/linalg.py            | 105 ++++
 fmralign/hyperalignment/local_template.py    | 256 +++++++++
 fmralign/hyperalignment/model.py             | 296 +++++++++++
 fmralign/hyperalignment/procrustes.py        |  74 +++
 fmralign/hyperalignment/regions.py           | 517 +++++++++++++++++++
 fmralign/hyperalignment/regions_alignment.py | 212 ++++++++
 fmralign/template_alignment.py               |  19 -
 11 files changed, 1856 insertions(+), 20 deletions(-)
 create mode 100644 fmralign/hyperalignment/__init__.py
 create mode 100644 fmralign/hyperalignment/hyperalignment.py
 create mode 100644 fmralign/hyperalignment/linalg.py
 create mode 100644 fmralign/hyperalignment/local_template.py
 create mode 100644 fmralign/hyperalignment/model.py
 create mode 100644 fmralign/hyperalignment/procrustes.py
 create mode 100644 fmralign/hyperalignment/regions.py
 create mode 100644 fmralign/hyperalignment/regions_alignment.py

diff --git a/.flake8 b/.flake8
index c4e987c..5c06765 100644
--- a/.flake8
+++ b/.flake8
@@ -9,7 +9,7 @@ per-file-ignores =
 	examples/*/*: D103, D205, D301, D400
     # - docstrings rules that should not be applied to doc
 	doc/*: D100, D103, F401
-ignore = D105, D107, E402, W503, W504, W605, BLK100
+ignore = D105, D107, E402, W503, W504, W605, BLK100, E501
 # for compatibility with black
 # https://black.readthedocs.io/en/stable/guides/using_black_with_other_tools.html#flake8
 extend-ignore = E203
\ No newline at end of file
diff --git a/fmralign/alignment_methods.py b/fmralign/alignment_methods.py
index ca0ad6e..38022bc 100644
--- a/fmralign/alignment_methods.py
+++ b/fmralign/alignment_methods.py
@@ -13,6 +13,9 @@
 from sklearn.base import BaseEstimator, TransformerMixin
 from sklearn.linear_model import RidgeCV
 from sklearn.metrics.pairwise import pairwise_distances
+import os
+from .hyperalignment.regions_alignment import RegionAlignment
+from .hyperalignment.linalg import safe_svd, svd_pca
 
 
 def scaled_procrustes(X, Y, scaling=False, primal=None):
@@ -465,3 +468,289 @@ def fit(self, X, Y):
     def transform(self, X):
         """Transform X using optimal coupling computed during fit."""
         return X.dot(self.R)
+
+
+class Hyperalignment(Alignment):
+    def __init__(
+        self,
+        tmpl_kind="pca",
+        decomp_method=None,
+        latent_dim=None,
+        alignment_method="searchlight",
+        id=None,
+        cache=True,
+        n_jobs=-1,
+    ):
+        """
+        Initialize the IndividualizedNeuralTuning object.
+
+        Parameters:
+        --------
+
+        - tmpl_kind (str): The type of template to use for alignment. Default is "pca".
+        - decomp_method (str): The decomposition method to use. Default is None.
+        - latent_dim (int): The number of latent dimensions to use in the shared stimulus information matrix. Default is None.
+        - n_jobs (int): The number of parallel jobs to run. Default is -1.
+
+        Returns:
+        --------
+        None
+        """
+
+        self.n_s = None
+        self.n_t = None
+        self.n_v = None
+        self.labels = None
+        self.alphas = None
+        self.alignment_method = alignment_method
+        if alignment_method == "parcel":
+            self.parcels = None
+
+        elif (
+            alignment_method == "searchlight"
+            or alignment_method == "ensemble_searchlight"
+        ):
+            self.searchlights = None
+            self.distances = None
+            self.radius = None
+
+        self.path = None
+        self.tuning_data = []
+        self.denoised_signal = []
+        self.decomp_method = decomp_method
+        self.tmpl_kind = tmpl_kind
+        self.latent_dim = latent_dim
+        self.n_jobs = n_jobs
+        self.cache = cache
+
+        if cache:
+            if id is None:
+                self.id = "default"
+            else:
+                self.id = id
+
+            path = os.path.join(os.getcwd(), f"cache/int/{self.id}")
+            # Check if cache folder exists
+            if not os.path.exists(path):
+                os.makedirs(path)
+
+            self.path = path
+
+    def fit(
+        self,
+        X_train,
+        searchlights=None,
+        parcels=None,
+        dists=None,
+        radius: int = 20,
+        verbose=True,
+    ):
+        """
+        Fits the IndividualizedNeuralTuning model to the training data.
+
+        Parameters:
+        --------
+
+        - X_train (array-like):
+            The training data of shape (n_subjects, n_samples, n_voxels).
+        - searchlights (array-like):
+            The searchlight indices for each subject, of shape (n_s, n_searchlights).
+        - parcels (array-like):
+            The parcel indices for each subject, of shape (n_s, n_parcels) (if not using searchlights)
+        - dists (array-like):
+            The distances of vertices to the center of their searchlight, of shape (n_searchlights, n_vertices_sl)
+        - radius (int, optional):
+            The radius of the searchlight sphere, in milimeters. Defaults to 20.
+        - verbose (bool, optional):
+            Whether to print progress information. Defaults to True.
+        - id (str, optional):
+            An identifier for caching purposes. Defaults to None.
+
+        Returns:
+        --------
+
+        - self (IndividualizedNeuralTuning):
+            The fitted model.
+        """
+
+        X_train_ = np.array(X_train, copy=True, dtype=np.float32)
+
+        self.n_s, self.n_t, self.n_v = X_train_.shape
+
+        self.tuning_data = np.empty(self.n_s, dtype=np.float32)
+        self.denoised_signal = np.empty(self.n_s, dtype=np.float32)
+
+        if searchlights is None:
+            self.regions = parcels
+        else:
+            self.regions = searchlights
+            self.distances = dists
+            self.radius = radius
+
+        # check for cached data
+        try:
+            self.denoised_signal = np.load(self.path + "/train_data_denoised.npy")
+            if verbose:
+                print("Loaded denoised data from cache")
+
+        except:  # noqa: E722
+            denoiser = RegionAlignment(
+                alignment_method=self.alignment_method,
+                n_jobs=self.n_jobs,
+                verbose=verbose,
+                path=self.path,
+            )
+            self.denoised_signal = denoiser.fit_transform(
+                X_train_,
+                regions=self.regions,
+                dists=dists,
+                radius=radius,
+            )
+            # Clear memory of the SearchlightAlignment object
+            denoiser = None
+
+        iterm = range(self.n_s)
+
+        # Stimulus matrix computation
+        if self.decomp_method is None:
+            full_signal = np.concatenate(self.denoised_signal, axis=1)
+            self.shared_response = stimulus_estimator(
+                full_signal, self.n_t, self.n_s, self.latent_dim
+            )
+
+            self.tuning_data = Parallel(n_jobs=self.n_jobs)(
+                delayed(tuning_estimator)(
+                    self.shared_response,
+                    self.denoised_signal[i],
+                    latent_dim=self.latent_dim,
+                )
+                for i in iterm
+            )
+
+        return self
+
+    def transform(self, X_test_, verbose=False):
+        """
+        Transforms the input test data using the hyperalignment model.
+
+        Args:
+            X_test_ (list of arrays):
+                The input test data.
+            verbose (bool, optional):
+                Whether to print verbose output. Defaults to False.
+            id (int, optional):
+                Identifier for the transformation. Defaults to None.
+
+        Returns:
+            numpy.ndarray: The transformed test data.
+        """
+
+        full_signal = np.concatenate(X_test_, axis=1, dtype=np.float32)
+
+        if verbose:
+            print("Predict : Computing stimulus matrix...")
+
+        if self.decomp_method is None:
+            S = stimulus_estimator(full_signal, self.n_t, self.n_s, self.latent_dim)
+
+        if verbose:
+            print("Predict : stimulus matrix shape: ", S.shape)
+
+        reconstructed_signal = [
+            reconstruct_signal(S, T_est) for T_est in self.tuning_data
+        ]
+        return np.array(reconstructed_signal, dtype=np.float32)
+
+    def get_shared_stimulus(self):
+        """
+        Returns the shared stimulus used for individualized neural tuning.
+
+        Returns:
+            The shared stimulus of shape (n_t, latent_dim) or (n_t, n_t).
+        """
+        return self.shared_response
+
+    def get_tuning_matrices(self):
+        """
+        Returns the tuning matrices as a NumPy array.
+
+        Returns:
+            numpy.ndarray: The tuning matrices of shape (n_s, latent_dim, n_v) or (n_s, n_t, n_v).
+        """
+        return np.array(self.tuning_data)
+
+    def clean_cache(self, id):
+        """
+        Removes the cache file associated with the given ID.
+
+        Args:
+            id (int): The ID of the cache file to be removed.
+        """
+        try:
+            os.remove("cache")
+        except:  # noqa: E722
+            print("No cache to remove")
+
+
+#######################################################################################
+# Computing decomposition
+
+
+def tuning_estimator(shared_response, target, latent_dim=None):
+    """
+    Estimate the tuning weights for individualized neural tuning.
+
+    Parameters:
+    --------
+        - shared_response (array-like):
+            The shared response matrix.
+        - target (array-like):
+            The target matrix.
+        - latent_dim (int, optional):
+            The number of latent dimensions. Defaults to None.
+
+    Returns:
+    --------
+        array-like: The estimated tuning weights.
+
+    """
+    if latent_dim is None:
+        return np.linalg.inv(shared_response).dot(target)
+    return np.linalg.pinv(shared_response).dot(target).astype(np.float32)
+
+
+def stimulus_estimator(full_signal, n_t, n_s, latent_dim=None):
+    """
+    Estimates the stimulus response using the given parameters.
+
+    Args:
+        full_signal (np.ndarray): The full signal data.
+        n_t (int): The number of time points.
+        n_s (int): The number of stimuli.
+        latent_dim (int, optional): The number of latent dimensions. Defaults to None.
+
+    Returns:
+        np.ndarray: The estimated shared response.
+    """
+    if latent_dim is not None and latent_dim < n_t:
+        U = svd_pca(full_signal)
+        U = U[:, :latent_dim]
+    else:
+        U, _, _ = safe_svd(full_signal)
+
+    shared_response = np.sqrt(n_s) * U
+    return shared_response.astype(np.float32)
+
+
+def reconstruct_signal(shared_response, individual_tuning):
+    """
+    Reconstructs the signal using the shared response and individual tuning.
+
+    Args:
+        shared_response (numpy.ndarray): The shared response of shape (n_t, n_t) or (n_t, latent_dim).
+        individual_tuning (numpy.ndarray): The individual tuning of shape (latent_dim, n_v) or (n_t, n_v).
+
+    Returns:
+        numpy.ndarray: The reconstructed signal of shape (n_t, n_v) (same shape as the original signal)
+    """
+    return (shared_response @ individual_tuning).astype(np.float32)
diff --git a/fmralign/hyperalignment/__init__.py b/fmralign/hyperalignment/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/fmralign/hyperalignment/hyperalignment.py b/fmralign/hyperalignment/hyperalignment.py
new file mode 100644
index 0000000..c5a16cd
--- /dev/null
+++ b/fmralign/hyperalignment/hyperalignment.py
@@ -0,0 +1,106 @@
+"""
+A wrapper for the IndividualTuningModel class to be used in fmralign (taking Nifti1 images as input).
+"""
+from .model import INT
+from .regions import compute_searchlights
+from nilearn.maskers import NiftiMasker
+from nibabel import Nifti1Image
+import numpy as np
+
+
+class HyperAlignment(INT):
+    def __init__(self, method="searchlight", n_jobs=-1):
+        """
+        Initialize the Hyperalignment object.
+
+        Parameters:
+        --------
+        - method (str): The method used for hyperalignment. Can be either "searchlight" or "parcellation". Default is "searchlight".
+        - n_jobs (int): The number of parallel jobs to run. Default is -1, which uses all available processors.
+
+        Returns:
+        --------
+        None
+        """
+        super().__init__(n_jobs=n_jobs)
+        self.mask_img = None
+        self.masker = None
+        self.preds = []
+        self.method = method
+
+    def fit(
+        self,
+        imgs,
+        masker: NiftiMasker = None,
+        mask_img: Nifti1Image = None,
+        y=None,
+        verbose=0,
+    ):
+        """
+        Fit the model to the data.
+        Can either take as entry a masking image or a masker object.
+        This information will be kept for transforming the data.
+
+        Parameters
+        ----------
+        imgs : list of Nifti1Image or np.ndarray
+            The images to be fit.
+        masker : NiftiMasker
+            The masker to be used to transform the images into the common space.
+        mask_img : Nifti1Image
+            The mask to be used to transform the images into the common space.
+        y : None
+            Not used.
+        verbose : int
+            The verbosity level.
+        """
+
+        if mask_img is None:
+            mask_img = masker.mask_img
+        self.mask_img = mask_img
+        self.masker = masker
+
+        if self.method == "searchlight":
+            data, searchlights, dists = compute_searchlights(
+                niimg=imgs[0],
+                mask_img=mask_img,
+            )
+
+        elif self.method == "parcellation":
+            raise NotImplementedError("Parcellation method not implemented yet.")
+
+        if isinstance(imgs, np.ndarray):
+            X = imgs
+        else:
+            X = data
+
+        super().fit(X, searchlights, dists, verbose=verbose)
+        return self
+
+    def transform(self, imgs, y=None, verbose=0):
+        """
+        Transform the data into the common space and return the Nifti1Image objects.
+
+        Parameters
+        ----------
+        imgs : list of Nifti1Image
+            The images to be transformed into the common space.
+        y : None
+            Not used.
+        verbose : int
+            The verbosity level.
+
+        Returns
+        -------
+        preds : list of Nifti1Image
+            The images in the common space.
+        """
+        if self.masker is None:
+            self.masker = NiftiMasker(mask_img=self.mask_img)
+
+        X = self.masker.fit_transform(imgs)
+        Y = super().transform(X, verbose=verbose)
+        for y in Y:
+            self.preds.append(self.masker.inverse_transform(y))
+
+        return self.preds
diff --git a/fmralign/hyperalignment/linalg.py b/fmralign/hyperalignment/linalg.py
new file mode 100644
index 0000000..48b3bd6
--- /dev/null
+++ b/fmralign/hyperalignment/linalg.py
@@ -0,0 +1,105 @@
+"""
+The linear algebra module. This module contains functions that are often
+used in hyperalignment algorithms. Specifically, the robustness of
+singular value decomposition (SVD) is enhanced in ``safe_svd`` to avoid
+occasional crashes when the operation is performed many times (e.g., in a
+searchlight algorithm), and ``svd_pca`` performs PCA based on
+``safe_svd``.
+"""
+
+import numpy as np
+from scipy.linalg import svd, LinAlgError
+
+__all__ = ["safe_svd", "svd_pca", "ridge"]
+
+
+def safe_svd(X, remove_mean=True):
+    """
+    Singular value decomposition without occasional LinAlgError crashes.
+
+    The default ``lapack_driver`` of ``scipy.linalg.svd`` is ``'gesdd'``,
+    which occassionaly crashes even if the input matrix is not singular.
+    This function automatically handles the ``LinAlgError`` when it's
+    raised and switches to the ``'gesvd'`` driver in this case.
+
+    The input matrix ``X`` is factorized as ``U @ np.diag(s) @ Vt``.
+
+    Parameters
+    ----------
+    X : ndarray of shape (M, N)
+        The matrix to be decomposed in NumPy array format.
+    remove_mean : bool, default=True
+        Whether to subtract the mean of each column before the actual SVD
+        (True) or not (False). Setting `remove_mean=True` is helpful when
+        the SVD is used to perform PCA.
+
+    Returns
+    -------
+    U : ndarray of shape (M, K)
+        Unitary matrix.
+    s : ndarray of shape (K,)
+        The singular values.
+    Vt : ndarray of shape (K, N)
+        Unitary matrix.
+    """
+    if remove_mean:
+        X = X - X.mean(axis=0, keepdims=True)
+    try:
+        U, s, Vt = svd(X, full_matrices=False)
+    except LinAlgError:
+        U, s, Vt = svd(X, full_matrices=False, lapack_driver="gesvd")
+    return U, s, Vt
+
+
+def svd_pca(X, remove_mean=True):
+    """
+    Principal component analysis (PCA) based on SVD.
+
+    This function performs a rotation and returns the transformed data in
+    PC space. Therefore, its behavior is similar to the ``fit_transform``
+    method of ``sklearn.decomposition.PCA``.
+
+    It does not throw away any PCs, and therefore there is no
+    dimensionality reduction in the PC space. However, the number of PCs
+    might be less than the number of features in ``X``, depending on the
+    rank of ``X``.
+
+    Parameters
+    ----------
+    X : ndarray of shape (M, N)
+        The data matrix to be transformed into PC space.
+    remove_mean : bool, default=True
+        Whether to subtract the mean of each column before the SVD (True)
+        or not (False). This parameter should be set to True unless the
+        columns already have zero mean.
+
+    Returns
+    -------
+    X_new : ndarray of shape (M, K)
+        The transformed data matrix in PC space.
+    """
+    U, s, Vt = safe_svd(X, remove_mean=remove_mean)
+    X_new = U * s[np.newaxis]
+    return X_new
+
+
+def ridge(X, Y, alpha):
+    """Solve ridge regression problem for matrix target using SVD.
+    Parameters
+    ----------
+    X : ndarray
+        The data matrix of shape (n_samples, n_features).
+    Y : ndarray of shape (n_samples, n_targets)
+        The target matrix.
+    alpha : float
+        The regularization parameter.
+    Returns
+    -------
+    betas : ndarray of shape (n_features, n_targets)
+        The solution to the ridge regression problem.
+    """
+    U, s, Vt = safe_svd(X, remove_mean=False)
+    d = s / (alpha + s**2)
+    d_UT_Y = d[:, np.newaxis] * (U.T @ Y)
+    betas = Vt.T @ d_UT_Y
+    return betas
diff --git a/fmralign/hyperalignment/local_template.py b/fmralign/hyperalignment/local_template.py
new file mode 100644
index 0000000..ec51ea4
--- /dev/null
+++ b/fmralign/hyperalignment/local_template.py
@@ -0,0 +1,256 @@
+""" Local template computation functions. Those functions are part of the warp hyperalignment
+introducted by Feilong Ma et al. 2023. The functions are adapted from the original code and
+adapted for more general regionations approches.
+Authors: Feilong Ma (Haxby lab, Dartmouth College), Denis Fouchard (MIND, INRIA Saclay).
+"""
+
+
+import numpy as np
+from scipy.stats import zscore
+from sklearn.decomposition import PCA
+from sklearn.utils.extmath import randomized_svd
+from tqdm import tqdm
+
+from .linalg import safe_svd
+from .procrustes import procrustes
+
+
+def PCA_decomposition(X, max_npc=None, flavor="sklearn", adjust_ns=False, demean=True):
+    """Decompose concatenated data matrices using PCA/SVD.
+
+    Parameters
+    ----------
+    dss : ndarray of shape (ns, nt, nv)
+    max_npc : integer or None
+    flavor : {'sklearn', 'svd'}
+    adjust_ns : bool
+        Whether to adjust the variance of the output so that it doesn't increase with the number of subjects.
+    demean : bool
+        Whether to remove the mean of the columns prior to SVD.
+
+    Returns
+    -------
+    XX : ndarray of shape (nt, npc)
+    cc : ndarray of shape (npc, ns, nv)
+    """
+    ns, nt, nv = X.shape
+    X = X.transpose(1, 0, 2).reshape(nt, ns * nv).astype(np.float32)
+    if max_npc is not None:
+        # max_npc = min(max_npc, min(X.shape[0], X.shape[1]))
+        pass
+    if flavor == "sklearn":
+        try:
+            if demean:
+                pca = PCA(n_components=max_npc, random_state=0)
+                XX = pca.fit_transform(X)
+                cc = pca.components_.reshape(-1, ns, nv)
+                if adjust_ns:
+                    XX /= np.sqrt(ns)
+                return XX.astype(np.float32), cc
+            else:
+                U, s, Vt = randomized_svd(
+                    X,
+                    (max_npc if max_npc is not None else min(X.shape)),
+                    random_state=0,
+                )
+                if adjust_ns:
+                    XX = U[:, :max_npc] * (s[np.newaxis, :max_npc] / np.sqrt(ns))
+                else:
+                    XX = U[:, :max_npc] * (s[np.newaxis, :max_npc])
+                cc = Vt[:max_npc].reshape(-1, ns, nv)
+                return XX.astype(np.float32), cc
+        except:  # noqa: E722
+            return PCA_decomposition(
+                X,
+                max_npc=max_npc,
+                flavor="svd",
+                adjust_ns=adjust_ns,
+                demean=demean,
+            )
+    elif flavor == "svd":
+        U, s, Vt = safe_svd(X)
+        if adjust_ns:
+            XX = U[:, :max_npc] * (s[np.newaxis, :max_npc] / np.sqrt(ns))
+        else:
+            XX = U[:, :max_npc] * (s[np.newaxis, :max_npc])
+        cc = Vt[:max_npc].reshape(-1, ns, nv)
+        return XX.astype(np.float32), cc
+    else:
+        raise NotImplementedError
+
+
+def compute_PCA_template(X, sl=None, max_npc=None, flavor="sklearn", demean=False):
+    """
+    Compute the PCA template from the input data.
+
+    Parameters:
+    -----------
+
+    - X (ndarray):
+        The input data array of shape (n_samples, n_features, n_timepoints).
+    - sl (slice, optional):
+        The slice of timepoints to consider. Defaults to None.
+    - max_npc (int, optional):
+        The maximum number of principal components to keep. Defaults to None.
+    - flavor (str, optional):
+        The flavor of PCA algorithm to use. Defaults to "sklearn".
+    - demean (bool, optional):
+        Whether to demean the data before performing PCA. Defaults to False.
+
+    Returns:
+    --------
+
+    ndarray:
+        The PCA template array of shape (n_samples, n_features, n_components).
+    """
+    if sl is not None:
+        dss = X[:, :, sl]
+    else:
+        dss = X
+    max_npc = min(dss.shape[1], dss.shape[2])
+    XX, cc = PCA_decomposition(
+        dss, max_npc=max_npc, flavor=flavor, adjust_ns=True, demean=demean
+    )
+    return XX.astype(np.float32)
+
+
+def compute_procrustes_template(
+    X,
+    region=None,
+    reflection=True,
+    scaling=False,
+    zscore_common=True,
+    level2_iter=1,
+    dss2=None,
+    debug=False,
+):
+    """
+    Compute the Procrustes template for a given set of data.
+
+    Args:
+        X (ndarray): The input data array of shape (n_samples, n_features, n_regions).
+        region (int or None, optional): The index of the region to consider. If None, all regions are considered. Defaults to None.
+        reflection (bool, optional): Whether to allow reflection in the Procrustes alignment. Defaults to True.
+        scaling (bool, optional): Whether to allow scaling in the Procrustes alignment. Defaults to False.
+        zscore_common (bool, optional): Whether to z-score the aligned data to have zero mean and unit variance. Defaults to True.
+        level2_iter (int, optional): The number of iterations for the level 2 alignment. Defaults to 1.
+        dss2 (ndarray or None, optional): The second set of input data array of shape (n_samples, n_features, n_regions). Only used for level 2 alignment. Defaults to None.
+        debug (bool, optional): Whether to display progress bars during alignment. Defaults to False.
+
+    Returns:
+        ndarray: The computed Procrustes template.
+
+    """
+    if region is not None:
+        X = X[:, :, region]
+    common_space = np.copy(X[0])
+    aligned_dss = [X[0]]
+    if debug:
+        iter = tqdm(X[1:])
+        iter.set_description("Computing procrustes alignment (level 1)...")
+    else:
+        iter = X[1:]
+    for ds in iter:
+        T = procrustes(ds, common_space, reflection=reflection, scaling=scaling)
+        aligned_ds = ds.dot(T)
+        if zscore_common:
+            aligned_ds = np.nan_to_num(zscore(aligned_ds, axis=0))
+        aligned_dss.append(aligned_ds)
+        common_space = (common_space + aligned_ds) * 0.5
+        if zscore_common:
+            common_space = np.nan_to_num(zscore(common_space, axis=0))
+
+    aligned_dss2 = []
+
+    if debug:
+        iter2 = tqdm(range(level2_iter))
+        iter2.set_description("Computing procrustes alignment (level 2)...")
+    else:
+        iter2 = range(level2_iter)
+
+    for level2 in iter2:
+        common_space = np.zeros_like(X[0])
+        for ds in aligned_dss:
+            common_space += ds
+        for i, ds in enumerate(X):
+            reference = (common_space - aligned_dss[i]) / float(len(X) - 1)
+            if zscore_common:
+                reference = np.nan_to_num(zscore(reference, axis=0))
+            T = procrustes(ds, reference, reflection=reflection, scaling=scaling)
+            if level2 == level2_iter - 1 and dss2 is not None:
+                aligned_dss2.append(dss2[i].dot(T))
+            aligned_dss[i] = ds.dot(T)
+
+    common_space = np.sum(aligned_dss, axis=0)
+    if zscore_common:
+        common_space = np.nan_to_num(zscore(common_space, axis=0))
+    else:
+        common_space /= float(len(X))
+    if dss2 is not None:
+        common_space2 = np.zeros_like(dss2[0])
+        for ds in aligned_dss2:
+            common_space2 += ds
+        if zscore_common:
+            common_space2 = np.nan_to_num(zscore(common_space2, axis=0))
+        else:
+            common_space2 /= float(len(X))
+        return common_space, common_space2
+
+    return common_space
+
+
+def compute_template(
+    X,
+    region,
+    kind="searchlight_pca",
+    max_npc=None,
+    common_topography=True,
+    demean=True,
+):
+    """
+    Compute a template from a set of datasets.
+
+    ----------
+    Parameters:
+
+       - dss (ndarray): The input datasets.
+       - region (ndarray or None): The region indices for searchlight or region-based template computation.
+       - sl (ndarray or None): The searchlight indices for searchlight-based template computation.
+       - region (int or None, optional): The index of the region to consider. If None, all regions are considered (or searchlights). Defaults to None.
+       - kind (str): The type of template computation algorithm to use. Can be "pca", "pcav1", "pcav2", or "cls".
+       - max_npc (int or None): The maximum number of principal components to use for PCA-based template computation.
+       - common_topography (bool): Whether to enforce common topography across datasets.
+       - demean (bool): Whether to demean the datasets before template computation.
+
+    ----------
+    Returns:
+        tmpl : The computed template on all parcels (or searchlights).
+    """
+    mapping = {
+        "searchlight_pca": compute_PCA_template,
+        "parcels_pca": compute_PCA_template,
+        "cls": compute_procrustes_template,
+    }
+
+    if kind == "procrustes":
+        tmpl = compute_procrustes_template(
+            X=X,
+            region=region,
+            reflection=True,
+            scaling=False,
+            zscore_common=True,
+        )
+    elif kind in mapping:
+        tmpl = mapping[kind](X, sl=region, max_npc=max_npc, demean=demean)
+    else:
+        raise ValueError("Unknown template kind")
+
+    if common_topography:
+        if region is not None:
+            dss_ = X[:, :, region]
+        else:
+            dss_ = np.copy(X)
+        ns, nt, nv = dss_.shape
+        T = procrustes(np.tile(tmpl, (ns, 1)), dss_.reshape(ns * nt, nv))
+        tmpl = tmpl @ T
+    return tmpl.astype(np.float32)
diff --git a/fmralign/hyperalignment/model.py b/fmralign/hyperalignment/model.py
new file mode 100644
index 0000000..15fe165
--- /dev/null
+++ b/fmralign/hyperalignment/model.py
@@ -0,0 +1,296 @@
+import numpy as np
+from tqdm import tqdm
+from sklearn.base import BaseEstimator, TransformerMixin
+from .regions_alignment import RegionAlignment
+from .linalg import safe_svd, svd_pca
+from joblib import Parallel, delayed
+import os
+
+
+class INT(BaseEstimator, TransformerMixin):
+    def __init__(
+        self,
+        tmpl_kind="pca",
+        decomp_method=None,
+        latent_dim=None,
+        alignment_method="searchlight",
+        id=None,
+        cache=True,
+        n_jobs=-1,
+    ):
+        """
+        Initialize the IndividualizedNeuralTuning object.
+
+        Parameters:
+        --------
+
+        - tmpl_kind (str): The type of template to use for alignment. Default is "pca".
+        - decomp_method (str): The decomposition method to use. Default is None.
+        - latent_dim (int): The number of latent dimensions to use in the shared stimulus information matrix. Default is None.
+        - n_jobs (int): The number of parallel jobs to run. Default is -1.
+
+        Returns:
+        --------
+        None
+        """
+
+        self.n_s = None
+        self.n_t = None
+        self.n_v = None
+        self.labels = None
+        self.alphas = None
+        self.alignment_method = alignment_method
+        if alignment_method == "parcel":
+            self.parcels = None
+
+        elif (
+            alignment_method == "searchlight"
+            or alignment_method == "ensemble_searchlight"
+        ):
+            self.searchlights = None
+            self.distances = None
+            self.radius = None
+
+        self.path = None
+        self.tuning_data = []
+        self.denoised_signal = []
+        self.decomp_method = decomp_method
+        self.tmpl_kind = tmpl_kind
+        self.latent_dim = latent_dim
+        self.n_jobs = n_jobs
+        self.cache = cache
+
+        if cache:
+            if id is None:
+                self.id = "default"
+            else:
+                self.id = id
+
+            path = os.path.join(os.getcwd(), f"cache/int/{self.id}")
+            # Check if cache folder exists
+            if not os.path.exists(path):
+                os.makedirs(path)
+
+            self.path = path
+
+    def fit(
+        self,
+        X_train,
+        searchlights=None,
+        parcels=None,
+        dists=None,
+        radius: int = 20,
+        verbose=True,
+    ):
+        """
+        Fits the IndividualizedNeuralTuning model to the training data.
+
+        Parameters:
+        --------
+
+        - X_train (array-like):
+            The training data of shape (n_subjects, n_samples, n_voxels).
+        - searchlights (array-like):
+            The searchlight indices for each subject, of shape (n_s, n_searchlights).
+        - parcels (array-like):
+            The parcel indices for each subject, of shape (n_s, n_parcels) (if not using searchlights)
+        - dists (array-like):
+            The distances of vertices to the center of their searchlight, of shape (n_searchlights, n_vertices_sl)
+        - radius (int, optional):
+            The radius of the searchlight sphere, in milimeters. Defaults to 20.
+        - verbose (bool, optional):
+            Whether to print progress information. Defaults to True.
+        - id (str, optional):
+            An identifier for caching purposes. Defaults to None.
+
+        Returns:
+        --------
+
+        - self (IndividualizedNeuralTuning):
+            The fitted model.
+        """
+
+        X_train_ = np.array(X_train, copy=True, dtype=np.float32)
+
+        self.n_s, self.n_t, self.n_v = X_train_.shape
+
+        self.tuning_data = np.empty(self.n_s, dtype=np.float32)
+        self.denoised_signal = np.empty(self.n_s, dtype=np.float32)
+
+        if searchlights is None:
+            self.regions = parcels
+        else:
+            self.regions = searchlights
+            self.distances = dists
+            self.radius = radius
+
+        # check for cached data
+        try:
+            self.denoised_signal = np.load(self.path + "/train_data_denoised.npy")
+            if verbose:
+                print("Loaded denoised data from cache")
+
+        except:  # noqa: E722
+            denoiser = RegionAlignment(
+                alignment_method=self.alignment_method,
+                n_jobs=self.n_jobs,
+                verbose=verbose,
+                path=self.path,
+            )
+            self.denoised_signal = denoiser.fit_transform(
+                X_train_,
+                regions=self.regions,
+                dists=dists,
+                radius=radius,
+            )
+            # Clear memory of the SearchlightAlignment object
+            denoiser = None
+
+        iterm = range(self.n_s)
+
+        if verbose:
+            iterm = tqdm(iterm)
+
+        # Stimulus matrix computation
+        if self.decomp_method is None:
+            full_signal = np.concatenate(self.denoised_signal, axis=1)
+            self.shared_response = stimulus_estimator(
+                full_signal, self.n_t, self.n_s, self.latent_dim
+            )
+
+            self.tuning_data = Parallel(n_jobs=self.n_jobs)(
+                delayed(tuning_estimator)(
+                    self.shared_response,
+                    self.denoised_signal[i],
+                    latent_dim=self.latent_dim,
+                )
+                for i in iterm
+            )
+
+        return self
+
+    def transform(self, X_test_, verbose=False):
+        """
+        Transforms the input test data using the hyperalignment model.
+
+        Args:
+            X_test_ (list of arrays):
+                The input test data.
+            verbose (bool, optional):
+                Whether to print verbose output. Defaults to False.
+            id (int, optional):
+                Identifier for the transformation. Defaults to None.
+
+        Returns:
+            numpy.ndarray: The transformed test data.
+        """
+
+        full_signal = np.concatenate(X_test_, axis=1, dtype=np.float32)
+
+        if verbose:
+            print("Predict : Computing stimulus matrix...")
+
+        if self.decomp_method is None:
+            S = stimulus_estimator(full_signal, self.n_t, self.n_s, self.latent_dim)
+
+        if verbose:
+            print("Predict : stimulus matrix shape: ", S.shape)
+
+        reconstructed_signal = [
+            reconstruct_signal(S, T_est) for T_est in self.tuning_data
+        ]
+        return np.array(reconstructed_signal, dtype=np.float32)
+
+    def get_shared_stimulus(self):
+        """
+        Returns the shared stimulus used for individualized neural tuning.
+
+        Returns:
+            The shared stimulus of shape (n_t, latent_dim) or (n_t, n_t).
+        """
+        return self.shared_response
+
+    def get_tuning_matrices(self):
+        """
+        Returns the tuning matrices as a NumPy array.
+
+        Returns:
+            numpy.ndarray: The tuning matrices of shape (n_s, latent_dim, n_v) or (n_s, n_t, n_v).
+        """
+        return np.array(self.tuning_data)
+
+    def clean_cache(self, id):
+        """
+        Removes the cache file associated with the given ID.
+
+        Args:
+            id (int): The ID of the cache file to be removed.
+        """
+        try:
+            os.remove("cache")
+        except:  # noqa: E722
+            print("No cache to remove")
+
+
+#######################################################################################
+# Computing decomposition
+
+
+def tuning_estimator(shared_response, target, latent_dim=None):
+    """
+    Estimate the tuning weights for individualized neural tuning.
+
+    Parameters:
+    --------
+        - shared_response (array-like):
+            The shared response matrix.
+        - target (array-like):
+            The target matrix.
+        - latent_dim (int, optional):
+            The number of latent dimensions. Defaults to None.
+
+    Returns:
+    --------
+        array-like: The estimated tuning weights.
+
+    """
+    if latent_dim is None:
+        return np.linalg.inv(shared_response).dot(target)
+    return np.linalg.pinv(shared_response).dot(target).astype(np.float32)
+
+
+def stimulus_estimator(full_signal, n_t, n_s, latent_dim=None):
+    """
+    Estimates the stimulus response using the given parameters.
+
+    Args:
+        full_signal (np.ndarray): The full signal data.
+        n_t (int): The number of time points.
+        n_s (int): The number of stimuli.
+        latent_dim (int, optional): The number of latent dimensions. Defaults to None.
+
+    Returns:
+        np.ndarray: The estimated shared response.
+    """
+    if latent_dim is not None and latent_dim < n_t:
+        U = svd_pca(full_signal)
+        U = U[:, :latent_dim]
+    else:
+        U, _, _ = safe_svd(full_signal)
+
+    shared_response = np.sqrt(n_s) * U
+    return shared_response.astype(np.float32)
+
+
+def reconstruct_signal(shared_response, individual_tuning):
+    """
+    Reconstructs the signal using the shared response and individual tuning.
+
+    Args:
+        shared_response (numpy.ndarray): The shared response of shape (n_t, n_t) or (n_t, latent_dim).
+        individual_tuning (numpy.ndarray): The individual tuning of shape (latent_dim, n_v) or (n_t, n_v).
+
+    Returns:
+        numpy.ndarray: The reconstructed signal of shape (n_t, n_v) (same shape as the original signal)
+    """
+    return (shared_response @ individual_tuning).astype(np.float32)
diff --git a/fmralign/hyperalignment/procrustes.py b/fmralign/hyperalignment/procrustes.py
new file mode 100644
index 0000000..4a4ca62
--- /dev/null
+++ b/fmralign/hyperalignment/procrustes.py
@@ -0,0 +1,74 @@
+import numpy as np
+from .linalg import safe_svd
+
+
+def procrustes(X, Y, reflection=True, scaling=False):
+    r"""
+    The orthogonal Procrustes algorithm.
+
+    The orthogonal Procrustes algorithm, also known as the classic
+    hyperalignment algorithm, is the first hyperalignment algorithm, which
+    was introduced in Haxby et al. (2011). It tries to align two
+    configurations in a high-dimensional space through an orthogonal
+    transformation. The transformation is an improper rotation, which is a
+    rotation with an optional reflection. Neither rotation nor reflection
+    changes the geometry of the configuration. Therefore, the geometry,
+    such as a representational dissimilarity matrix (RDM), remains the
+    same in the process.
+
+    Optionally, a global scaling can be added to the algorithm to further
+    improve alignment quality. That is, scaling the data with the same
+    factor for all directions. Different from rotation and reflection,
+    global scaling can change the geometry of the data.
+
+    Parameters
+    ----------
+    X : ndarray
+        The data matrix to be aligned to Y.
+    Y : ndarray
+        The "target" data matrix -- the matrix to be aligned to.
+    reflection : bool, default=True
+        Whether allows reflection in the transformation (True) or not
+        (False). Note that even with ``reflection=True``, the solution
+        may not contain a reflection if the alignment cannot be improved
+        by adding a reflection to the rotation.
+    scaling : bool, default=False
+        Whether allows global scaling (True) or not (False). Allowing
+        scaling can improve alignment quality, but it also changes the
+        geometry of data.
+
+    Returns
+    -------
+    T : ndarray
+        The transformation matrix which can be used to align X to Y.
+        Depending on the parameters ``reflection`` and ``scaling``, the
+        transformation can be a pure rotation, an improper rotation, or a
+        pure/improper rotation with global scaling.
+
+    Notes
+    -----
+    The algorithm tries to minimize the Frobenius norm of the difference
+    between transformed data ``X @ T`` and the target ``Y``:
+
+    .. math:: \underset{T}{\arg\min} \lVert XT - Y \rVert_F
+
+    The solution ``T`` differs depending on whether reflection and global
+    scaling are allowed. When it's a rotation (pure or improper), it's
+    often denoted as ``R``.
+    """
+
+    A = Y.T.dot(X).T
+    U, s, Vt = safe_svd(A, remove_mean=False)
+    T = np.dot(U, Vt)
+
+    if not reflection:
+        sign = np.sign(np.linalg.det(T))
+        s[-1] *= sign
+        if sign < 0:
+            T -= np.outer(U[:, -1], Vt[-1, :]) * 2
+
+    if scaling:
+        scale = s.sum() / (X.var(axis=0).sum() * X.shape[0])
+        T *= scale
+
+    return T
diff --git a/fmralign/hyperalignment/regions.py b/fmralign/hyperalignment/regions.py
new file mode 100644
index 0000000..d90069d
--- /dev/null
+++ b/fmralign/hyperalignment/regions.py
@@ -0,0 +1,517 @@
+"""Utilities for computing searchlights. Adapted from nilearn.\n
+Author: Denis Fouchard, INRIA Saclay, MIND, 2023.
+"""
+
+import functools
+import numpy as np
+from joblib import Parallel, delayed
+from nilearn import image, masking
+from nilearn._utils import check_niimg_4d, check_niimg_3d
+from nilearn.image.resampling import coord_transform
+import warnings
+from sklearn import neighbors
+from scipy.spatial import distance_matrix
+from nilearn._utils.niimg_conversions import (
+    _safe_get_data,
+)
+from tqdm import tqdm
+from .procrustes import procrustes
+from .linalg import ridge
+from .local_template import compute_template
+from fmralign._utils import _make_parcellation
+from nilearn.maskers import NiftiMasker
+from nibabel.nifti1 import Nifti1Image
+
+###################################################################################################
+# Compute parcels
+###################################################################################################
+
+
+def create_parcels_from_labels(labels: np.ndarray):
+    """
+    Create parcels from labels.
+
+    Args:
+        labels (np.ndarray): Array of labels.
+
+    Returns:
+        list: List of parcels, where each parcel is an array of indices.
+
+    """
+    n_labels = labels.max()
+    parcels = []
+    for i in range(1, n_labels + 1):
+        parcels.append(np.where(labels == i)[0])
+    return parcels
+
+
+def compute_parcels(
+    niimg, mask, n_parcels=100, verbose=True, smoothing_fwhm=5, n_jobs=1
+):
+    """
+    Compute parcels using a given mask and input image.
+
+    Parameters:
+    - niimg: Input image to be parcellated.
+    - mask: Mask image defining the region of interest.
+    - n_parcels: Number of parcels to be created (default: 100).
+    - verbose: Whether to print progress messages (default: True).
+    - smoothing_fwhm: Full Width at Half Maximum (FWHM) for smoothing (default: 5).
+    - n_jobs: Number of parallel jobs to run (default: 1).
+
+    Returns:
+    - Parcels created from the input image and mask.
+    """
+    if verbose:
+        print("[Loading/Parcel] : Parcellating...")
+
+    if isinstance(mask, Nifti1Image):
+        mask = NiftiMasker(mask_img=mask, standardize=True)
+    # Parcellation
+    indexes = [1]
+    labels = _make_parcellation(
+        imgs=niimg,
+        clustering_index=indexes,
+        clustering="kmeans",
+        n_pieces=n_parcels,
+        masker=mask,
+        smoothing_fwhm=mask.smoothing_fwhm,
+    )
+
+    parcels = create_parcels_from_labels(labels)
+
+    print("Minimum length parcel: ", min([len(p) for p in parcels]))
+
+    return parcels
+
+
+###################################################################################################
+# Computing searchlights
+###################################################################################################
+
+
+def _apply_mask_and_get_affinity(
+    seeds, niimg, radius, allow_overlap, mask_img=None, n_jobs=1
+):
+    """Get only the rows which are occupied by sphere \
+    at given seed locations and the provided radius.
+
+    Rows are in target_affine and target_shape space.
+
+    Parameters
+    ----------
+    seeds : List of triplets of coordinates in native space
+        Seed definitions. List of coordinates of the seeds in the same space
+        as target_affine.
+
+    niimg : 3D/4D Niimg-like object
+        See :ref:`extracting_data`.
+        Images to process.
+        If a 3D niimg is provided, a singleton dimension will be added to
+        the output to represent the single scan in the niimg.
+
+    radius : float
+        Indicates, in millimeters, the radius for the sphere around the seed.
+
+    allow_overlap : boolean
+        If False, a ValueError is raised if VOIs overlap
+
+    mask_img : Niimg-like object, optional
+        Mask to apply to regions before extracting signals. If niimg is None,
+        mask_img is used as a reference space in which the spheres 'indices are
+        placed.
+
+    Returns
+    -------
+    X : 2D numpy.ndarray
+        Signal for each brain voxel in the (masked) niimgs.
+        shape: (number of scans, number of voxels)
+
+    A : scipy.sparse.lil_matrix
+        Contains the boolean indices for each sphere.
+        shape: (number of seeds, number of voxels)
+
+    """
+    seeds = list(seeds)
+
+    # Compute world coordinates of all in-mask voxels.
+    if niimg is None:
+        mask, affine = masking._load_mask_img(mask_img)
+        # Get coordinate for all voxels inside of mask
+        mask_coords = np.asarray(np.nonzero(mask)).T.tolist()
+        X = None
+
+    elif mask_img is not None:
+        affine = niimg.affine
+        mask_img = check_niimg_3d(mask_img)
+        mask_img = image.resample_img(
+            mask_img,
+            target_affine=affine,
+            target_shape=niimg.shape[:3],
+            interpolation="nearest",
+        )
+        mask, _ = masking._load_mask_img(mask_img)
+        mask_coords = list(zip(*np.where(mask != 0)))
+
+        X = masking._apply_mask_fmri(niimg, mask_img)
+
+    elif niimg is not None:
+        affine = niimg.affine
+        if np.isnan(np.sum(_safe_get_data(niimg))):
+            warnings.warn(
+                "The imgs you have fed into fit_transform() contains NaN "
+                "values which will be converted to zeroes."
+            )
+            X = _safe_get_data(niimg, True).reshape([-1, niimg.shape[3]]).T
+        else:
+            X = _safe_get_data(niimg).reshape([-1, niimg.shape[3]]).T
+
+        mask_coords = list(np.ndindex(niimg.shape[:3]))
+
+    else:
+        raise ValueError("Either a niimg or a mask_img must be provided.")
+
+    # For each seed, get coordinates of nearest voxel
+    nearests = []
+    for sx, sy, sz in seeds:
+        nearest = np.round(
+            image.resampling.coord_transform(sx, sy, sz, np.linalg.inv(affine))
+        )
+        nearest = nearest.astype(int)
+        nearest = (nearest[0], nearest[1], nearest[2])
+        try:
+            nearests.append(mask_coords.index(nearest))
+        except ValueError:
+            nearests.append(None)
+
+    mask_coords = np.asarray(list(zip(*mask_coords)))
+    mask_coords = image.resampling.coord_transform(
+        mask_coords[0], mask_coords[1], mask_coords[2], affine
+    )
+    mask_coords = np.asarray(mask_coords).T
+
+    clf = neighbors.NearestNeighbors(radius=radius, n_jobs=n_jobs)
+    A = clf.fit(mask_coords).radius_neighbors_graph(seeds)
+    A = A.tolil()
+    for i, nearest in enumerate(nearests):
+        if nearest is None:
+            continue
+
+        A[i, nearest] = True
+
+    mask_coords_floats = mask_coords.copy()
+
+    # Include the voxel containing the seed itself if not masked
+    mask_coords = mask_coords.astype(int).tolist()
+    for i, seed in enumerate(seeds):
+        try:
+            A[i, mask_coords.index(list(map(int, seed)))] = True
+        except ValueError:
+            # seed is not in the mask
+            pass
+
+    sphere_sizes = np.asarray(A.tocsr().sum(axis=1)).ravel()
+    empty_spheres = np.nonzero(sphere_sizes == 0)[0]
+    if len(empty_spheres) != 0:
+        raise ValueError(f"These spheres are empty: {empty_spheres}")
+
+    if (not allow_overlap) and np.any(A.sum(axis=0) >= 2):
+        raise ValueError("Overlap detected between spheres")
+
+    return X, A, mask_coords_floats
+
+
+def compute_searchlights(
+    niimg,
+    mask_img,
+    process_mask_img=None,
+    radius=20,
+    return_dist_mat=False,
+    n_jobs=1,
+):
+    """Implement search_light analysis using an arbitrary type of classifier.
+
+    Parameters
+    ----------
+    miimg : Niimg-like object
+        See :ref:`extracting_data`.
+        4D image.
+
+    mask_img : Niimg-like object
+        See :ref:`extracting_data`.
+        Boolean image giving location of voxels containing usable signals.
+
+    process_mask_img : Niimg-like object, optional
+        See :ref:`extracting_data`.
+        Boolean image giving voxels on which searchlight should be
+        computed.
+
+    radius : float, optional
+        radius of the searchlight ball, in millimeters. Defaults to 20.
+
+    return_dist_mat : bool, optional
+        Whether to return the distance matrix between voxels in the mask.
+        Defaults to False.
+
+    groups : array-like of shape (n_samples,), optional
+        Labels of samples for each subject. If provided, the searchlights
+        will be computed within each group separately.
+        Defaults to None.
+
+    verbose : int, optional
+        Verbosity level (0 means no message).
+        Defaults to 0.
+
+    Returns
+    -------
+    X : 2D numpy.ndarray
+        Signal for each brain voxel in the (masked) niimgs.
+        shape: (number of scans, number of voxels)
+
+    A_list : list of lists
+        Contains the boolean indices for each sphere.
+        shape: (number of seeds, number of voxels)
+
+    dist_matrix : 2D numpy.ndarray
+        Distance matrix between voxels in the mask.
+        shape: (number of voxels, number of voxels)
+
+    dists : list of lists
+        Contains the distance between each voxel and the seed.
+        shape: (number of seeds, number of voxels)
+
+    """
+
+    # check if image is 4D
+    niimg = check_niimg_4d(niimg)
+
+    # Get the seeds
+    if process_mask_img is None:
+        process_mask_img = mask_img
+
+    # Compute world coordinates of the seeds
+    process_mask, process_mask_affine = masking._load_mask_img(process_mask_img)
+    process_mask_coords = np.where(process_mask != 0)
+    process_mask_coords = coord_transform(
+        process_mask_coords[0],
+        process_mask_coords[1],
+        process_mask_coords[2],
+        process_mask_affine,
+    )
+    process_mask_coords = np.asarray(process_mask_coords).T
+
+    X, A, mask_coords = _apply_mask_and_get_affinity(
+        process_mask_coords,
+        niimg,
+        radius=radius,
+        allow_overlap=True,
+        mask_img=mask_img,
+        n_jobs=n_jobs,
+    )
+
+    A_list = []
+    for i in range(A.shape[0]):
+        A_list.append(A[i].nonzero()[1].tolist())
+
+    dist_matrix = distance_matrix(mask_coords, mask_coords)
+    dists = []
+    for i, sl in enumerate(A_list):
+        dists.append(dist_matrix[i, sl])
+
+    if return_dist_mat:
+        return X, A_list, dist_matrix
+
+    return X, A_list, dists
+
+
+def searchlight_weights(searchlights, dists, radius):
+    nv = np.concatenate(searchlights).max() + 1
+    weights_sum = np.zeros((nv,))
+    for sl, d in zip(searchlights, dists):
+        w = (radius - d) / radius
+        weights_sum[sl] += w
+    # print(np.percentile(weights_sum, np.linspace(0, 100, 11)))
+    weights = []
+    for sl, d in zip(searchlights, dists):
+        w = (radius - d) / radius
+        w /= weights_sum[sl]
+        weights.append(w)
+    return weights
+
+
+###################################################################################################
+# Hyperalignment
+###################################################################################################
+
+
+def iter_hyperalignment(
+    X,
+    Y,
+    searchlights,
+    sl_func,
+    weights=None,
+    verbose=False,
+):
+    """
+    Perform searchlight hyperalignment on the given data.
+
+    Args:
+        X (ndarray): The source data matrix of shape (n_samples, n_features).
+        Y (ndarray): The target data matrix of shape (n_samples, n_features).
+        searchlights (list): List of searchlight indices.
+        dists (ndarray): distances of vertices to the center of their searchlight, of shape (n_searchlights, n_vertices_sl)
+        radius (float): Radius of the searchlight.
+        sl_func (function): Function to compute the searchlight transformation.
+        weighted (bool, optional): Whether to use weighted searchlights. Defaults to True.
+
+    Returns:
+        ndarray: The transformed matrix T of shape (n_features, n_features).
+    """
+    Yhat = np.zeros_like(X, dtype=np.float32)
+
+    if weights is not None:
+        zip_iter = zip(searchlights, weights)
+        if verbose:
+            zip_iter = tqdm(zip_iter, leave=False)
+        for sl, w in zip_iter:
+            x, y = X[:, sl], Y[:, sl]
+            t = sl_func(x, y)
+            Yhat[:, sl] += x @ t * w[np.newaxis]
+            del t
+    else:
+        searchlights_iter = searchlights
+        if verbose:
+            searchlights_iter = tqdm(searchlights_iter, leave=False)
+        for sl in searchlights_iter:
+            x, y = X[:, sl], Y[:, sl]
+            t = sl_func(x, y)
+            Yhat[:, sl] += x @ t
+            del t
+
+    return Yhat
+
+
+def piece_procrustes(
+    X,
+    Y,
+    regions,
+    weights=None,
+    T0=None,
+    reflection=True,
+    scaling=False,
+):
+    """
+    Applies searchlight hyperalignment using Procrustes alignment.
+
+    Args:
+        X (array-like): The source data matrix of shape (n_samples, n_features).
+        Y (array-like): The target data matrix of shape (n_samples, n_features).
+        searchlights (array-like): The indices of the searchlight regions.
+        dists (array-like): distances of vertices to the center of their searchlight, of shape (n_searchlights, n_vertices_sl)
+        radius (float): The radius of the searchlight region.
+        T0 (array-like, optional): The initial transformation matrix. Defaults to None.
+        reflection (bool, optional): Whether to allow reflection in the alignment. Defaults to True.
+        scaling (bool, optional): Whether to allow scaling in the alignment. Defaults to False.
+        weighted (bool, optional): Whether to use weighted Procrustes alignment. Defaults to True.
+
+    Returns:
+        array-like: The transformation matrix T.
+
+    """
+    sl_func = functools.partial(procrustes, reflection=reflection, scaling=scaling)
+    T = iter_hyperalignment(
+        X,
+        Y,
+        regions,
+        T0=T0,
+        weights=weights,
+        sl_func=sl_func,
+    )
+    return T
+
+
+def piece_ridge(
+    X,
+    Y,
+    regions,
+    alpha=1e3,
+    weights=None,
+    verbose=False,
+):
+    """
+    Perform searchlight ridge regression for hyperalignment.
+
+    Args:
+        X (array-like): The source data matrix of shape (n_samples, n_features).
+        Y (array-like): The target data matrix of shape (n_samples, n_features).
+        searchlights (array-like): The indices of the searchlight regions.
+        dists (array-like): distances of vertices to the center of their searchlight, of shape (n_searchlights, n_vertices_sl)
+        radius (float): The radius of the searchlight region.
+        T0 (array-like, optional): The initial transformation matrix. Defaults to None.
+        alpha (float, optional): The regularization parameter for ridge regression. Defaults to 1e3.
+        weighted (bool, optional): Whether to use weighted ridge regression. Defaults to True.
+
+    Returns:
+        array-like: The transformation matrix T.
+
+    """
+    sl_func = functools.partial(ridge, alpha=alpha)
+    T = iter_hyperalignment(
+        X,
+        Y,
+        regions,
+        sl_func=sl_func,
+        weights=weights,
+        verbose=verbose,
+    )
+    return T
+
+
+def template(
+    X,
+    regions,
+    n_jobs=1,
+    template_kind="searchlight_pca",
+    verbose=False,
+    weights=None,
+):
+    """
+    Compute a template by aggregating local templates within searchlights.
+
+    Args:
+        X (numpy.ndarray): The input data matrix of shape (n_subjects, n_samples, n_features).
+        regions (list): The indices of the searchlight/parcels regions.
+        n_jobs (int, optional): The number of parallel jobs to run. Defaults to -1.
+        template_kind (str, optional): The kind of template to compute. Defaults to "pca".
+        verbose (bool, optional): Whether to display progress. Defaults to False.
+
+    Returns:
+        numpy.ndarray: The computed template of shape (n_features,).
+
+    """
+
+    if verbose:
+        iterator = tqdm(regions, leave=False)
+        iterator.set_description("Computing local templates")
+    else:
+        iterator = regions
+    with Parallel(n_jobs=n_jobs, batch_size=1, verbose=1) as parallel:
+        local_templates = parallel(
+            delayed(compute_template)(
+                X,
+                region=region,
+                kind=template_kind,
+                max_npc=None,
+                common_topography=True,
+            )
+            for region in iterator
+        )
+
+    template = np.zeros_like(X[0])
+    if weights is not None:
+        for local_template, w, region in zip(local_templates, weights, regions):
+            template[:, region] += local_template * w[np.newaxis]
+    else:
+        for local_template, region in zip(local_templates, regions):
+            template[:, region] += local_template
+    return template
diff --git a/fmralign/hyperalignment/regions_alignment.py b/fmralign/hyperalignment/regions_alignment.py
new file mode 100644
index 0000000..4ba62bf
--- /dev/null
+++ b/fmralign/hyperalignment/regions_alignment.py
@@ -0,0 +1,212 @@
+import numpy as np
+from sklearn.base import BaseEstimator, TransformerMixin
+from .regions import (
+    template,
+    piece_ridge,
+    searchlight_weights,
+)
+import os
+from joblib import Parallel, delayed
+
+
+class RegionAlignment(BaseEstimator, TransformerMixin):
+    """Searchlight alignment model. This model decomposes the data into a
+    global template and a linear transformation for each subject.
+    The global template is computed using a searchlight approach.
+    The linear transformation is computed using a ridge regression.
+    This step is enssential to the hyperalignment model, as it is
+    used as a denoiser for the data.
+    Parameters
+    ----------
+    alignment_method : str, default="ridge"
+        The alignment method to use. Can be "ridge" or "ensemble_ridge".
+    template_kind : str, default="pca"
+        The kind of template to use. Can be "pca" or "mean".
+    demean : bool, default=False
+        Whether to demean the data before alignment.
+    verbose : bool, default=True
+        Whether to display progress bar.
+    n_jobs : int, default=-1
+    """
+
+    def __init__(
+        self,
+        alignment_method="searchlight_ridge",
+        template_kind="searchlight_pca",
+        verbose=True,
+        path="cache/int/",
+        cache=True,
+        n_jobs=-1,
+    ):
+        self.W = []
+        self.Xhat = []
+        self.n_s = None
+        self.n_t = None
+        self.n_v = None
+        self.warp_alignment_method = alignment_method
+        self.template_kind = template_kind
+        self.verbose = verbose
+        self.n_jobs = n_jobs
+        self.regions = None
+        self.distances = None
+        self.radius = None
+        self.weights = None
+        self.path = path
+        self.cache = (path is not None) and (cache)
+
+        if self.cache:
+            # Check if cache folder exists
+            if not os.path.exists(self.path):
+                os.makedirs(self.path)
+
+    def compute_linear_transformation(self, x_i, template, i: int = 0, save=True):
+        """Compute the linear transformation W_i for a given subject.
+        ----------
+        Parameters
+        ----------
+        x_i : ndarray of shape (n_samples, n_voxels)
+            The brain images for one subject.
+            Those are the B_1, ..., B_n in the paper.
+        template : ndarray of shape (n_samples, n_voxels)
+            The global template M.
+
+        Returns
+        -------
+        Xhat : ndarray of shape (n_samples, n_voxels)
+            The denoised estimation signal for each subject.
+        """
+        try:
+            W_p = np.load(self.path + (f"/train_data_W_{i}.npy"))
+            if self.verbose:
+                print(f"Loaded W_{i} from cache")
+            x_hat = template.dot(W_p)
+            del W_p
+            return x_hat
+
+        except:  # noqa E722
+            if self.verbose:
+                print(f"No cache found, computing W_{i}")
+
+            x_hat = piece_ridge(
+                X=x_i,
+                Y=template,
+                regions=self.regions,
+                weights=self.weights,
+                verbose=self.verbose,
+            )
+
+            if self.cache:
+                np.save(self.path + (f"/train_data_W_{i}.npy"), x_hat)
+        return x_hat
+
+    def fit_transform(
+        self,
+        X: np.ndarray,
+        regions=None,
+        dists=None,
+        radius=None,
+        weights=None,
+        id=None,
+    ):
+        """From brain imgs compute the INT model (M, Ws, S)
+          with the given parameters)
+
+        Parameters
+        ----------
+        X : list of ndarray of shape (n_samples, n_voxels)
+            The brain images for one subject.
+        searchlights : list of searchlights
+            The searchlight indices.
+        dists : list of distances
+
+        radius : int
+            The radius of the searchlight (in millimeters)
+
+        Returns
+        -------
+        Xhat : list of ndarray of shape (n_samples, n_voxels)
+            The denoised estimations B_1, ... B_p for each subject.
+        """
+
+        self.FUNC = "RegionAlignment"
+        if dists is not None and radius is not None:
+            self.FUNC = "SearchlightAlignment"
+
+        if self.verbose:
+            print(f"[{self.FUNC}] Shape of input data: ", X.shape)
+
+        try:
+            self.Xhat = np.load(self.path + "/train_data_denoised.npy")
+            if self.verbose:
+                print(f"[{self.FUNC}] Loaded denoised data from cache")
+            return self.Xhat
+        except:  # noqa E722
+            if self.verbose:
+                print(f"[{self.FUNC}] No cache found, computing denoised data")
+
+        self.n_s, self.n_t, self.n_v = X.shape
+        self.regions = regions
+        self.FUNC = "ParcelAlignment"
+
+        if weights is None and dists is not None and radius is not None:
+            self.distances = dists
+            self.radius = radius
+            self.FUNC = "SearchlightAlignment"
+
+        # Compute global template M (sl_template)
+        if self.verbose:
+            print(f"[{self.FUNC}]Computing global template M ...")
+
+        try:
+            sl_template = np.load(self.path + ("/train_data_template.npy"))
+            if self.verbose:
+                print("Loaded template from cache")
+        except:  # noqa E722
+            if self.verbose:
+                print(f"[{self.FUNC}] No cache found, computing template")
+            if dists is None or radius is None:
+                self.weights = None
+            else:
+                self.weights = searchlight_weights(
+                    searchlights=regions, dists=dists, radius=radius
+                )
+            sl_template = template(
+                X,
+                regions=regions,
+                n_jobs=self.n_jobs,
+                template_kind=self.template_kind,
+                verbose=self.verbose,
+                weights=self.weights,
+            )
+
+        if self.cache:
+            np.save(self.path + ("/train_data_template.npy"), sl_template)
+            if self.verbose:
+                print(f"[{self.FUNC}] Saved template to cache")
+
+        self.Xhat = Parallel(n_jobs=self.n_jobs)(
+            delayed(self.compute_linear_transformation)(X[i], sl_template, i)
+            for i in range(self.n_s)
+        )
+
+        if id is None:
+            id = np.random.randint(0, 1000000)
+
+        if self.cache:
+            np.save(self.path + ("/train_data_denoised.npy"), self.Xhat)
+
+        return np.array(self.Xhat)
+
+    def get_linear_transformations(self):
+        """Return the linear transformations W_1, ... W_p for each subject.
+        ----------
+        Returns
+        -------
+        W : list of ndarray of shape (n_voxels, n_voxels)
+            The linear transformations W_1, ... W_p for each subject.
+        """
+        return np.array(self.W)
+
+    def get_denoised_estimation(self):
+        """Return the denoised estimations B_1, ... B_p for each subject."""
+        return self.Xhat
diff --git a/fmralign/template_alignment.py b/fmralign/template_alignment.py
index 95f0473..0730ca4 100644
--- a/fmralign/template_alignment.py
+++ b/fmralign/template_alignment.py
@@ -11,9 +11,7 @@
 from nilearn.image import concat_imgs, index_img, load_img
 from nilearn.maskers._masker_validation import _check_embedded_nifti_masker
 from sklearn.base import BaseEstimator, TransformerMixin
-
 from fmralign.pairwise_alignment import PairwiseAlignment
-from hyperalignment.hyperalignment import HyperAlignment
 
 
 def _rescaled_euclidean_mean(imgs, masker, scale_average=False):
@@ -379,23 +377,6 @@ def fit(self, imgs):
         """
 
         # Alignment method is hyperalignment (too different from other methods)
-        if self.alignment_method == "slha":
-            self.model = HyperAlignment(
-                method="searchlight",
-                n_jobs=self.n_jobs,
-            )
-            return self.model.fit(
-                imgs, self.masker_, self.masker_.mask_img_, verbose=self.verbose
-            )
-
-        elif self.alignment_method == "pcha":
-            self.model = HyperAlignment(
-                method="parcels",
-                n_jobs=self.n_jobs,
-            )
-            return self.model.fit(
-                imgs, self.masker_, self.masker_.mask_img_, verbose=self.verbose
-            )
 
         # Check if the input is a list, if list of lists, concatenate each subjects
         # data into one unique image.

From 4ce6f10fe2f24d4f6d5acd0f3fdea6bbd84857f6 Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Wed, 3 Jan 2024 12:06:06 +0100
Subject: [PATCH 05/69] fix ot/ott import problems

---
 fmralign/alignment_methods.py | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/fmralign/alignment_methods.py b/fmralign/alignment_methods.py
index 38022bc..380d9d4 100644
--- a/fmralign/alignment_methods.py
+++ b/fmralign/alignment_methods.py
@@ -2,7 +2,7 @@
 """Module implementing alignment estimators on ndarrays."""
 import warnings
 
-import ot
+import ott
 import numpy as np
 import scipy
 from joblib import Parallel, delayed
@@ -17,6 +17,11 @@
 from .hyperalignment.regions_alignment import RegionAlignment
 from .hyperalignment.linalg import safe_svd, svd_pca
 
+import jax
+from ott.geometry import costs, geometry
+from ott.problems.linear import linear_problem
+from ott.solvers.linear import sinkhorn
+
 
 def scaled_procrustes(X, Y, scaling=False, primal=None):
     """
@@ -389,10 +394,10 @@ def fit(self, X, Y):
             M = cdist(X.T, Y.T, metric=self.metric)
 
             if self.solver == "exact":
-                self.R = ot.lp.emd(a, b, M) * n
+                self.R = ott.lp.emd(a, b, M) * n
             else:
                 self.R = (
-                    ot.sinkhorn(
+                    ott.sinkhorn(
                         a,
                         b,
                         M,
@@ -445,10 +450,6 @@ def fit(self, X, Y):
         Y: (n_samples, n_features) nd array
             target data
         """
-        import jax
-        from ott.geometry import costs, geometry
-        from ott.problems.linear import linear_problem
-        from ott.solvers.linear import sinkhorn
 
         if self.metric == "euclidean":
             cost_matrix = costs.Euclidean().all_pairs(x=X.T, y=Y.T)

From 21201dd97544b84bb7629a06ad13b8acfd6b0b65 Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Wed, 3 Jan 2024 13:24:14 +0100
Subject: [PATCH 06/69] fix: rm hyperalignment requirements and fix tests

---
 fmralign/__init__.py                     |   2 +-
 fmralign/alignment_methods.py            |   7 +-
 fmralign/generate_data.py                | 141 +++++++++++++++++++++++
 fmralign/tests/test_alignment_methods.py |  73 ++++++++++++
 pyproject.toml                           |   1 -
 5 files changed, 218 insertions(+), 6 deletions(-)
 create mode 100644 fmralign/generate_data.py

diff --git a/fmralign/__init__.py b/fmralign/__init__.py
index 3a8d6d5..8e6a9c1 100644
--- a/fmralign/__init__.py
+++ b/fmralign/__init__.py
@@ -1 +1 @@
-from ._version import __version__  # noqa: F401
+# from ._version import __version__  # noqa: F401
diff --git a/fmralign/alignment_methods.py b/fmralign/alignment_methods.py
index 380d9d4..4bf4f7e 100644
--- a/fmralign/alignment_methods.py
+++ b/fmralign/alignment_methods.py
@@ -1,8 +1,6 @@
 # -*- coding: utf-8 -*-
 """Module implementing alignment estimators on ndarrays."""
 import warnings
-
-import ott
 import numpy as np
 import scipy
 from joblib import Parallel, delayed
@@ -377,6 +375,7 @@ def fit(self, X, Y):
         Y: (n_samples, n_features) nd array
             target data
         """
+        import ot
 
         n = len(X.T)
         if n > 5000:
@@ -394,10 +393,10 @@ def fit(self, X, Y):
             M = cdist(X.T, Y.T, metric=self.metric)
 
             if self.solver == "exact":
-                self.R = ott.lp.emd(a, b, M) * n
+                self.R = ot.lp.emd(a, b, M) * n
             else:
                 self.R = (
-                    ott.sinkhorn(
+                    ot.sinkhorn(
                         a,
                         b,
                         M,
diff --git a/fmralign/generate_data.py b/fmralign/generate_data.py
new file mode 100644
index 0000000..9539f88
--- /dev/null
+++ b/fmralign/generate_data.py
@@ -0,0 +1,141 @@
+import numpy as np
+from fastsrm.srm import projection
+
+
+def generate_dummy_signal(
+    n_s: int,
+    n_t: int,
+    n_v: int,
+    S_std=1,
+    latent_dim=None,
+    T_mean=0,
+    T_std=1,
+    SNR=1,
+    generative_method="custom",
+    seed=0,
+):
+    """Generate dummy signal for testing INT model
+
+    Parameters
+    ----------
+    n_s : int
+        Number of subjects.
+    n_t : int
+        Number of timepoints.
+    n_v : int
+        Number of voxels.
+    S_std : float, default=1
+        Standard deviation of latent variables.
+    latent_dim: int, defult=None
+        Number of latent dimensions. Defualts to n_t
+    T_mean : float
+        Mean of weights.
+    T_std : float
+        Standard deviation of weights.
+    SNR : float
+        Signal-to-noise ratio.
+
+
+    Returns
+    -------
+    imgs_train : ndarray of shape (n_s, n_t, n_v)
+        Training data.
+    imgs_test : ndarray of shape (n_s, n_t, n_v)
+        Testing data.
+    S_train : ndarray of shape (n_t, latent_dim)
+        Training latent variables.
+    S_test : ndarray of shape (n_t, latent_dim)
+        Testing latent variables.
+    Ts : ndarray of shape (n_s, latent_dim , n_v)
+        Tuning matrices.
+    """
+    if latent_dim is None:
+        latent_dim = n_t
+
+    rng = np.random.RandomState(seed=seed)
+
+    if generative_method == "custom":
+        sigma = n_s * np.arange(1, latent_dim + 1)
+        # np.random.shuffle(sigma)
+        # Generate common signal matrix
+        S_train = S_std * np.random.randn(n_t, latent_dim)
+        # Normalize each row to have unit norm
+        S_train = S_train / np.linalg.norm(S_train, axis=0, keepdims=True)
+        S_train = S_train @ np.diag(sigma)
+        S_test = S_std * np.random.randn(n_t, latent_dim)
+        S_test = S_test / np.linalg.norm(S_test, axis=0, keepdims=True)
+        S_test = S_test @ np.diag(sigma)
+
+    elif generative_method == "fastsrm":
+        Sigma = rng.dirichlet(np.ones(latent_dim), 1).flatten()
+        S_train = np.sqrt(Sigma)[:, None] * rng.randn(n_t, latent_dim)
+        S_test = np.sqrt(Sigma)[:, None] * rng.randn(n_t, latent_dim)
+
+    elif generative_method == "multiviewica":
+        S_train = np.random.laplace(size=(n_t, latent_dim))
+        S_test = np.random.laplace(size=(n_t, latent_dim))
+
+    # Generate indiivdual spatial components
+    data_train, data_test = [], []
+    Ts = []
+    for _ in range(n_s):
+        if generative_method == "custom" or generative_method == "multiviewica":
+            W = T_mean + T_std * np.random.randn(latent_dim, n_v)
+        else:
+            W = projection(rng.randn(latent_dim, n_v))
+
+        Ts.append(W)
+        X_train = S_train @ W
+        N = np.random.randn(n_t, n_v)
+        N = (
+            N
+            * np.linalg.norm(X_train)
+            / (SNR * np.linalg.norm(N, axis=0, keepdims=True))
+        )
+        X_train += N
+        data_train.append(X_train)
+        X_test = S_test @ W
+        N = np.random.randn(n_t, n_v)
+        N = (
+            N
+            * np.linalg.norm(X_test)
+            / (SNR * np.linalg.norm(N, axis=0, keepdims=True))
+        )
+        X_test += N
+        data_test.append(X_test)
+
+    data_train = np.array(data_train)
+    data_test = np.array(data_test)
+    return data_train, data_test, S_train, S_test, Ts
+
+
+def generate_dummy_searchlights(
+    n_searchlights: int,
+    n_v: int,
+    radius: int,
+    seed: int = 0,
+):
+    """Generate dummy searchlights for testing INT model
+
+    Parameters
+    ----------
+    n_searchlights : int
+        Number of searchlights.
+    n_v : int
+        Number of voxels.
+    radius : int
+        Radius of searchlights.
+    seed : int
+        Random seed.
+
+    Returns
+    -------
+    searchlights : ndarray of shape (n_searchlights, 5)
+        Searchlights.
+    dists : ndarray of shape (n_searchlights, 5)
+        Distances.
+    """
+    rng = np.random.RandomState(seed=seed)
+    searchlights = rng.randint(n_v, size=(n_searchlights, 5))
+    dists = rng.randint(radius, size=searchlights.shape)
+    return searchlights, dists
diff --git a/fmralign/tests/test_alignment_methods.py b/fmralign/tests/test_alignment_methods.py
index 7a9059e..4f9aa2a 100644
--- a/fmralign/tests/test_alignment_methods.py
+++ b/fmralign/tests/test_alignment_methods.py
@@ -5,6 +5,12 @@
 from scipy.sparse import csc_matrix
 from scipy.linalg import orthogonal_procrustes
 
+# add directory to path
+import sys
+import os
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
 from fmralign.alignment_methods import (
     DiagonalAlignment,
     Hungarian,
@@ -13,11 +19,20 @@
     POTAlignment,
     RidgeAlignment,
     ScaledOrthogonalAlignment,
+    Hyperalignment as INT,
     _voxelwise_signal_projection,
     optimal_permutation,
     scaled_procrustes,
 )
 from fmralign.tests.utils import zero_mean_coefficient_determination
+from fmralign.generate_data import (
+    generate_dummy_signal,
+    generate_dummy_searchlights,
+)
+
+from hyperalignment.regions import create_parcels_from_labels
+import os
+import shutil
 
 
 def test_scaled_procrustes_algorithmic():
@@ -198,3 +213,61 @@ def test_ott_backend():
     algo.fit(X, Y)
     old_implem.fit(X, Y)
     assert_array_almost_equal(algo.R, old_implem.R, decimal=3)
+
+
+def test_searchlight_alignment_with_ridge():
+    n_voxels = 99
+    n_time_points = 149
+    n_searchlights = 92
+    n_subjects = 5
+
+    radius = 20
+    searchlights, dists = generate_dummy_searchlights(
+        n_searchlights=n_searchlights,
+        n_v=n_voxels,
+        radius=radius,
+        seed=0,
+    )
+
+    X_train, X_test, _, _, _ = generate_dummy_signal(
+        n_t=n_time_points, n_v=n_voxels, n_s=n_subjects
+    )
+
+    model = INT()
+    model.fit(X_train, searchlights, dists, radius=radius)
+    X_pred = model.transform(X_test)
+
+    # assert that the saved cached files exist
+    b = os.path.exists("cache/")
+    shutil.rmtree("cache")
+    assert b
+    assert X_pred.shape == X_test.shape
+
+
+def test_parcel_alignment():
+    n_voxels = 99
+    n_time_points = 149
+    n_subjects = 5
+
+    n_parcels = 10
+    labels = np.arange(n_voxels) % n_parcels + 1
+    parcels = create_parcels_from_labels(labels)
+
+    X_train, X_test, _, _, _ = generate_dummy_signal(
+        n_t=n_time_points, n_v=n_voxels, n_s=n_subjects
+    )
+
+    model = INT(n_jobs=-1, alignment_method="parcel")
+    model.fit(X_train, parcels=parcels)
+    X_pred = model.transform(X_test)
+
+    # assert that the saved cached files exist
+    b = os.path.exists("cache/")
+    shutil.rmtree("cache")
+    assert b
+    assert X_pred.shape == X_test.shape
+
+
+# Cleaning up the cache folder
+if os.path.exists("cache/"):
+    shutil.rmtree("cache")
diff --git a/pyproject.toml b/pyproject.toml
index 289c829..82467eb 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -26,7 +26,6 @@ dependencies = [
     "nibabel",
     "nilearn@git+https://github.com/nilearn/nilearn#egg=1b11edf5789813d20995e7744bc58f3a2b6bc4ea",
     "POT",
-    "hyperalignment"
 ]
 dynamic = ["version"]
 

From 15078fd61200f9f2a3af9785a559679e89fd1384 Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Wed, 3 Jan 2024 13:30:38 +0100
Subject: [PATCH 07/69] fix: doc

---
 fmralign/generate_data.py | 21 +++++++++++++--------
 1 file changed, 13 insertions(+), 8 deletions(-)

diff --git a/fmralign/generate_data.py b/fmralign/generate_data.py
index 9539f88..d85f26b 100644
--- a/fmralign/generate_data.py
+++ b/fmralign/generate_data.py
@@ -1,3 +1,5 @@
+"""Some functions to generate toy fMRI data using shared stimulus information."""
+
 import numpy as np
 from fastsrm.srm import projection
 
@@ -34,6 +36,10 @@ def generate_dummy_signal(
         Standard deviation of weights.
     SNR : float
         Signal-to-noise ratio.
+    generative_method : str, default="custom"
+        Method for generating data. Options are "custom", "fastsrm".
+    seed : int
+        Random seed.
 
 
     Returns
@@ -71,15 +77,11 @@ def generate_dummy_signal(
         S_train = np.sqrt(Sigma)[:, None] * rng.randn(n_t, latent_dim)
         S_test = np.sqrt(Sigma)[:, None] * rng.randn(n_t, latent_dim)
 
-    elif generative_method == "multiviewica":
-        S_train = np.random.laplace(size=(n_t, latent_dim))
-        S_test = np.random.laplace(size=(n_t, latent_dim))
-
     # Generate indiivdual spatial components
     data_train, data_test = [], []
     Ts = []
     for _ in range(n_s):
-        if generative_method == "custom" or generative_method == "multiviewica":
+        if generative_method == "custom":
             W = T_mean + T_std * np.random.randn(latent_dim, n_v)
         else:
             W = projection(rng.randn(latent_dim, n_v))
@@ -113,6 +115,7 @@ def generate_dummy_searchlights(
     n_searchlights: int,
     n_v: int,
     radius: int,
+    sl_size: int = 5,
     seed: int = 0,
 ):
     """Generate dummy searchlights for testing INT model
@@ -125,17 +128,19 @@ def generate_dummy_searchlights(
         Number of voxels.
     radius : int
         Radius of searchlights.
+    sl_size : int, default=5
+        Size of each searchlight (easier for dummy signal generation).
     seed : int
         Random seed.
 
     Returns
     -------
-    searchlights : ndarray of shape (n_searchlights, 5)
+    searchlights : ndarray of shape (n_searchlights, sl_size)
         Searchlights.
-    dists : ndarray of shape (n_searchlights, 5)
+    dists : ndarray of shape (n_searchlights, sl_size)
         Distances.
     """
     rng = np.random.RandomState(seed=seed)
-    searchlights = rng.randint(n_v, size=(n_searchlights, 5))
+    searchlights = rng.randint(n_v, size=(n_searchlights, sl_size))
     dists = rng.randint(radius, size=searchlights.shape)
     return searchlights, dists

From 12fe4a83f1a1cdf50898531a0ea4a99401ad8639 Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Wed, 3 Jan 2024 14:28:23 +0100
Subject: [PATCH 08/69] fix: delete procrustes.py

---
 fmralign/hyperalignment/linalg.py         | 72 ++++++++++++++++++++++
 fmralign/hyperalignment/local_template.py |  2 +-
 fmralign/hyperalignment/procrustes.py     | 74 -----------------------
 fmralign/hyperalignment/regions.py        |  2 +-
 4 files changed, 74 insertions(+), 76 deletions(-)
 delete mode 100644 fmralign/hyperalignment/procrustes.py

diff --git a/fmralign/hyperalignment/linalg.py b/fmralign/hyperalignment/linalg.py
index 48b3bd6..3536e10 100644
--- a/fmralign/hyperalignment/linalg.py
+++ b/fmralign/hyperalignment/linalg.py
@@ -103,3 +103,75 @@ def ridge(X, Y, alpha):
     d_UT_Y = d[:, np.newaxis] * (U.T @ Y)
     betas = Vt.T @ d_UT_Y
     return betas
+
+
+def procrustes(X, Y, reflection=True, scaling=False):
+    r"""
+    The orthogonal Procrustes algorithm.
+
+    The orthogonal Procrustes algorithm, also known as the classic
+    hyperalignment algorithm, is the first hyperalignment algorithm, which
+    was introduced in Haxby et al. (2011). It tries to align two
+    configurations in a high-dimensional space through an orthogonal
+    transformation. The transformation is an improper rotation, which is a
+    rotation with an optional reflection. Neither rotation nor reflection
+    changes the geometry of the configuration. Therefore, the geometry,
+    such as a representational dissimilarity matrix (RDM), remains the
+    same in the process.
+
+    Optionally, a global scaling can be added to the algorithm to further
+    improve alignment quality. That is, scaling the data with the same
+    factor for all directions. Different from rotation and reflection,
+    global scaling can change the geometry of the data.
+
+    Parameters
+    ----------
+    X : ndarray
+        The data matrix to be aligned to Y.
+    Y : ndarray
+        The "target" data matrix -- the matrix to be aligned to.
+    reflection : bool, default=True
+        Whether allows reflection in the transformation (True) or not
+        (False). Note that even with ``reflection=True``, the solution
+        may not contain a reflection if the alignment cannot be improved
+        by adding a reflection to the rotation.
+    scaling : bool, default=False
+        Whether allows global scaling (True) or not (False). Allowing
+        scaling can improve alignment quality, but it also changes the
+        geometry of data.
+
+    Returns
+    -------
+    T : ndarray
+        The transformation matrix which can be used to align X to Y.
+        Depending on the parameters ``reflection`` and ``scaling``, the
+        transformation can be a pure rotation, an improper rotation, or a
+        pure/improper rotation with global scaling.
+
+    Notes
+    -----
+    The algorithm tries to minimize the Frobenius norm of the difference
+    between transformed data ``X @ T`` and the target ``Y``:
+
+    .. math:: \underset{T}{\arg\min} \lVert XT - Y \rVert_F
+
+    The solution ``T`` differs depending on whether reflection and global
+    scaling are allowed. When it's a rotation (pure or improper), it's
+    often denoted as ``R``.
+    """
+
+    A = Y.T.dot(X).T
+    U, s, Vt = safe_svd(A, remove_mean=False)
+    T = np.dot(U, Vt)
+
+    if not reflection:
+        sign = np.sign(np.linalg.det(T))
+        s[-1] *= sign
+        if sign < 0:
+            T -= np.outer(U[:, -1], Vt[-1, :]) * 2
+
+    if scaling:
+        scale = s.sum() / (X.var(axis=0).sum() * X.shape[0])
+        T *= scale
+
+    return T
diff --git a/fmralign/hyperalignment/local_template.py b/fmralign/hyperalignment/local_template.py
index ec51ea4..3c99fed 100644
--- a/fmralign/hyperalignment/local_template.py
+++ b/fmralign/hyperalignment/local_template.py
@@ -12,7 +12,7 @@
 from tqdm import tqdm
 
 from .linalg import safe_svd
-from .procrustes import procrustes
+from .linalg import procrustes
 
 
 def PCA_decomposition(X, max_npc=None, flavor="sklearn", adjust_ns=False, demean=True):
diff --git a/fmralign/hyperalignment/procrustes.py b/fmralign/hyperalignment/procrustes.py
deleted file mode 100644
index 4a4ca62..0000000
--- a/fmralign/hyperalignment/procrustes.py
+++ /dev/null
@@ -1,74 +0,0 @@
-import numpy as np
-from .linalg import safe_svd
-
-
-def procrustes(X, Y, reflection=True, scaling=False):
-    r"""
-    The orthogonal Procrustes algorithm.
-
-    The orthogonal Procrustes algorithm, also known as the classic
-    hyperalignment algorithm, is the first hyperalignment algorithm, which
-    was introduced in Haxby et al. (2011). It tries to align two
-    configurations in a high-dimensional space through an orthogonal
-    transformation. The transformation is an improper rotation, which is a
-    rotation with an optional reflection. Neither rotation nor reflection
-    changes the geometry of the configuration. Therefore, the geometry,
-    such as a representational dissimilarity matrix (RDM), remains the
-    same in the process.
-
-    Optionally, a global scaling can be added to the algorithm to further
-    improve alignment quality. That is, scaling the data with the same
-    factor for all directions. Different from rotation and reflection,
-    global scaling can change the geometry of the data.
-
-    Parameters
-    ----------
-    X : ndarray
-        The data matrix to be aligned to Y.
-    Y : ndarray
-        The "target" data matrix -- the matrix to be aligned to.
-    reflection : bool, default=True
-        Whether allows reflection in the transformation (True) or not
-        (False). Note that even with ``reflection=True``, the solution
-        may not contain a reflection if the alignment cannot be improved
-        by adding a reflection to the rotation.
-    scaling : bool, default=False
-        Whether allows global scaling (True) or not (False). Allowing
-        scaling can improve alignment quality, but it also changes the
-        geometry of data.
-
-    Returns
-    -------
-    T : ndarray
-        The transformation matrix which can be used to align X to Y.
-        Depending on the parameters ``reflection`` and ``scaling``, the
-        transformation can be a pure rotation, an improper rotation, or a
-        pure/improper rotation with global scaling.
-
-    Notes
-    -----
-    The algorithm tries to minimize the Frobenius norm of the difference
-    between transformed data ``X @ T`` and the target ``Y``:
-
-    .. math:: \underset{T}{\arg\min} \lVert XT - Y \rVert_F
-
-    The solution ``T`` differs depending on whether reflection and global
-    scaling are allowed. When it's a rotation (pure or improper), it's
-    often denoted as ``R``.
-    """
-
-    A = Y.T.dot(X).T
-    U, s, Vt = safe_svd(A, remove_mean=False)
-    T = np.dot(U, Vt)
-
-    if not reflection:
-        sign = np.sign(np.linalg.det(T))
-        s[-1] *= sign
-        if sign < 0:
-            T -= np.outer(U[:, -1], Vt[-1, :]) * 2
-
-    if scaling:
-        scale = s.sum() / (X.var(axis=0).sum() * X.shape[0])
-        T *= scale
-
-    return T
diff --git a/fmralign/hyperalignment/regions.py b/fmralign/hyperalignment/regions.py
index d90069d..d012f54 100644
--- a/fmralign/hyperalignment/regions.py
+++ b/fmralign/hyperalignment/regions.py
@@ -15,7 +15,7 @@
     _safe_get_data,
 )
 from tqdm import tqdm
-from .procrustes import procrustes
+from .linalg import procrustes
 from .linalg import ridge
 from .local_template import compute_template
 from fmralign._utils import _make_parcellation

From 8a14eed7060c494699a599cc547bd1a7e376847d Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Wed, 3 Jan 2024 14:29:06 +0100
Subject: [PATCH 09/69] fix: minor fixes to template

---
 fmralign/hyperalignment/local_template.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/fmralign/hyperalignment/local_template.py b/fmralign/hyperalignment/local_template.py
index 3c99fed..568a91b 100644
--- a/fmralign/hyperalignment/local_template.py
+++ b/fmralign/hyperalignment/local_template.py
@@ -35,9 +35,6 @@ def PCA_decomposition(X, max_npc=None, flavor="sklearn", adjust_ns=False, demean
     """
     ns, nt, nv = X.shape
     X = X.transpose(1, 0, 2).reshape(nt, ns * nv).astype(np.float32)
-    if max_npc is not None:
-        # max_npc = min(max_npc, min(X.shape[0], X.shape[1]))
-        pass
     if flavor == "sklearn":
         try:
             if demean:

From a0295ea55922679ea238601cc6525ce0bf1158d4 Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Wed, 3 Jan 2024 14:31:13 +0100
Subject: [PATCH 10/69] fix: rm dss variable (strange name)

---
 fmralign/hyperalignment/local_template.py | 48 +++++++++++------------
 1 file changed, 24 insertions(+), 24 deletions(-)

diff --git a/fmralign/hyperalignment/local_template.py b/fmralign/hyperalignment/local_template.py
index 568a91b..b5e53f6 100644
--- a/fmralign/hyperalignment/local_template.py
+++ b/fmralign/hyperalignment/local_template.py
@@ -20,7 +20,7 @@ def PCA_decomposition(X, max_npc=None, flavor="sklearn", adjust_ns=False, demean
 
     Parameters
     ----------
-    dss : ndarray of shape (ns, nt, nv)
+    X : ndarray of shape (ns, nt, nv)
     max_npc : integer or None
     flavor : {'sklearn', 'svd'}
     adjust_ns : bool
@@ -101,12 +101,12 @@ def compute_PCA_template(X, sl=None, max_npc=None, flavor="sklearn", demean=Fals
         The PCA template array of shape (n_samples, n_features, n_components).
     """
     if sl is not None:
-        dss = X[:, :, sl]
+        X = X[:, :, sl]
     else:
-        dss = X
-    max_npc = min(dss.shape[1], dss.shape[2])
+        X = X
+    max_npc = min(X.shape[1], X.shape[2])
     XX, cc = PCA_decomposition(
-        dss, max_npc=max_npc, flavor=flavor, adjust_ns=True, demean=demean
+        X, max_npc=max_npc, flavor=flavor, adjust_ns=True, demean=demean
     )
     return XX.astype(np.float32)
 
@@ -118,7 +118,7 @@ def compute_procrustes_template(
     scaling=False,
     zscore_common=True,
     level2_iter=1,
-    dss2=None,
+    X2=None,
     debug=False,
 ):
     """
@@ -131,7 +131,7 @@ def compute_procrustes_template(
         scaling (bool, optional): Whether to allow scaling in the Procrustes alignment. Defaults to False.
         zscore_common (bool, optional): Whether to z-score the aligned data to have zero mean and unit variance. Defaults to True.
         level2_iter (int, optional): The number of iterations for the level 2 alignment. Defaults to 1.
-        dss2 (ndarray or None, optional): The second set of input data array of shape (n_samples, n_features, n_regions). Only used for level 2 alignment. Defaults to None.
+        X2 (ndarray or None, optional): The second set of input data array of shape (n_samples, n_features, n_regions). Only used for level 2 alignment. Defaults to None.
         debug (bool, optional): Whether to display progress bars during alignment. Defaults to False.
 
     Returns:
@@ -141,7 +141,7 @@ def compute_procrustes_template(
     if region is not None:
         X = X[:, :, region]
     common_space = np.copy(X[0])
-    aligned_dss = [X[0]]
+    aligned_X = [X[0]]
     if debug:
         iter = tqdm(X[1:])
         iter.set_description("Computing procrustes alignment (level 1)...")
@@ -152,12 +152,12 @@ def compute_procrustes_template(
         aligned_ds = ds.dot(T)
         if zscore_common:
             aligned_ds = np.nan_to_num(zscore(aligned_ds, axis=0))
-        aligned_dss.append(aligned_ds)
+        aligned_X.append(aligned_ds)
         common_space = (common_space + aligned_ds) * 0.5
         if zscore_common:
             common_space = np.nan_to_num(zscore(common_space, axis=0))
 
-    aligned_dss2 = []
+    aligned_X2 = []
 
     if debug:
         iter2 = tqdm(range(level2_iter))
@@ -167,25 +167,25 @@ def compute_procrustes_template(
 
     for level2 in iter2:
         common_space = np.zeros_like(X[0])
-        for ds in aligned_dss:
+        for ds in aligned_X:
             common_space += ds
         for i, ds in enumerate(X):
-            reference = (common_space - aligned_dss[i]) / float(len(X) - 1)
+            reference = (common_space - aligned_X[i]) / float(len(X) - 1)
             if zscore_common:
                 reference = np.nan_to_num(zscore(reference, axis=0))
             T = procrustes(ds, reference, reflection=reflection, scaling=scaling)
-            if level2 == level2_iter - 1 and dss2 is not None:
-                aligned_dss2.append(dss2[i].dot(T))
-            aligned_dss[i] = ds.dot(T)
+            if level2 == level2_iter - 1 and X2 is not None:
+                aligned_X2.append(X2[i].dot(T))
+            aligned_X[i] = ds.dot(T)
 
-    common_space = np.sum(aligned_dss, axis=0)
+    common_space = np.sum(aligned_X, axis=0)
     if zscore_common:
         common_space = np.nan_to_num(zscore(common_space, axis=0))
     else:
         common_space /= float(len(X))
-    if dss2 is not None:
-        common_space2 = np.zeros_like(dss2[0])
-        for ds in aligned_dss2:
+    if X2 is not None:
+        common_space2 = np.zeros_like(X2[0])
+        for ds in aligned_X2:
             common_space2 += ds
         if zscore_common:
             common_space2 = np.nan_to_num(zscore(common_space2, axis=0))
@@ -210,7 +210,7 @@ def compute_template(
     ----------
     Parameters:
 
-       - dss (ndarray): The input datasets.
+       - X (ndarray): The input datasets.
        - region (ndarray or None): The region indices for searchlight or region-based template computation.
        - sl (ndarray or None): The searchlight indices for searchlight-based template computation.
        - region (int or None, optional): The index of the region to consider. If None, all regions are considered (or searchlights). Defaults to None.
@@ -244,10 +244,10 @@ def compute_template(
 
     if common_topography:
         if region is not None:
-            dss_ = X[:, :, region]
+            X_ = X[:, :, region]
         else:
-            dss_ = np.copy(X)
-        ns, nt, nv = dss_.shape
-        T = procrustes(np.tile(tmpl, (ns, 1)), dss_.reshape(ns * nt, nv))
+            X_ = np.copy(X)
+        ns, nt, nv = X_.shape
+        T = procrustes(np.tile(tmpl, (ns, 1)), X_.reshape(ns * nt, nv))
         tmpl = tmpl @ T
     return tmpl.astype(np.float32)

From 16bde565b73c6c17c9a253db2a75c951d7b38827 Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Wed, 3 Jan 2024 14:36:05 +0100
Subject: [PATCH 11/69] black formating

---
 examples/plot_alignment_methods_benchmark.py | 10 ++++-
 examples/plot_alignment_simulated_2D_data.py | 14 ++++--
 examples/plot_pairwise_roi_alignment.py      |  4 +-
 examples/plot_template_alignment.py          |  4 +-
 fmralign/_utils.py                           | 22 +++++++---
 fmralign/alignment_methods.py                | 12 +++--
 fmralign/fetch_example_data.py               | 20 +++++++--
 fmralign/hyperalignment/hyperalignment.py    |  4 +-
 fmralign/hyperalignment/local_template.py    | 46 +++++++++++---------
 fmralign/hyperalignment/model.py             |  8 +++-
 fmralign/hyperalignment/regions.py           |  4 +-
 fmralign/hyperalignment/regions_alignment.py |  4 +-
 fmralign/metrics.py                          | 12 +++--
 fmralign/pairwise_alignment.py               | 18 +++++---
 fmralign/template_alignment.py               | 12 +++--
 fmralign/tests/test_alignment_methods.py     | 18 ++++++--
 fmralign/tests/test_metrics.py               | 12 +++--
 fmralign/tests/test_pairwise_alignment.py    | 11 ++++-
 fmralign/tests/test_template_alignment.py    | 37 ++++++++++++----
 fmralign/tests/test_utils.py                 |  4 +-
 fmralign/tests/utils.py                      |  7 ++-
 21 files changed, 206 insertions(+), 77 deletions(-)

diff --git a/examples/plot_alignment_methods_benchmark.py b/examples/plot_alignment_methods_benchmark.py
index 246a27b..2fe0fac 100644
--- a/examples/plot_alignment_methods_benchmark.py
+++ b/examples/plot_alignment_methods_benchmark.py
@@ -48,7 +48,9 @@
 # Select visual cortex, create a mask and resample it to the right resolution
 
 mask_visual = new_img_like(atlas, atlas.get_fdata() == 1)
-resampled_mask_visual = resample_to_img(mask_visual, mask, interpolation="nearest")
+resampled_mask_visual = resample_to_img(
+    mask_visual, mask, interpolation="nearest"
+)
 
 # Plot the mask we will use
 plotting.plot_roi(
@@ -147,7 +149,11 @@
     aligned_score = roi_masker.inverse_transform(method_error)
     title = f"Correlation of prediction after {method} alignment"
     display = plotting.plot_stat_map(
-        aligned_score, display_mode="z", cut_coords=[-15, -5], vmax=1, title=title
+        aligned_score,
+        display_mode="z",
+        cut_coords=[-15, -5],
+        vmax=1,
+        title=title,
     )
 
 ###############################################################################
diff --git a/examples/plot_alignment_simulated_2D_data.py b/examples/plot_alignment_simulated_2D_data.py
index ca18ec6..393bc43 100644
--- a/examples/plot_alignment_simulated_2D_data.py
+++ b/examples/plot_alignment_simulated_2D_data.py
@@ -161,7 +161,9 @@ def _plot_distributions_and_alignment(
 Y = np.roll(Y, 6, axis=0)
 # We plot them and observe that their initial matching is wrong
 R_identity = np.eye(n_points, dtype=np.float64)
-_plot_distributions_and_alignment(X, Y, R=R_identity, title="Initial Matching", thr=0.1)
+_plot_distributions_and_alignment(
+    X, Y, R=R_identity, title="Initial Matching", thr=0.1
+)
 
 ###############################################################################
 # Alignment : finding the right transform
@@ -193,7 +195,9 @@ def _plot_distributions_and_alignment(
     title="Procrustes between distributions",
     thr=0.1,
 )
-_plot_mixing_matrix(R=scaled_orthogonal_alignment.R.T, title="Orthogonal mixing matrix")
+_plot_mixing_matrix(
+    R=scaled_orthogonal_alignment.R.T, title="Orthogonal mixing matrix"
+)
 
 ###############################################################################
 # Ridge alignment
@@ -206,7 +210,11 @@ def _plot_distributions_and_alignment(
 ridge_alignment = RidgeAlignment(alphas=[0.01, 0.1], cv=2).fit(X.T, Y.T)
 
 _plot_distributions_and_alignment(
-    X, Y, R=ridge_alignment.R.coef_, title="Ridge between distributions", thr=0.1
+    X,
+    Y,
+    R=ridge_alignment.R.coef_,
+    title="Ridge between distributions",
+    thr=0.1,
 )
 _plot_mixing_matrix(R=ridge_alignment.R.coef_, title="Ridge coefficients")
 
diff --git a/examples/plot_pairwise_roi_alignment.py b/examples/plot_pairwise_roi_alignment.py
index 5c6f365..0baa6ea 100644
--- a/examples/plot_pairwise_roi_alignment.py
+++ b/examples/plot_pairwise_roi_alignment.py
@@ -54,7 +54,9 @@
 # Select visual cortex, create a mask and resample it to the right resolution
 
 mask_visual = new_img_like(atlas, atlas.get_fdata() == 1)
-resampled_mask_visual = resample_to_img(mask_visual, mask, interpolation="nearest")
+resampled_mask_visual = resample_to_img(
+    mask_visual, mask, interpolation="nearest"
+)
 
 # Plot the mask we will use
 plot_roi(
diff --git a/examples/plot_template_alignment.py b/examples/plot_template_alignment.py
index e2458c8..7170aa8 100644
--- a/examples/plot_template_alignment.py
+++ b/examples/plot_template_alignment.py
@@ -155,7 +155,9 @@
     score_voxelwise(target_test, prediction_from_average, masker, loss="corr")
 )
 template_score = masker.inverse_transform(
-    score_voxelwise(target_test, prediction_from_template[0], masker, loss="corr")
+    score_voxelwise(
+        target_test, prediction_from_template[0], masker, loss="corr"
+    )
 )
 
 
diff --git a/fmralign/_utils.py b/fmralign/_utils.py
index eafe96b..9aca728 100644
--- a/fmralign/_utils.py
+++ b/fmralign/_utils.py
@@ -42,7 +42,9 @@ def piecewise_transform(labels, estimators, X):
 
     for i in range(len(unique_labels)):
         label = unique_labels[i]
-        X_transform[:, labels == label] = estimators[i].transform(X[:, labels == label])
+        X_transform[:, labels == label] = estimators[i].transform(
+            X[:, labels == label]
+        )
     return X_transform
 
 
@@ -71,7 +73,13 @@ def _check_labels(labels, threshold=1000):
 
 
 def _make_parcellation(
-    imgs, clustering_index, clustering, n_pieces, masker, smoothing_fwhm=5, verbose=0
+    imgs,
+    clustering_index,
+    clustering,
+    n_pieces,
+    masker,
+    smoothing_fwhm=5,
+    verbose=0,
 ):
     """
     Use nilearn Parcellation class in our pipeline.
@@ -106,7 +114,9 @@ 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) or os.path.isfile(
+        clustering
+    ):
         _check_same_fov(masker.mask_img_, clustering)
         labels = _apply_mask_fmri(clustering, masker.mask_img_).astype(int)
 
@@ -140,9 +150,9 @@ def _make_parcellation(
             )
             err.args += (errmsg,)
             raise err
-        labels = _apply_mask_fmri(parcellation.labels_img_, masker.mask_img_).astype(
-            int
-        )
+        labels = _apply_mask_fmri(
+            parcellation.labels_img_, masker.mask_img_
+        ).astype(int)
 
     if verbose > 0:
         unique_labels, counts = np.unique(labels, return_counts=True)
diff --git a/fmralign/alignment_methods.py b/fmralign/alignment_methods.py
index 4bf4f7e..cdf7866 100644
--- a/fmralign/alignment_methods.py
+++ b/fmralign/alignment_methods.py
@@ -94,7 +94,9 @@ def optimal_permutation(X, Y):
     dist = pairwise_distances(X.T, Y.T)
     u = linear_sum_assignment(dist)
     u = np.array(list(zip(*u)))
-    permutation = scipy.sparse.csr_matrix((np.ones(X.shape[1]), (u[:, 0], u[:, 1]))).T
+    permutation = scipy.sparse.csr_matrix(
+        (np.ones(X.shape[1]), (u[:, 0], u[:, 1]))
+    ).T
     return permutation
 
 
@@ -589,7 +591,9 @@ def fit(
 
         # check for cached data
         try:
-            self.denoised_signal = np.load(self.path + "/train_data_denoised.npy")
+            self.denoised_signal = np.load(
+                self.path + "/train_data_denoised.npy"
+            )
             if verbose:
                 print("Loaded denoised data from cache")
 
@@ -651,7 +655,9 @@ def transform(self, X_test_, verbose=False):
             print("Predict : Computing stimulus matrix...")
 
         if self.decomp_method is None:
-            S = stimulus_estimator(full_signal, self.n_t, self.n_s, self.latent_dim)
+            S = stimulus_estimator(
+                full_signal, self.n_t, self.n_s, self.latent_dim
+            )
 
         if verbose:
             print("Predict : stimulus matrix shape: ", S.shape)
diff --git a/fmralign/fetch_example_data.py b/fmralign/fetch_example_data.py
index b91a78d..6c745b9 100644
--- a/fmralign/fetch_example_data.py
+++ b/fmralign/fetch_example_data.py
@@ -43,9 +43,13 @@ def fetch_ibc_subjects_contrasts(subjects, data_dir=None, verbose=1):
     """
     # The URLs can be retrieved from the nilearn account on OSF
     if subjects == "all":
-        subjects = ["sub-{i:02d}" for i in [1, 2, 4, 5, 6, 7, 8, 9, 11, 12, 13, 14, 15]]
+        subjects = [
+            "sub-{i:02d}" for i in [1, 2, 4, 5, 6, 7, 8, 9, 11, 12, 13, 14, 15]
+        ]
     dataset_name = "ibc"
-    data_dir = _get_dataset_dir(dataset_name, data_dir=data_dir, verbose=verbose)
+    data_dir = _get_dataset_dir(
+        dataset_name, data_dir=data_dir, verbose=verbose
+    )
 
     # download or retrieve metadatas, put it in a dataframe,
     # list all condition and specify path to the right directory
@@ -62,14 +66,22 @@ def fetch_ibc_subjects_contrasts(subjects, data_dir=None, verbose=1):
     )
     metadata_df = pd.read_csv(metadata_path[0])
     conditions = metadata_df.condition.unique()
-    metadata_df["path"] = metadata_df["path"].str.replace("path_to_dir", data_dir)
+    metadata_df["path"] = metadata_df["path"].str.replace(
+        "path_to_dir", data_dir
+    )
     # filter the dataframe to return only rows relevant for subjects argument
     metadata_df = metadata_df[metadata_df.subject.isin(subjects)]
 
     # download / retrieve mask niimg and find its path
     mask = _fetch_files(
         data_dir,
-        [("gm_mask_3mm.nii.gz", "https://osf.io/yvju3/download", {"uncompress": True})],
+        [
+            (
+                "gm_mask_3mm.nii.gz",
+                "https://osf.io/yvju3/download",
+                {"uncompress": True},
+            )
+        ],
         verbose=verbose,
     )[0]
 
diff --git a/fmralign/hyperalignment/hyperalignment.py b/fmralign/hyperalignment/hyperalignment.py
index c5a16cd..7b89a3b 100644
--- a/fmralign/hyperalignment/hyperalignment.py
+++ b/fmralign/hyperalignment/hyperalignment.py
@@ -67,7 +67,9 @@ def fit(
             )
 
         elif self.method == "parcellation":
-            raise NotImplementedError("Parcellation method not implemented yet.")
+            raise NotImplementedError(
+                "Parcellation method not implemented yet."
+            )
 
         if isinstance(imgs, np.ndarray):
             X = imgs
diff --git a/fmralign/hyperalignment/local_template.py b/fmralign/hyperalignment/local_template.py
index b5e53f6..e801849 100644
--- a/fmralign/hyperalignment/local_template.py
+++ b/fmralign/hyperalignment/local_template.py
@@ -15,7 +15,9 @@
 from .linalg import procrustes
 
 
-def PCA_decomposition(X, max_npc=None, flavor="sklearn", adjust_ns=False, demean=True):
+def PCA_decomposition(
+    X, max_npc=None, flavor="sklearn", adjust_ns=False, demean=True
+):
     """Decompose concatenated data matrices using PCA/SVD.
 
     Parameters
@@ -51,7 +53,9 @@ def PCA_decomposition(X, max_npc=None, flavor="sklearn", adjust_ns=False, demean
                     random_state=0,
                 )
                 if adjust_ns:
-                    XX = U[:, :max_npc] * (s[np.newaxis, :max_npc] / np.sqrt(ns))
+                    XX = U[:, :max_npc] * (
+                        s[np.newaxis, :max_npc] / np.sqrt(ns)
+                    )
                 else:
                     XX = U[:, :max_npc] * (s[np.newaxis, :max_npc])
                 cc = Vt[:max_npc].reshape(-1, ns, nv)
@@ -76,7 +80,9 @@ def PCA_decomposition(X, max_npc=None, flavor="sklearn", adjust_ns=False, demean
         raise NotImplementedError
 
 
-def compute_PCA_template(X, sl=None, max_npc=None, flavor="sklearn", demean=False):
+def compute_PCA_template(
+    X, sl=None, max_npc=None, flavor="sklearn", demean=False
+):
     """
     Compute the PCA template from the input data.
 
@@ -97,7 +103,7 @@ def compute_PCA_template(X, sl=None, max_npc=None, flavor="sklearn", demean=Fals
     Returns:
     --------
 
-    ndarray:
+    XX (ndarray):
         The PCA template array of shape (n_samples, n_features, n_components).
     """
     if sl is not None:
@@ -143,17 +149,17 @@ def compute_procrustes_template(
     common_space = np.copy(X[0])
     aligned_X = [X[0]]
     if debug:
-        iter = tqdm(X[1:])
-        iter.set_description("Computing procrustes alignment (level 1)...")
+        iter_X = tqdm(X[1:])
+        iter_X.set_description("Computing procrustes alignment (level 1)...")
     else:
-        iter = X[1:]
-    for ds in iter:
-        T = procrustes(ds, common_space, reflection=reflection, scaling=scaling)
-        aligned_ds = ds.dot(T)
+        iter_X = X[1:]
+    for x in iter_X:
+        T = procrustes(x, common_space, reflection=reflection, scaling=scaling)
+        aligned_x = x.dot(T)
         if zscore_common:
-            aligned_ds = np.nan_to_num(zscore(aligned_ds, axis=0))
-        aligned_X.append(aligned_ds)
-        common_space = (common_space + aligned_ds) * 0.5
+            aligned_x = np.nan_to_num(zscore(aligned_x, axis=0))
+        aligned_X.append(aligned_x)
+        common_space = (common_space + aligned_x) * 0.5
         if zscore_common:
             common_space = np.nan_to_num(zscore(common_space, axis=0))
 
@@ -167,16 +173,16 @@ def compute_procrustes_template(
 
     for level2 in iter2:
         common_space = np.zeros_like(X[0])
-        for ds in aligned_X:
-            common_space += ds
-        for i, ds in enumerate(X):
+        for x in aligned_X:
+            common_space += x
+        for i, x in enumerate(X):
             reference = (common_space - aligned_X[i]) / float(len(X) - 1)
             if zscore_common:
                 reference = np.nan_to_num(zscore(reference, axis=0))
-            T = procrustes(ds, reference, reflection=reflection, scaling=scaling)
+            T = procrustes(x, reference, reflection=reflection, scaling=scaling)
             if level2 == level2_iter - 1 and X2 is not None:
                 aligned_X2.append(X2[i].dot(T))
-            aligned_X[i] = ds.dot(T)
+            aligned_X[i] = x.dot(T)
 
     common_space = np.sum(aligned_X, axis=0)
     if zscore_common:
@@ -185,8 +191,8 @@ def compute_procrustes_template(
         common_space /= float(len(X))
     if X2 is not None:
         common_space2 = np.zeros_like(X2[0])
-        for ds in aligned_X2:
-            common_space2 += ds
+        for x in aligned_X2:
+            common_space2 += x
         if zscore_common:
             common_space2 = np.nan_to_num(zscore(common_space2, axis=0))
         else:
diff --git a/fmralign/hyperalignment/model.py b/fmralign/hyperalignment/model.py
index 15fe165..8adec93 100644
--- a/fmralign/hyperalignment/model.py
+++ b/fmralign/hyperalignment/model.py
@@ -126,7 +126,9 @@ def fit(
 
         # check for cached data
         try:
-            self.denoised_signal = np.load(self.path + "/train_data_denoised.npy")
+            self.denoised_signal = np.load(
+                self.path + "/train_data_denoised.npy"
+            )
             if verbose:
                 print("Loaded denoised data from cache")
 
@@ -191,7 +193,9 @@ def transform(self, X_test_, verbose=False):
             print("Predict : Computing stimulus matrix...")
 
         if self.decomp_method is None:
-            S = stimulus_estimator(full_signal, self.n_t, self.n_s, self.latent_dim)
+            S = stimulus_estimator(
+                full_signal, self.n_t, self.n_s, self.latent_dim
+            )
 
         if verbose:
             print("Predict : stimulus matrix shape: ", S.shape)
diff --git a/fmralign/hyperalignment/regions.py b/fmralign/hyperalignment/regions.py
index d012f54..5d579b7 100644
--- a/fmralign/hyperalignment/regions.py
+++ b/fmralign/hyperalignment/regions.py
@@ -418,7 +418,9 @@ def piece_procrustes(
         array-like: The transformation matrix T.
 
     """
-    sl_func = functools.partial(procrustes, reflection=reflection, scaling=scaling)
+    sl_func = functools.partial(
+        procrustes, reflection=reflection, scaling=scaling
+    )
     T = iter_hyperalignment(
         X,
         Y,
diff --git a/fmralign/hyperalignment/regions_alignment.py b/fmralign/hyperalignment/regions_alignment.py
index 4ba62bf..082b0a5 100644
--- a/fmralign/hyperalignment/regions_alignment.py
+++ b/fmralign/hyperalignment/regions_alignment.py
@@ -59,7 +59,9 @@ def __init__(
             if not os.path.exists(self.path):
                 os.makedirs(self.path)
 
-    def compute_linear_transformation(self, x_i, template, i: int = 0, save=True):
+    def compute_linear_transformation(
+        self, x_i, template, i: int = 0, save=True
+    ):
         """Compute the linear transformation W_i for a given subject.
         ----------
         Parameters
diff --git a/fmralign/metrics.py b/fmralign/metrics.py
index cc2d29e..d9b330c 100644
--- a/fmralign/metrics.py
+++ b/fmralign/metrics.py
@@ -4,7 +4,9 @@
 from sklearn.metrics import r2_score
 
 
-def score_voxelwise(ground_truth, prediction, masker, loss, multioutput="raw_values"):
+def score_voxelwise(
+    ground_truth, prediction, masker, loss, multioutput="raw_values"
+):
     """
     Calculate loss function for predicted, ground truth arrays.
     Supported scores are R2, correlation, and normalized
@@ -51,7 +53,9 @@ def score_voxelwise(ground_truth, prediction, masker, loss, multioutput="raw_val
     if loss == "R2":
         score = r2_score(X_gt, X_pred, multioutput=multioutput)
     elif loss == "n_reconstruction_err":
-        score = normalized_reconstruction_error(X_gt, X_pred, multioutput=multioutput)
+        score = normalized_reconstruction_error(
+            X_gt, X_pred, multioutput=multioutput
+        )
     elif loss == "corr":
         score = np.array(
             [
@@ -124,7 +128,9 @@ def normalized_reconstruction_error(
 
     # Calculate reconstruction error
     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]
+    )
     if multioutput == "raw_values":
         # return scores individually
         return output_scores
diff --git a/fmralign/pairwise_alignment.py b/fmralign/pairwise_alignment.py
index 9753d87..68ff222 100644
--- a/fmralign/pairwise_alignment.py
+++ b/fmralign/pairwise_alignment.py
@@ -53,7 +53,9 @@ def generate_Xi_Yi(labels, X, Y, masker, verbose):
         label = unique_labels[k]
         i = label == labels
         if (k + 1) % 25 == 0 and verbose > 0:
-            print("Fitting parcel: " + str(k + 1) + "/" + str(len(unique_labels)))
+            print(
+                "Fitting parcel: " + str(k + 1) + "/" + str(len(unique_labels))
+            )
         # should return X_i Y_i
         yield X_[:, i], Y_[:, i]
 
@@ -321,9 +323,9 @@ def fit(self, X, Y):
         else:
             self.masker_.fit()
 
-        if isinstance(self.clustering, nib.nifti1.Nifti1Image) or os.path.isfile(
-            self.clustering
-        ):
+        if isinstance(
+            self.clustering, nib.nifti1.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(
@@ -350,7 +352,9 @@ def fit(self, X, Y):
         self.fit_, self.labels_ = [], []
         rs = ShuffleSplit(n_splits=self.n_bags, test_size=0.8, random_state=0)
 
-        outputs = Parallel(n_jobs=self.n_jobs, prefer="threads", verbose=self.verbose)(
+        outputs = Parallel(
+            n_jobs=self.n_jobs, prefer="threads", verbose=self.verbose
+        )(
             delayed(fit_one_parcellation)(
                 X_,
                 Y_,
@@ -389,7 +393,9 @@ def transform(self, X):
 
         X_transform = np.zeros_like(X_)
         for i in range(self.n_bags):
-            X_transform += piecewise_transform(self.labels_[i], self.fit_[i], X_)
+            X_transform += piecewise_transform(
+                self.labels_[i], self.fit_[i], X_
+            )
 
         X_transform /= self.n_bags
 
diff --git a/fmralign/template_alignment.py b/fmralign/template_alignment.py
index 0730ca4..0edd642 100644
--- a/fmralign/template_alignment.py
+++ b/fmralign/template_alignment.py
@@ -132,7 +132,9 @@ def _create_template(
     aligned_imgs = imgs
     template_history = []
     for iter in range(n_iter):
-        template = _rescaled_euclidean_mean(aligned_imgs, masker, scale_template)
+        template = _rescaled_euclidean_mean(
+            aligned_imgs, masker, scale_template
+        )
         if 0 < iter < n_iter - 1:
             template_history.append(template)
         aligned_imgs = _align_images_to_template(
@@ -477,7 +479,9 @@ def transform(self, imgs, train_index, test_index):
                 "greater index in train_index or test_index."
             )
 
-        fitted_mappings = Parallel(self.n_jobs, prefer="threads", verbose=self.verbose)(
+        fitted_mappings = Parallel(
+            self.n_jobs, prefer="threads", verbose=self.verbose
+        )(
             delayed(_map_template_to_image)(
                 img,
                 train_index,
@@ -495,7 +499,9 @@ def transform(self, imgs, train_index, test_index):
             for img in imgs
         )
 
-        predicted_imgs = Parallel(self.n_jobs, prefer="threads", verbose=self.verbose)(
+        predicted_imgs = Parallel(
+            self.n_jobs, prefer="threads", verbose=self.verbose
+        )(
             delayed(_predict_from_template_and_mapping)(
                 self.template, test_index, mapping
             )
diff --git a/fmralign/tests/test_alignment_methods.py b/fmralign/tests/test_alignment_methods.py
index 4f9aa2a..2effde9 100644
--- a/fmralign/tests/test_alignment_methods.py
+++ b/fmralign/tests/test_alignment_methods.py
@@ -94,19 +94,27 @@ def test_scaled_procrustes_on_simple_exact_cases():
     assert_array_almost_equal(R_test.T, R)
 
     """Scaled Matrix"""
-    X = np.array([[1.0, 2.0, 3.0, 4.0], [5.0, 3.0, 4.0, 6.0], [7.0, 8.0, -5.0, -2.0]])
+    X = np.array(
+        [[1.0, 2.0, 3.0, 4.0], [5.0, 3.0, 4.0, 6.0], [7.0, 8.0, -5.0, -2.0]]
+    )
 
     X = X - X.mean(axis=1, keepdims=True)
 
     Y = 2 * X
     Y = Y - Y.mean(axis=1, keepdims=True)
 
-    assert_array_almost_equal(scaled_procrustes(X.T, Y.T, scaling=True)[0], np.eye(3))
+    assert_array_almost_equal(
+        scaled_procrustes(X.T, Y.T, scaling=True)[0], np.eye(3)
+    )
     assert_array_almost_equal(scaled_procrustes(X.T, Y.T, scaling=True)[1], 2)
 
     """3D Rotation"""
     R = np.array(
-        [[1.0, 0.0, 0.0], [0.0, np.cos(1), -np.sin(1)], [0.0, np.sin(1), np.cos(1)]]
+        [
+            [1.0, 0.0, 0.0],
+            [0.0, np.cos(1), -np.sin(1)],
+            [0.0, np.sin(1), np.cos(1)],
+        ]
     )
     X = np.random.rand(3, 4)
     X = X - X.mean(axis=1, keepdims=True)
@@ -209,7 +217,9 @@ def test_ott_backend():
     algo = OptimalTransportAlignment(
         reg=epsilon, metric="euclidean", tol=1e-5, max_iter=10000
     )
-    old_implem = POTAlignment(reg=epsilon, metric="euclidean", tol=1e-5, max_iter=10000)
+    old_implem = POTAlignment(
+        reg=epsilon, metric="euclidean", tol=1e-5, max_iter=10000
+    )
     algo.fit(X, Y)
     old_implem.fit(X, Y)
     assert_array_almost_equal(algo.R, old_implem.R, decimal=3)
diff --git a/fmralign/tests/test_metrics.py b/fmralign/tests/test_metrics.py
index fc72f74..62ebcef 100644
--- a/fmralign/tests/test_metrics.py
+++ b/fmralign/tests/test_metrics.py
@@ -7,8 +7,12 @@
 
 
 def test_score_voxelwise():
-    A = np.asarray([[[[1, 1.2, 1, 1.2, 1]], [[1, 1, 1, 0.2, 1]], [[1, -1, 1, -1, 1]]]])
-    B = np.asarray([[[[0, 0.2, 0, 0.2, 0]], [[0.2, 1, 1, 1, 1]], [[-1, 1, -1, 1, -1]]]])
+    A = np.asarray(
+        [[[[1, 1.2, 1, 1.2, 1]], [[1, 1, 1, 0.2, 1]], [[1, -1, 1, -1, 1]]]]
+    )
+    B = np.asarray(
+        [[[[0, 0.2, 0, 0.2, 0]], [[0.2, 1, 1, 1, 1]], [[-1, 1, -1, 1, -1]]]]
+    )
     im_A = nib.Nifti1Image(A, np.eye(4))
     im_B = nib.Nifti1Image(B, np.eye(4))
     mask_img = nib.Nifti1Image(np.ones(im_A.shape[0:3]), np.eye(4))
@@ -29,7 +33,9 @@ def test_score_voxelwise():
     assert_array_almost_equal(r2, [-1.0, -1.0, -1.0])
 
     # check normalized reconstruction
-    norm_rec = metrics.score_voxelwise(im_A, im_B, masker, loss="n_reconstruction_err")
+    norm_rec = metrics.score_voxelwise(
+        im_A, im_B, masker, loss="n_reconstruction_err"
+    )
     assert_array_almost_equal(norm_rec, [0.14966, 0.683168, -1.0])
 
 
diff --git a/fmralign/tests/test_pairwise_alignment.py b/fmralign/tests/test_pairwise_alignment.py
index 1420bf7..5698d7d 100644
--- a/fmralign/tests/test_pairwise_alignment.py
+++ b/fmralign/tests/test_pairwise_alignment.py
@@ -29,7 +29,12 @@ def test_pairwise_identity():
     args_list = [
         {"alignment_method": "identity", "mask": mask_img},
         {"alignment_method": "identity", "n_pieces": 3, "mask": mask_img},
-        {"alignment_method": "identity", "n_pieces": 3, "n_bags": 4, "mask": mask_img},
+        {
+            "alignment_method": "identity",
+            "n_pieces": 3,
+            "n_bags": 4,
+            "mask": mask_img,
+        },
         {
             "alignment_method": "identity",
             "n_pieces": 3,
@@ -58,7 +63,9 @@ def test_pairwise_identity():
     )
     with pytest.warns(UserWarning):
         algo.fit(img1, img1)
-    assert (algo.mask.get_fdata() > 0).sum() == (clustering.get_fdata() > 0).sum()
+    assert (algo.mask.get_fdata() > 0).sum() == (
+        clustering.get_fdata() > 0
+    ).sum()
 
     # test warning raised if parcel is 0 :
     null_im = new_img_like(img1, np.zeros_like(img1.get_fdata()))
diff --git a/fmralign/tests/test_template_alignment.py b/fmralign/tests/test_template_alignment.py
index 8798b4a..f04b150 100644
--- a/fmralign/tests/test_template_alignment.py
+++ b/fmralign/tests/test_template_alignment.py
@@ -4,8 +4,14 @@
 from nilearn.maskers import NiftiMasker
 from numpy.testing import assert_array_almost_equal
 
-from fmralign.template_alignment import TemplateAlignment, _rescaled_euclidean_mean
-from fmralign.tests.utils import random_niimg, zero_mean_coefficient_determination
+from fmralign.template_alignment import (
+    TemplateAlignment,
+    _rescaled_euclidean_mean,
+)
+from fmralign.tests.utils import (
+    random_niimg,
+    zero_mean_coefficient_determination,
+)
 
 
 def test_template_identity():
@@ -24,7 +30,9 @@ def test_template_identity():
 
     # test euclidian mean function
     euclidian_template = _rescaled_euclidean_mean(subs, masker)
-    assert_array_almost_equal(ref_template.get_fdata(), euclidian_template.get_fdata())
+    assert_array_almost_equal(
+        ref_template.get_fdata(), euclidian_template.get_fdata()
+    )
 
     # test different fit() accept list of list of 3D Niimgs as input.
     algo = TemplateAlignment(alignment_method="identity", mask=masker)
@@ -37,7 +45,12 @@ def test_template_identity():
         {"alignment_method": "identity", "mask": masker},
         {"alignment_method": "identity", "mask": masker, "n_jobs": 2},
         {"alignment_method": "identity", "n_pieces": 3, "mask": masker},
-        {"alignment_method": "identity", "n_pieces": 3, "n_bags": 2, "mask": masker},
+        {
+            "alignment_method": "identity",
+            "n_pieces": 3,
+            "n_bags": 2,
+            "mask": masker,
+        },
     ]
 
     for args in args_list:
@@ -45,9 +58,13 @@ def test_template_identity():
         # Learning a template which is
         algo.fit(subs)
         # test template
-        assert_array_almost_equal(ref_template.get_fdata(), algo.template.get_fdata())
+        assert_array_almost_equal(
+            ref_template.get_fdata(), algo.template.get_fdata()
+        )
         predicted_imgs = algo.transform(
-            [index_img(sub_1, range(8))], train_index=range(8), test_index=range(8, 10)
+            [index_img(sub_1, range(8))],
+            train_index=range(8),
+            test_index=range(8, 10),
         )
         ground_truth = index_img(ref_template, range(8, 10))
         assert_array_almost_equal(
@@ -65,7 +82,9 @@ def test_template_identity():
     for train_ind, test_ind in zip(train_inds, test_inds):
         with pytest.raises(Exception):
             assert algo.transform(
-                [index_img(sub_1, range(2))], train_index=train_ind, test_index=test_ind
+                [index_img(sub_1, range(2))],
+                train_index=train_ind,
+                test_index=test_ind,
             )
 
     # test wrong images input in fit() and transform method
@@ -116,4 +135,6 @@ def test_template_closer_to_target():
             avg_data, template_data
         )
         assert template_mean_distance >= mean_distance_1
-        assert template_mean_distance >= mean_distance_2 - 1.0e-3  # for robustness
+        assert (
+            template_mean_distance >= mean_distance_2 - 1.0e-3
+        )  # for robustness
diff --git a/fmralign/tests/test_utils.py b/fmralign/tests/test_utils.py
index d82359a..0b70802 100644
--- a/fmralign/tests/test_utils.py
+++ b/fmralign/tests/test_utils.py
@@ -31,7 +31,9 @@ def test_make_parcellation():
 
         # check that not inputing n_pieces yields problems
         with pytest.raises(Exception):
-            assert _make_parcellation(img, indexes, clustering_method, 0, masker)
+            assert _make_parcellation(
+                img, indexes, clustering_method, 0, masker
+            )
 
     clustering = nibabel.Nifti1Image(
         np.hstack([np.ones((7, 3, 8)), 2 * np.ones((7, 3, 8))]), np.eye(4)
diff --git a/fmralign/tests/utils.py b/fmralign/tests/utils.py
index e2f9fc5..b30dbf2 100644
--- a/fmralign/tests/utils.py
+++ b/fmralign/tests/utils.py
@@ -41,7 +41,9 @@ 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":
@@ -52,7 +54,8 @@ 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):

From 3d685767eb53292ef6fa6065e4657dcb74314ef6 Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Wed, 3 Jan 2024 14:41:25 +0100
Subject: [PATCH 12/69] fix: memory issues

---
 fmralign/hyperalignment/linalg.py         |  9 ++++++---
 fmralign/hyperalignment/local_template.py | 12 +++---------
 fmralign/hyperalignment/regions.py        |  4 +---
 3 files changed, 10 insertions(+), 15 deletions(-)

diff --git a/fmralign/hyperalignment/linalg.py b/fmralign/hyperalignment/linalg.py
index 3536e10..71c30ed 100644
--- a/fmralign/hyperalignment/linalg.py
+++ b/fmralign/hyperalignment/linalg.py
@@ -43,11 +43,14 @@ def safe_svd(X, remove_mean=True):
         Unitary matrix.
     """
     if remove_mean:
-        X = X - X.mean(axis=0, keepdims=True)
+        X_ = X - X.mean(axis=0, keepdims=True)
+    else:
+        X_ = X.copy()
     try:
-        U, s, Vt = svd(X, full_matrices=False)
+        U, s, Vt = svd(X_, full_matrices=False)
     except LinAlgError:
-        U, s, Vt = svd(X, full_matrices=False, lapack_driver="gesvd")
+        U, s, Vt = svd(X_, full_matrices=False, lapack_driver="gesvd")
+    del X_
     return U, s, Vt
 
 
diff --git a/fmralign/hyperalignment/local_template.py b/fmralign/hyperalignment/local_template.py
index e801849..a5da094 100644
--- a/fmralign/hyperalignment/local_template.py
+++ b/fmralign/hyperalignment/local_template.py
@@ -15,9 +15,7 @@
 from .linalg import procrustes
 
 
-def PCA_decomposition(
-    X, max_npc=None, flavor="sklearn", adjust_ns=False, demean=True
-):
+def PCA_decomposition(X, max_npc=None, flavor="sklearn", adjust_ns=False, demean=True):
     """Decompose concatenated data matrices using PCA/SVD.
 
     Parameters
@@ -53,9 +51,7 @@ def PCA_decomposition(
                     random_state=0,
                 )
                 if adjust_ns:
-                    XX = U[:, :max_npc] * (
-                        s[np.newaxis, :max_npc] / np.sqrt(ns)
-                    )
+                    XX = U[:, :max_npc] * (s[np.newaxis, :max_npc] / np.sqrt(ns))
                 else:
                     XX = U[:, :max_npc] * (s[np.newaxis, :max_npc])
                 cc = Vt[:max_npc].reshape(-1, ns, nv)
@@ -80,9 +76,7 @@ def PCA_decomposition(
         raise NotImplementedError
 
 
-def compute_PCA_template(
-    X, sl=None, max_npc=None, flavor="sklearn", demean=False
-):
+def compute_PCA_template(X, sl=None, max_npc=None, flavor="sklearn", demean=False):
     """
     Compute the PCA template from the input data.
 
diff --git a/fmralign/hyperalignment/regions.py b/fmralign/hyperalignment/regions.py
index 5d579b7..d012f54 100644
--- a/fmralign/hyperalignment/regions.py
+++ b/fmralign/hyperalignment/regions.py
@@ -418,9 +418,7 @@ def piece_procrustes(
         array-like: The transformation matrix T.
 
     """
-    sl_func = functools.partial(
-        procrustes, reflection=reflection, scaling=scaling
-    )
+    sl_func = functools.partial(procrustes, reflection=reflection, scaling=scaling)
     T = iter_hyperalignment(
         X,
         Y,

From b25f9057e61bd7c693856b9699fe2630d4ade9df Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Wed, 3 Jan 2024 14:47:15 +0100
Subject: [PATCH 13/69] rm all refs to the hyperalignment package

---
 fmralign/hyperalignment/hyperalignment.py |  4 +---
 fmralign/template_alignment.py            | 14 ++++----------
 2 files changed, 5 insertions(+), 13 deletions(-)

diff --git a/fmralign/hyperalignment/hyperalignment.py b/fmralign/hyperalignment/hyperalignment.py
index 7b89a3b..c5a16cd 100644
--- a/fmralign/hyperalignment/hyperalignment.py
+++ b/fmralign/hyperalignment/hyperalignment.py
@@ -67,9 +67,7 @@ def fit(
             )
 
         elif self.method == "parcellation":
-            raise NotImplementedError(
-                "Parcellation method not implemented yet."
-            )
+            raise NotImplementedError("Parcellation method not implemented yet.")
 
         if isinstance(imgs, np.ndarray):
             X = imgs
diff --git a/fmralign/template_alignment.py b/fmralign/template_alignment.py
index 0edd642..f9d6e1b 100644
--- a/fmralign/template_alignment.py
+++ b/fmralign/template_alignment.py
@@ -132,9 +132,7 @@ def _create_template(
     aligned_imgs = imgs
     template_history = []
     for iter in range(n_iter):
-        template = _rescaled_euclidean_mean(
-            aligned_imgs, masker, scale_template
-        )
+        template = _rescaled_euclidean_mean(aligned_imgs, masker, scale_template)
         if 0 < iter < n_iter - 1:
             template_history.append(template)
         aligned_imgs = _align_images_to_template(
@@ -266,7 +264,7 @@ def __init__(
         alignment_method: string
             Algorithm used to perform alignment between X_i and Y_i :
             * either 'identity', 'scaled_orthogonal', 'optimal_transport',
-            'ridge_cv', 'permutation', 'diagonal', 'pcha', 'slha'
+            'ridge_cv', 'permutation', 'diagonal',
             * or an instance of one of alignment classes (imported from
             functional_alignment.alignment_methods)
         n_pieces: int, optional (default = 1)
@@ -479,9 +477,7 @@ def transform(self, imgs, train_index, test_index):
                 "greater index in train_index or test_index."
             )
 
-        fitted_mappings = Parallel(
-            self.n_jobs, prefer="threads", verbose=self.verbose
-        )(
+        fitted_mappings = Parallel(self.n_jobs, prefer="threads", verbose=self.verbose)(
             delayed(_map_template_to_image)(
                 img,
                 train_index,
@@ -499,9 +495,7 @@ def transform(self, imgs, train_index, test_index):
             for img in imgs
         )
 
-        predicted_imgs = Parallel(
-            self.n_jobs, prefer="threads", verbose=self.verbose
-        )(
+        predicted_imgs = Parallel(self.n_jobs, prefer="threads", verbose=self.verbose)(
             delayed(_predict_from_template_and_mapping)(
                 self.template, test_index, mapping
             )

From f034a552efdca824c0d0acd8c9b02f241a56ceda Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Thu, 4 Jan 2024 10:53:54 +0100
Subject: [PATCH 14/69] PR adjustments

---
 fmralign/alignment_methods.py             | 45 ++++--------------
 fmralign/hyperalignment/hyperalignment.py |  6 +--
 fmralign/hyperalignment/local_template.py | 31 +++++++-----
 fmralign/hyperalignment/model.py          | 58 ++++++++---------------
 4 files changed, 50 insertions(+), 90 deletions(-)

diff --git a/fmralign/alignment_methods.py b/fmralign/alignment_methods.py
index cdf7866..9ae12c1 100644
--- a/fmralign/alignment_methods.py
+++ b/fmralign/alignment_methods.py
@@ -94,9 +94,7 @@ def optimal_permutation(X, Y):
     dist = pairwise_distances(X.T, Y.T)
     u = linear_sum_assignment(dist)
     u = np.array(list(zip(*u)))
-    permutation = scipy.sparse.csr_matrix(
-        (np.ones(X.shape[1]), (u[:, 0], u[:, 1]))
-    ).T
+    permutation = scipy.sparse.csr_matrix((np.ones(X.shape[1]), (u[:, 0], u[:, 1]))).T
     return permutation
 
 
@@ -544,7 +542,7 @@ def fit(
         searchlights=None,
         parcels=None,
         dists=None,
-        radius: int = 20,
+        radius=20,
         verbose=True,
     ):
         """
@@ -591,9 +589,7 @@ def fit(
 
         # check for cached data
         try:
-            self.denoised_signal = np.load(
-                self.path + "/train_data_denoised.npy"
-            )
+            self.denoised_signal = np.load(self.path + "/train_data_denoised.npy")
             if verbose:
                 print("Loaded denoised data from cache")
 
@@ -613,8 +609,6 @@ def fit(
             # Clear memory of the SearchlightAlignment object
             denoiser = None
 
-        iterm = range(self.n_s)
-
         # Stimulus matrix computation
         if self.decomp_method is None:
             full_signal = np.concatenate(self.denoised_signal, axis=1)
@@ -628,7 +622,7 @@ def fit(
                     self.denoised_signal[i],
                     latent_dim=self.latent_dim,
                 )
-                for i in iterm
+                for i in range(self.n_s)
             )
 
         return self
@@ -655,9 +649,7 @@ def transform(self, X_test_, verbose=False):
             print("Predict : Computing stimulus matrix...")
 
         if self.decomp_method is None:
-            S = stimulus_estimator(
-                full_signal, self.n_t, self.n_s, self.latent_dim
-            )
+            S = stimulus_estimator(full_signal, self.n_t, self.n_s, self.latent_dim)
 
         if verbose:
             print("Predict : stimulus matrix shape: ", S.shape)
@@ -667,24 +659,6 @@ def transform(self, X_test_, verbose=False):
         ]
         return np.array(reconstructed_signal, dtype=np.float32)
 
-    def get_shared_stimulus(self):
-        """
-        Returns the shared stimulus used for individualized neural tuning.
-
-        Returns:
-            The shared stimulus of shape (n_t, latent_dim) or (n_t, n_t).
-        """
-        return self.shared_response
-
-    def get_tuning_matrices(self):
-        """
-        Returns the tuning matrices as a NumPy array.
-
-        Returns:
-            numpy.ndarray: The tuning matrices of shape (n_s, latent_dim, n_v) or (n_s, n_t, n_v).
-        """
-        return np.array(self.tuning_data)
-
     def clean_cache(self, id):
         """
         Removes the cache file associated with the given ID.
@@ -720,8 +694,6 @@ def tuning_estimator(shared_response, target, latent_dim=None):
         array-like: The estimated tuning weights.
 
     """
-    if latent_dim is None:
-        return np.linalg.inv(shared_response).dot(target)
     return np.linalg.pinv(shared_response).dot(target).astype(np.float32)
 
 
@@ -736,7 +708,7 @@ def stimulus_estimator(full_signal, n_t, n_s, latent_dim=None):
         latent_dim (int, optional): The number of latent dimensions. Defaults to None.
 
     Returns:
-        np.ndarray: The estimated shared response.
+        stimulus (np.ndarray): The stimulus response of shape (n_t, latent_dim) or (n_t, n_t).
     """
     if latent_dim is not None and latent_dim < n_t:
         U = svd_pca(full_signal)
@@ -744,8 +716,9 @@ def stimulus_estimator(full_signal, n_t, n_s, latent_dim=None):
     else:
         U, _, _ = safe_svd(full_signal)
 
-    shared_response = np.sqrt(n_s) * U
-    return shared_response.astype(np.float32)
+    stimulus = np.sqrt(n_s) * U
+    stimulus = stimulus.astype(np.float32)
+    return stimulus
 
 
 def reconstruct_signal(shared_response, individual_tuning):
diff --git a/fmralign/hyperalignment/hyperalignment.py b/fmralign/hyperalignment/hyperalignment.py
index c5a16cd..d941776 100644
--- a/fmralign/hyperalignment/hyperalignment.py
+++ b/fmralign/hyperalignment/hyperalignment.py
@@ -25,7 +25,6 @@ def __init__(self, method="searchlight", n_jobs=-1):
         super().__init__(n_jobs=n_jobs)
         self.mask_img = None
         self.masker = None
-        self.preds = []
         self.method = method
 
     def fit(
@@ -100,7 +99,8 @@ def transform(self, imgs, y=None, verbose=0):
 
         X = self.masker.fit_transform(imgs)
         Y = super().transform(X, verbose=verbose)
+        preds = []
         for y in Y:
-            self.preds.append(self.masker.inverse_transform(y))
+            preds.append(self.masker.inverse_transform(y))
 
-        return self.preds
+        return preds
diff --git a/fmralign/hyperalignment/local_template.py b/fmralign/hyperalignment/local_template.py
index a5da094..c16a92c 100644
--- a/fmralign/hyperalignment/local_template.py
+++ b/fmralign/hyperalignment/local_template.py
@@ -15,14 +15,19 @@
 from .linalg import procrustes
 
 
-def PCA_decomposition(X, max_npc=None, flavor="sklearn", adjust_ns=False, demean=True):
+def PCA_decomposition(
+    X, n_components=None, flavor="sklearn", adjust_ns=False, demean=True
+):
     """Decompose concatenated data matrices using PCA/SVD.
 
     Parameters
     ----------
     X : ndarray of shape (ns, nt, nv)
-    max_npc : integer or None
+        The input data array.
+    n_components : int or None
+        The number of components to keep. If None, all components are kept.
     flavor : {'sklearn', 'svd'}
+        Wethter to use sklearn or the custom SVD implementation.
     adjust_ns : bool
         Whether to adjust the variance of the output so that it doesn't increase with the number of subjects.
     demean : bool
@@ -38,7 +43,7 @@ def PCA_decomposition(X, max_npc=None, flavor="sklearn", adjust_ns=False, demean
     if flavor == "sklearn":
         try:
             if demean:
-                pca = PCA(n_components=max_npc, random_state=0)
+                pca = PCA(n_components=n_components, random_state=0)
                 XX = pca.fit_transform(X)
                 cc = pca.components_.reshape(-1, ns, nv)
                 if adjust_ns:
@@ -47,19 +52,21 @@ def PCA_decomposition(X, max_npc=None, flavor="sklearn", adjust_ns=False, demean
             else:
                 U, s, Vt = randomized_svd(
                     X,
-                    (max_npc if max_npc is not None else min(X.shape)),
+                    (n_components if n_components is not None else min(X.shape)),
                     random_state=0,
                 )
                 if adjust_ns:
-                    XX = U[:, :max_npc] * (s[np.newaxis, :max_npc] / np.sqrt(ns))
+                    XX = U[:, :n_components] * (
+                        s[np.newaxis, :n_components] / np.sqrt(ns)
+                    )
                 else:
-                    XX = U[:, :max_npc] * (s[np.newaxis, :max_npc])
-                cc = Vt[:max_npc].reshape(-1, ns, nv)
+                    XX = U[:, :n_components] * (s[np.newaxis, :n_components])
+                cc = Vt[:n_components].reshape(-1, ns, nv)
                 return XX.astype(np.float32), cc
         except:  # noqa: E722
             return PCA_decomposition(
                 X,
-                max_npc=max_npc,
+                n_components=n_components,
                 flavor="svd",
                 adjust_ns=adjust_ns,
                 demean=demean,
@@ -67,10 +74,10 @@ def PCA_decomposition(X, max_npc=None, flavor="sklearn", adjust_ns=False, demean
     elif flavor == "svd":
         U, s, Vt = safe_svd(X)
         if adjust_ns:
-            XX = U[:, :max_npc] * (s[np.newaxis, :max_npc] / np.sqrt(ns))
+            XX = U[:, :n_components] * (s[np.newaxis, :n_components] / np.sqrt(ns))
         else:
-            XX = U[:, :max_npc] * (s[np.newaxis, :max_npc])
-        cc = Vt[:max_npc].reshape(-1, ns, nv)
+            XX = U[:, :n_components] * (s[np.newaxis, :n_components])
+        cc = Vt[:n_components].reshape(-1, ns, nv)
         return XX.astype(np.float32), cc
     else:
         raise NotImplementedError
@@ -106,7 +113,7 @@ def compute_PCA_template(X, sl=None, max_npc=None, flavor="sklearn", demean=Fals
         X = X
     max_npc = min(X.shape[1], X.shape[2])
     XX, cc = PCA_decomposition(
-        X, max_npc=max_npc, flavor=flavor, adjust_ns=True, demean=demean
+        X, n_components=max_npc, flavor=flavor, adjust_ns=True, demean=demean
     )
     return XX.astype(np.float32)
 
diff --git a/fmralign/hyperalignment/model.py b/fmralign/hyperalignment/model.py
index 8adec93..c7c08a4 100644
--- a/fmralign/hyperalignment/model.py
+++ b/fmralign/hyperalignment/model.py
@@ -34,9 +34,9 @@ def __init__(
         None
         """
 
-        self.n_s = None
-        self.n_t = None
-        self.n_v = None
+        self.n_subjects = None
+        self.n_timepoints = None
+        self.n_voxels = None
         self.labels = None
         self.alphas = None
         self.alignment_method = alignment_method
@@ -112,10 +112,10 @@ def fit(
 
         X_train_ = np.array(X_train, copy=True, dtype=np.float32)
 
-        self.n_s, self.n_t, self.n_v = X_train_.shape
+        self.n_subjects, self.n_timepoints, self.n_voxels = X_train_.shape
 
-        self.tuning_data = np.empty(self.n_s, dtype=np.float32)
-        self.denoised_signal = np.empty(self.n_s, dtype=np.float32)
+        self.tuning_data = np.empty(self.n_subjects, dtype=np.float32)
+        self.denoised_signal = np.empty(self.n_subjects, dtype=np.float32)
 
         if searchlights is None:
             self.regions = parcels
@@ -126,9 +126,7 @@ def fit(
 
         # check for cached data
         try:
-            self.denoised_signal = np.load(
-                self.path + "/train_data_denoised.npy"
-            )
+            self.denoised_signal = np.load(self.path + "/train_data_denoised.npy")
             if verbose:
                 print("Loaded denoised data from cache")
 
@@ -148,7 +146,7 @@ def fit(
             # Clear memory of the SearchlightAlignment object
             denoiser = None
 
-        iterm = range(self.n_s)
+        iterm = range(self.n_subjects)
 
         if verbose:
             iterm = tqdm(iterm)
@@ -157,7 +155,7 @@ def fit(
         if self.decomp_method is None:
             full_signal = np.concatenate(self.denoised_signal, axis=1)
             self.shared_response = stimulus_estimator(
-                full_signal, self.n_t, self.n_s, self.latent_dim
+                full_signal, self.n_timepoints, self.n_subjects, self.latent_dim
             )
 
             self.tuning_data = Parallel(n_jobs=self.n_jobs)(
@@ -193,36 +191,18 @@ def transform(self, X_test_, verbose=False):
             print("Predict : Computing stimulus matrix...")
 
         if self.decomp_method is None:
-            S = stimulus_estimator(
-                full_signal, self.n_t, self.n_s, self.latent_dim
+            stimulus_ = stimulus_estimator(
+                full_signal, self.n_timepoints, self.n_subjects, self.latent_dim
             )
 
         if verbose:
-            print("Predict : stimulus matrix shape: ", S.shape)
+            print("Predict : stimulus matrix shape: ", stimulus_.shape)
 
         reconstructed_signal = [
-            reconstruct_signal(S, T_est) for T_est in self.tuning_data
+            reconstruct_signal(stimulus_, T_est) for T_est in self.tuning_data
         ]
         return np.array(reconstructed_signal, dtype=np.float32)
 
-    def get_shared_stimulus(self):
-        """
-        Returns the shared stimulus used for individualized neural tuning.
-
-        Returns:
-            The shared stimulus of shape (n_t, latent_dim) or (n_t, n_t).
-        """
-        return self.shared_response
-
-    def get_tuning_matrices(self):
-        """
-        Returns the tuning matrices as a NumPy array.
-
-        Returns:
-            numpy.ndarray: The tuning matrices of shape (n_s, latent_dim, n_v) or (n_s, n_t, n_v).
-        """
-        return np.array(self.tuning_data)
-
     def clean_cache(self, id):
         """
         Removes the cache file associated with the given ID.
@@ -246,12 +226,12 @@ def tuning_estimator(shared_response, target, latent_dim=None):
 
     Parameters:
     --------
-        - shared_response (array-like):
-            The shared response matrix.
-        - target (array-like):
-            The target matrix.
-        - latent_dim (int, optional):
-            The number of latent dimensions. Defaults to None.
+    - shared_response (array-like):
+        The shared response matrix.
+    - target (array-like):
+        The target matrix.
+    - latent_dim (int, optional):
+        The number of latent dimensions. Defaults to None.
 
     Returns:
     --------

From a0fd075564ee7cc450e614a331c1ca587f24db03 Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Thu, 4 Jan 2024 11:04:03 +0100
Subject: [PATCH 15/69] fix imports

---
 fmralign/_utils.py                       | 17 ++++++++---------
 fmralign/tests/test_alignment_methods.py | 14 +++++---------
 2 files changed, 13 insertions(+), 18 deletions(-)

diff --git a/fmralign/_utils.py b/fmralign/_utils.py
index 3b395cc..d712442 100644
--- a/fmralign/_utils.py
+++ b/fmralign/_utils.py
@@ -4,9 +4,9 @@
 
 import nibabel as nib
 import numpy as np
-from nilearn._utils.niimg_conversions import check_same_fov
+from nilearn._utils.niimg_conversions import _check_same_fov
 from nilearn.image import index_img, new_img_like, smooth_img
-from nilearn.masking import apply_mask_fmri, intersect_masks
+from nilearn.masking import _apply_mask_fmri, intersect_masks
 from nilearn.regions.parcellations import Parcellations
 
 
@@ -42,9 +42,7 @@ def piecewise_transform(labels, estimators, X):
 
     for i in range(len(unique_labels)):
         label = unique_labels[i]
-        X_transform[:, labels == label] = estimators[i].transform(
-            X[:, labels == label]
-        )
+        X_transform[:, labels == label] = estimators[i].transform(X[:, labels == label])
     return X_transform
 
 
@@ -115,8 +113,8 @@ def _make_parcellation(
     """
     # check if clustering is provided
     if isinstance(clustering, nib.nifti1.Nifti1Image) or os.path.isfile(clustering):
-        check_same_fov(masker.mask_img_, clustering)
-        labels = apply_mask_fmri(clustering, masker.mask_img_).astype(int)
+        _check_same_fov(masker.mask_img_, clustering)
+        labels = _apply_mask_fmri(clustering, masker.mask_img_).astype(int)
 
     # otherwise check it's needed, if not return 1 everywhere
     elif n_pieces == 1:
@@ -148,8 +146,9 @@ def _make_parcellation(
             )
             err.args += (errmsg,)
             raise err
-        labels = apply_mask_fmri(parcellation.labels_img_, masker.mask_img_).astype(int)
-
+        labels = _apply_mask_fmri(parcellation.labels_img_, masker.mask_img_).astype(
+            int
+        )
 
     if verbose > 0:
         unique_labels, counts = np.unique(labels, return_counts=True)
diff --git a/fmralign/tests/test_alignment_methods.py b/fmralign/tests/test_alignment_methods.py
index 2effde9..6a906f7 100644
--- a/fmralign/tests/test_alignment_methods.py
+++ b/fmralign/tests/test_alignment_methods.py
@@ -94,18 +94,14 @@ def test_scaled_procrustes_on_simple_exact_cases():
     assert_array_almost_equal(R_test.T, R)
 
     """Scaled Matrix"""
-    X = np.array(
-        [[1.0, 2.0, 3.0, 4.0], [5.0, 3.0, 4.0, 6.0], [7.0, 8.0, -5.0, -2.0]]
-    )
+    X = np.array([[1.0, 2.0, 3.0, 4.0], [5.0, 3.0, 4.0, 6.0], [7.0, 8.0, -5.0, -2.0]])
 
     X = X - X.mean(axis=1, keepdims=True)
 
     Y = 2 * X
     Y = Y - Y.mean(axis=1, keepdims=True)
 
-    assert_array_almost_equal(
-        scaled_procrustes(X.T, Y.T, scaling=True)[0], np.eye(3)
-    )
+    assert_array_almost_equal(scaled_procrustes(X.T, Y.T, scaling=True)[0], np.eye(3))
     assert_array_almost_equal(scaled_procrustes(X.T, Y.T, scaling=True)[1], 2)
 
     """3D Rotation"""
@@ -217,9 +213,7 @@ def test_ott_backend():
     algo = OptimalTransportAlignment(
         reg=epsilon, metric="euclidean", tol=1e-5, max_iter=10000
     )
-    old_implem = POTAlignment(
-        reg=epsilon, metric="euclidean", tol=1e-5, max_iter=10000
-    )
+    old_implem = POTAlignment(reg=epsilon, metric="euclidean", tol=1e-5, max_iter=10000)
     algo.fit(X, Y)
     old_implem.fit(X, Y)
     assert_array_almost_equal(algo.R, old_implem.R, decimal=3)
@@ -248,6 +242,7 @@ def test_searchlight_alignment_with_ridge():
     X_pred = model.transform(X_test)
 
     # assert that the saved cached files exist
+    assert_array_almost_equal(X_pred, X_test, decimal=3)
     b = os.path.exists("cache/")
     shutil.rmtree("cache")
     assert b
@@ -272,6 +267,7 @@ def test_parcel_alignment():
     X_pred = model.transform(X_test)
 
     # assert that the saved cached files exist
+    assert_array_almost_equal(X_pred, X_test, decimal=3)
     b = os.path.exists("cache/")
     shutil.rmtree("cache")
     assert b

From fde7b3ec3d3c7167152203184fc9c167176d7d6b Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Thu, 4 Jan 2024 13:45:48 +0100
Subject: [PATCH 16/69] add docstring to int

---
 fmralign/alignment_methods.py                 |   8 +-
 ...ent.py => individualized_neural_tuning.py} |   9 +-
 fmralign/hyperalignment/model.py              | 280 ------------------
 fmralign/tests/test_alignment_methods.py      |   2 +-
 4 files changed, 12 insertions(+), 287 deletions(-)
 rename fmralign/hyperalignment/{hyperalignment.py => individualized_neural_tuning.py} (86%)
 delete mode 100644 fmralign/hyperalignment/model.py

diff --git a/fmralign/alignment_methods.py b/fmralign/alignment_methods.py
index 9ae12c1..2ee22aa 100644
--- a/fmralign/alignment_methods.py
+++ b/fmralign/alignment_methods.py
@@ -470,7 +470,7 @@ def transform(self, X):
         return X.dot(self.R)
 
 
-class Hyperalignment(Alignment):
+class IndividualizedNeuralTuning(Alignment):
     def __init__(
         self,
         tmpl_kind="pca",
@@ -479,10 +479,12 @@ def __init__(
         alignment_method="searchlight",
         id=None,
         cache=True,
-        n_jobs=-1,
+        n_jobs=1,
     ):
         """
-        Initialize the IndividualizedNeuralTuning object.
+        Method of alignment based on the Individualized Neural Tuning model, by Feilong Ma et al. (2023).
+        It uses searchlight/parcelation alignment to denoise the data, and then computes the stimulus response matrix.
+        See article : https://doi.org/10.1162/imag_a_00032
 
         Parameters:
         --------
diff --git a/fmralign/hyperalignment/hyperalignment.py b/fmralign/hyperalignment/individualized_neural_tuning.py
similarity index 86%
rename from fmralign/hyperalignment/hyperalignment.py
rename to fmralign/hyperalignment/individualized_neural_tuning.py
index d941776..b9b2839 100644
--- a/fmralign/hyperalignment/hyperalignment.py
+++ b/fmralign/hyperalignment/individualized_neural_tuning.py
@@ -1,17 +1,20 @@
 """
 A wrapper for the IndividualTuningModel class to be used in fmralign (taking Nifti1 images as input).
 """
-from .model import INT
+from .model import INTEstimator
 from .regions import compute_searchlights
 from nilearn.maskers import NiftiMasker
 from nibabel import Nifti1Image
 import numpy as np
 
 
-class HyperAlignment(INT):
+class IndividualizedNeuralTuning(INTEstimator):
     def __init__(self, method="searchlight", n_jobs=-1):
         """
-        Initialize the Hyperalignment object.
+        A wrapper for the IndividualTuningModel class to be used in fmralign (taking Nifti1 images as input).
+        Method of alignment based on the Individualized Neural Tuning model, by Feilong Ma et al. (2023).
+        It uses searchlight/parcelation alignment to denoise the data, and then computes the stimulus response matrix.
+        See article : https://doi.org/10.1162/imag_a_00032
 
         Parameters:
         --------
diff --git a/fmralign/hyperalignment/model.py b/fmralign/hyperalignment/model.py
deleted file mode 100644
index c7c08a4..0000000
--- a/fmralign/hyperalignment/model.py
+++ /dev/null
@@ -1,280 +0,0 @@
-import numpy as np
-from tqdm import tqdm
-from sklearn.base import BaseEstimator, TransformerMixin
-from .regions_alignment import RegionAlignment
-from .linalg import safe_svd, svd_pca
-from joblib import Parallel, delayed
-import os
-
-
-class INT(BaseEstimator, TransformerMixin):
-    def __init__(
-        self,
-        tmpl_kind="pca",
-        decomp_method=None,
-        latent_dim=None,
-        alignment_method="searchlight",
-        id=None,
-        cache=True,
-        n_jobs=-1,
-    ):
-        """
-        Initialize the IndividualizedNeuralTuning object.
-
-        Parameters:
-        --------
-
-        - tmpl_kind (str): The type of template to use for alignment. Default is "pca".
-        - decomp_method (str): The decomposition method to use. Default is None.
-        - latent_dim (int): The number of latent dimensions to use in the shared stimulus information matrix. Default is None.
-        - n_jobs (int): The number of parallel jobs to run. Default is -1.
-
-        Returns:
-        --------
-        None
-        """
-
-        self.n_subjects = None
-        self.n_timepoints = None
-        self.n_voxels = None
-        self.labels = None
-        self.alphas = None
-        self.alignment_method = alignment_method
-        if alignment_method == "parcel":
-            self.parcels = None
-
-        elif (
-            alignment_method == "searchlight"
-            or alignment_method == "ensemble_searchlight"
-        ):
-            self.searchlights = None
-            self.distances = None
-            self.radius = None
-
-        self.path = None
-        self.tuning_data = []
-        self.denoised_signal = []
-        self.decomp_method = decomp_method
-        self.tmpl_kind = tmpl_kind
-        self.latent_dim = latent_dim
-        self.n_jobs = n_jobs
-        self.cache = cache
-
-        if cache:
-            if id is None:
-                self.id = "default"
-            else:
-                self.id = id
-
-            path = os.path.join(os.getcwd(), f"cache/int/{self.id}")
-            # Check if cache folder exists
-            if not os.path.exists(path):
-                os.makedirs(path)
-
-            self.path = path
-
-    def fit(
-        self,
-        X_train,
-        searchlights=None,
-        parcels=None,
-        dists=None,
-        radius: int = 20,
-        verbose=True,
-    ):
-        """
-        Fits the IndividualizedNeuralTuning model to the training data.
-
-        Parameters:
-        --------
-
-        - X_train (array-like):
-            The training data of shape (n_subjects, n_samples, n_voxels).
-        - searchlights (array-like):
-            The searchlight indices for each subject, of shape (n_s, n_searchlights).
-        - parcels (array-like):
-            The parcel indices for each subject, of shape (n_s, n_parcels) (if not using searchlights)
-        - dists (array-like):
-            The distances of vertices to the center of their searchlight, of shape (n_searchlights, n_vertices_sl)
-        - radius (int, optional):
-            The radius of the searchlight sphere, in milimeters. Defaults to 20.
-        - verbose (bool, optional):
-            Whether to print progress information. Defaults to True.
-        - id (str, optional):
-            An identifier for caching purposes. Defaults to None.
-
-        Returns:
-        --------
-
-        - self (IndividualizedNeuralTuning):
-            The fitted model.
-        """
-
-        X_train_ = np.array(X_train, copy=True, dtype=np.float32)
-
-        self.n_subjects, self.n_timepoints, self.n_voxels = X_train_.shape
-
-        self.tuning_data = np.empty(self.n_subjects, dtype=np.float32)
-        self.denoised_signal = np.empty(self.n_subjects, dtype=np.float32)
-
-        if searchlights is None:
-            self.regions = parcels
-        else:
-            self.regions = searchlights
-            self.distances = dists
-            self.radius = radius
-
-        # check for cached data
-        try:
-            self.denoised_signal = np.load(self.path + "/train_data_denoised.npy")
-            if verbose:
-                print("Loaded denoised data from cache")
-
-        except:  # noqa: E722
-            denoiser = RegionAlignment(
-                alignment_method=self.alignment_method,
-                n_jobs=self.n_jobs,
-                verbose=verbose,
-                path=self.path,
-            )
-            self.denoised_signal = denoiser.fit_transform(
-                X_train_,
-                regions=self.regions,
-                dists=dists,
-                radius=radius,
-            )
-            # Clear memory of the SearchlightAlignment object
-            denoiser = None
-
-        iterm = range(self.n_subjects)
-
-        if verbose:
-            iterm = tqdm(iterm)
-
-        # Stimulus matrix computation
-        if self.decomp_method is None:
-            full_signal = np.concatenate(self.denoised_signal, axis=1)
-            self.shared_response = stimulus_estimator(
-                full_signal, self.n_timepoints, self.n_subjects, self.latent_dim
-            )
-
-            self.tuning_data = Parallel(n_jobs=self.n_jobs)(
-                delayed(tuning_estimator)(
-                    self.shared_response,
-                    self.denoised_signal[i],
-                    latent_dim=self.latent_dim,
-                )
-                for i in iterm
-            )
-
-        return self
-
-    def transform(self, X_test_, verbose=False):
-        """
-        Transforms the input test data using the hyperalignment model.
-
-        Args:
-            X_test_ (list of arrays):
-                The input test data.
-            verbose (bool, optional):
-                Whether to print verbose output. Defaults to False.
-            id (int, optional):
-                Identifier for the transformation. Defaults to None.
-
-        Returns:
-            numpy.ndarray: The transformed test data.
-        """
-
-        full_signal = np.concatenate(X_test_, axis=1, dtype=np.float32)
-
-        if verbose:
-            print("Predict : Computing stimulus matrix...")
-
-        if self.decomp_method is None:
-            stimulus_ = stimulus_estimator(
-                full_signal, self.n_timepoints, self.n_subjects, self.latent_dim
-            )
-
-        if verbose:
-            print("Predict : stimulus matrix shape: ", stimulus_.shape)
-
-        reconstructed_signal = [
-            reconstruct_signal(stimulus_, T_est) for T_est in self.tuning_data
-        ]
-        return np.array(reconstructed_signal, dtype=np.float32)
-
-    def clean_cache(self, id):
-        """
-        Removes the cache file associated with the given ID.
-
-        Args:
-            id (int): The ID of the cache file to be removed.
-        """
-        try:
-            os.remove("cache")
-        except:  # noqa: E722
-            print("No cache to remove")
-
-
-#######################################################################################
-# Computing decomposition
-
-
-def tuning_estimator(shared_response, target, latent_dim=None):
-    """
-    Estimate the tuning weights for individualized neural tuning.
-
-    Parameters:
-    --------
-    - shared_response (array-like):
-        The shared response matrix.
-    - target (array-like):
-        The target matrix.
-    - latent_dim (int, optional):
-        The number of latent dimensions. Defaults to None.
-
-    Returns:
-    --------
-        array-like: The estimated tuning weights.
-
-    """
-    if latent_dim is None:
-        return np.linalg.inv(shared_response).dot(target)
-    return np.linalg.pinv(shared_response).dot(target).astype(np.float32)
-
-
-def stimulus_estimator(full_signal, n_t, n_s, latent_dim=None):
-    """
-    Estimates the stimulus response using the given parameters.
-
-    Args:
-        full_signal (np.ndarray): The full signal data.
-        n_t (int): The number of time points.
-        n_s (int): The number of stimuli.
-        latent_dim (int, optional): The number of latent dimensions. Defaults to None.
-
-    Returns:
-        np.ndarray: The estimated shared response.
-    """
-    if latent_dim is not None and latent_dim < n_t:
-        U = svd_pca(full_signal)
-        U = U[:, :latent_dim]
-    else:
-        U, _, _ = safe_svd(full_signal)
-
-    shared_response = np.sqrt(n_s) * U
-    return shared_response.astype(np.float32)
-
-
-def reconstruct_signal(shared_response, individual_tuning):
-    """
-    Reconstructs the signal using the shared response and individual tuning.
-
-    Args:
-        shared_response (numpy.ndarray): The shared response of shape (n_t, n_t) or (n_t, latent_dim).
-        individual_tuning (numpy.ndarray): The individual tuning of shape (latent_dim, n_v) or (n_t, n_v).
-
-    Returns:
-        numpy.ndarray: The reconstructed signal of shape (n_t, n_v) (same shape as the original signal)
-    """
-    return (shared_response @ individual_tuning).astype(np.float32)
diff --git a/fmralign/tests/test_alignment_methods.py b/fmralign/tests/test_alignment_methods.py
index 6a906f7..830e479 100644
--- a/fmralign/tests/test_alignment_methods.py
+++ b/fmralign/tests/test_alignment_methods.py
@@ -19,7 +19,7 @@
     POTAlignment,
     RidgeAlignment,
     ScaledOrthogonalAlignment,
-    Hyperalignment as INT,
+    IndividualizedNeuralTuning as INT,
     _voxelwise_signal_projection,
     optimal_permutation,
     scaled_procrustes,

From 5d08ebdde77e020234b6e01f36ba5d176f1e0e49 Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Thu, 4 Jan 2024 13:52:31 +0100
Subject: [PATCH 17/69] remove shady variable names

---
 fmralign/alignment_methods.py | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/fmralign/alignment_methods.py b/fmralign/alignment_methods.py
index 2ee22aa..f1f2902 100644
--- a/fmralign/alignment_methods.py
+++ b/fmralign/alignment_methods.py
@@ -483,6 +483,7 @@ def __init__(
     ):
         """
         Method of alignment based on the Individualized Neural Tuning model, by Feilong Ma et al. (2023).
+        It works on 4D fMRI data, and is based on the assumption that the neural response to a stimulus is shared across subjects.
         It uses searchlight/parcelation alignment to denoise the data, and then computes the stimulus response matrix.
         See article : https://doi.org/10.1162/imag_a_00032
 
@@ -614,12 +615,12 @@ def fit(
         # Stimulus matrix computation
         if self.decomp_method is None:
             full_signal = np.concatenate(self.denoised_signal, axis=1)
-            self.shared_response = stimulus_estimator(
+            self.shared_response = _stimulus_estimator(
                 full_signal, self.n_t, self.n_s, self.latent_dim
             )
 
             self.tuning_data = Parallel(n_jobs=self.n_jobs)(
-                delayed(tuning_estimator)(
+                delayed(_tuning_estimator)(
                     self.shared_response,
                     self.denoised_signal[i],
                     latent_dim=self.latent_dim,
@@ -651,13 +652,15 @@ def transform(self, X_test_, verbose=False):
             print("Predict : Computing stimulus matrix...")
 
         if self.decomp_method is None:
-            S = stimulus_estimator(full_signal, self.n_t, self.n_s, self.latent_dim)
+            stimulus_ = _stimulus_estimator(
+                full_signal, self.n_t, self.n_s, self.latent_dim
+            )
 
         if verbose:
-            print("Predict : stimulus matrix shape: ", S.shape)
+            print("Predict : stimulus matrix shape: ", stimulus_.shape)
 
         reconstructed_signal = [
-            reconstruct_signal(S, T_est) for T_est in self.tuning_data
+            _reconstruct_signal(stimulus_, T_est) for T_est in self.tuning_data
         ]
         return np.array(reconstructed_signal, dtype=np.float32)
 
@@ -678,7 +681,7 @@ def clean_cache(self, id):
 # Computing decomposition
 
 
-def tuning_estimator(shared_response, target, latent_dim=None):
+def _tuning_estimator(shared_response, target, latent_dim=None):
     """
     Estimate the tuning weights for individualized neural tuning.
 
@@ -699,7 +702,7 @@ def tuning_estimator(shared_response, target, latent_dim=None):
     return np.linalg.pinv(shared_response).dot(target).astype(np.float32)
 
 
-def stimulus_estimator(full_signal, n_t, n_s, latent_dim=None):
+def _stimulus_estimator(full_signal, n_t, n_s, latent_dim=None):
     """
     Estimates the stimulus response using the given parameters.
 
@@ -723,7 +726,7 @@ def stimulus_estimator(full_signal, n_t, n_s, latent_dim=None):
     return stimulus
 
 
-def reconstruct_signal(shared_response, individual_tuning):
+def _reconstruct_signal(shared_response, individual_tuning):
     """
     Reconstructs the signal using the shared response and individual tuning.
 

From 49ed356d3282f243787f412213d68f407cce7dd4 Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Fri, 5 Jan 2024 14:24:00 +0100
Subject: [PATCH 18/69] add: toy experiment + more precise test with
 correlation

---
 examples/plot_alignment_methods_benchmark.py  |  12 +-
 examples/plot_int_alignment.py                | 194 ++++++++++++
 examples/plot_template_alignment.py           |   4 +-
 fmralign/_utils.py                            |  12 +-
 fmralign/alignment_methods.py                 |  23 +-
 fmralign/hyperalignment/correlation.py        | 299 ++++++++++++++++++
 .../individualized_neural_tuning.py           |  65 ++--
 fmralign/hyperalignment/regions.py            |  18 +-
 fmralign/hyperalignment/regions_alignment.py  |   4 +-
 .../hyperalignment/test_hyperalignment.py     | 124 ++++++++
 fmralign/hyperalignment/toy_experiment.py     | 177 +++++++++++
 fmralign/pairwise_alignment.py                |  18 +-
 fmralign/tests/test_alignment_methods.py      |   4 +-
 13 files changed, 886 insertions(+), 68 deletions(-)
 create mode 100644 examples/plot_int_alignment.py
 create mode 100644 fmralign/hyperalignment/correlation.py
 create mode 100644 fmralign/hyperalignment/test_hyperalignment.py
 create mode 100644 fmralign/hyperalignment/toy_experiment.py

diff --git a/examples/plot_alignment_methods_benchmark.py b/examples/plot_alignment_methods_benchmark.py
index 2fe0fac..3f4eb77 100644
--- a/examples/plot_alignment_methods_benchmark.py
+++ b/examples/plot_alignment_methods_benchmark.py
@@ -48,9 +48,7 @@
 # Select visual cortex, create a mask and resample it to the right resolution
 
 mask_visual = new_img_like(atlas, atlas.get_fdata() == 1)
-resampled_mask_visual = resample_to_img(
-    mask_visual, mask, interpolation="nearest"
-)
+resampled_mask_visual = resample_to_img(mask_visual, mask, interpolation="nearest")
 
 # Plot the mask we will use
 plotting.plot_roi(
@@ -131,7 +129,13 @@
 from fmralign.metrics import score_voxelwise
 from fmralign.pairwise_alignment import PairwiseAlignment
 
-methods = ["identity", "scaled_orthogonal", "ridge_cv", "optimal_transport"]
+methods = [
+    "identity",
+    "scaled_orthogonal",
+    "ridge_cv",
+    "optimal_transport",
+    "individualized_neural_tuning",
+]
 
 for method in methods:
     alignment_estimator = PairwiseAlignment(
diff --git a/examples/plot_int_alignment.py b/examples/plot_int_alignment.py
new file mode 100644
index 0000000..b684641
--- /dev/null
+++ b/examples/plot_int_alignment.py
@@ -0,0 +1,194 @@
+# -*- coding: utf-8 -*-
+
+"""
+Template-based prediction.
+==========================
+
+In this tutorial, we show how to better predict new contrasts for a target
+subject using many source subjects corresponding contrasts. For this purpose,
+we create a template to which we align the target subject, using shared information.
+We then predict new images for the target and compare them to a baseline.
+
+We mostly rely on Python common packages and on nilearn to handle
+functional data in a clean fashion.
+
+
+To run this example, you must launch IPython via ``ipython
+--matplotlib`` in a terminal, or use ``jupyter-notebook``.
+
+.. contents:: **Contents**
+    :local:
+    :depth: 1
+
+"""
+
+###############################################################################
+# Retrieve the data
+# -----------------
+# In this example we use the IBC dataset, which includes a large number of
+# different contrasts maps for 12 subjects.
+# We download the images for subjects sub-01, sub-02, sub-04, sub-05, sub-06
+# and sub-07 (or retrieve them if they were already downloaded).
+# imgs is the list of paths to available statistical images for each subjects.
+# df is a dataframe with metadata about each of them.
+# mask is a binary image used to extract grey matter regions.
+#
+
+from fmralign.fetch_example_data import fetch_ibc_subjects_contrasts
+
+imgs, df, mask_img = fetch_ibc_subjects_contrasts(
+    ["sub-01", "sub-02", "sub-04", "sub-05", "sub-06", "sub-07"]
+)
+
+###############################################################################
+# Definine a masker
+# -----------------
+# We define a nilearn masker that will be used to handle relevant data.
+#   For more information, visit :
+#   'http://nilearn.github.io/manipulating_images/masker_objects.html'
+#
+
+from nilearn.maskers import NiftiMasker
+
+masker = NiftiMasker(mask_img=mask_img).fit()
+
+###############################################################################
+# Prepare the data
+# ----------------
+# For each subject, we will use two series of contrasts acquired during
+# two independent sessions with a different phase encoding:
+# Antero-posterior(AP) or Postero-anterior(PA).
+#
+
+
+# To infer a template for subjects sub-01 to sub-06 for both AP and PA data,
+# we make a list of 4D niimgs from our list of list of files containing 3D images
+
+from nilearn.image import concat_imgs
+
+template_train = []
+for i in range(5):
+    template_train.append(concat_imgs(imgs[i]))
+target_train = df[df.subject == "sub-07"][df.acquisition == "ap"].path.values
+
+# For subject sub-07, we split it in two folds:
+#   - target train: sub-07 AP contrasts, used to learn alignment to template
+#   - target test: sub-07 PA contrasts, used as a ground truth to score predictions
+# We make a single 4D Niimg from our list of 3D filenames
+
+target_train = concat_imgs(target_train)
+target_test = df[df.subject == "sub-07"][df.acquisition == "pa"].path.values
+
+###############################################################################
+# Compute a baseline (average of subjects)
+# ----------------------------------------
+# We create an image with as many contrasts as any subject representing for
+# each contrast the average of all train subjects maps.
+#
+
+import numpy as np
+
+masked_imgs = [masker.transform(img) for img in template_train]
+average_img = np.mean(masked_imgs, axis=0)
+average_subject = masker.inverse_transform(average_img)
+
+###############################################################################
+# Create a template from the training subjects.
+# ---------------------------------------------
+# We define an estimator using the class TemplateAlignment:
+#   * We align the whole brain through 'multiple' local alignments.
+#   * These alignments are calculated on a parcellation of the brain in 150 pieces,
+#     this parcellation creates group of functionnally similar voxels.
+#   * The template is created iteratively, aligning all subjects data into a
+#     common space, from which the template is inferred and aligning again to this
+#     new template space.
+#
+
+from nilearn.image import index_img
+
+from fmralign.alignment_methods import IndividualizedNeuralTuning
+
+from fmralign.hyperalignment.regions import compute_parcels
+
+parcels = compute_parcels(niimg=template_train[0], mask=masker, n_parcels=150, n_jobs=5)
+
+train_index = range(53)
+model = IndividualizedNeuralTuning(n_jobs=5, alignment_method="parcelation")
+model.fit(np.array(masked_imgs)[:, train_index, :], parcels=parcels, verbose=1)
+stimulus = model.shared_response
+
+
+###############################################################################
+# Predict new data for left-out subject
+# -------------------------------------
+# We use target_train data to fit the transform, indicating it corresponds to
+# the contrasts indexed by train_index and predict from this learnt alignment
+# contrasts corresponding to template test_index numbers.
+# For each train subject and for the template, the AP contrasts are sorted from
+# 0, to 53, and then the PA contrasts from 53 to 106.
+#
+
+train_index = range(53)
+test_index = range(53, 106)
+
+# We input the mapping image target_train in a list, we could have input more
+# than one subject for which we'd want to predict : [train_1, train_2 ...]
+target_train_array = np.array([masker.transform(target_train)[train_index]])
+
+prediction_from_template = model.fit(target_train_array, parcels=parcels, verbose=1)
+tuning_target = model.tuning_data
+
+prediction_from_template = stimulus @ tuning_target[0]
+prediction_from_template = prediction_from_template[test_index, :]
+prediction_from_template = [masker.inverse_transform(prediction_from_template)]
+
+# As a baseline prediction, let's just take the average of activations across subjects.
+
+prediction_from_average = index_img(average_subject, test_index)
+
+###############################################################################
+# Score the baseline and the prediction
+# -------------------------------------
+# We use a utility scoring function to measure the voxelwise correlation
+# between the prediction and the ground truth. That is, for each voxel, we
+# measure the correlation between its profile of activation without and with
+# alignment, to see if alignment was able to predict a signal more alike the ground truth.
+#
+
+from fmralign.metrics import score_voxelwise
+
+# Now we use this scoring function to compare the correlation of predictions
+# made from group average and from template with the real PA contrasts of sub-07
+
+average_score = masker.inverse_transform(
+    score_voxelwise(target_test, prediction_from_average, masker, loss="corr")
+)
+template_score = masker.inverse_transform(
+    score_voxelwise(target_test, prediction_from_template[0], masker, loss="corr")
+)
+
+
+###############################################################################
+# Plotting the measures
+# ---------------------
+# Finally we plot both scores
+#
+
+from nilearn import plotting
+
+baseline_display = plotting.plot_stat_map(
+    average_score, display_mode="z", vmax=1, cut_coords=[-15, -5]
+)
+baseline_display.title("Group average correlation wt ground truth")
+display = plotting.plot_stat_map(
+    template_score, display_mode="z", cut_coords=[-15, -5], vmax=1
+)
+display.title("Template-based prediction correlation wt ground truth")
+
+###############################################################################
+# We observe that creating a template and aligning a new subject to it yields
+# a prediction that is better correlated with the ground truth than just using
+# the average activations of subjects.
+#
+
+plotting.show()
diff --git a/examples/plot_template_alignment.py b/examples/plot_template_alignment.py
index 7170aa8..e2458c8 100644
--- a/examples/plot_template_alignment.py
+++ b/examples/plot_template_alignment.py
@@ -155,9 +155,7 @@
     score_voxelwise(target_test, prediction_from_average, masker, loss="corr")
 )
 template_score = masker.inverse_transform(
-    score_voxelwise(
-        target_test, prediction_from_template[0], masker, loss="corr"
-    )
+    score_voxelwise(target_test, prediction_from_template[0], masker, loss="corr")
 )
 
 
diff --git a/fmralign/_utils.py b/fmralign/_utils.py
index d712442..f31641f 100644
--- a/fmralign/_utils.py
+++ b/fmralign/_utils.py
@@ -4,9 +4,9 @@
 
 import nibabel as nib
 import numpy as np
-from nilearn._utils.niimg_conversions import _check_same_fov
+from nilearn._utils.niimg_conversions import check_same_fov
 from nilearn.image import index_img, new_img_like, smooth_img
-from nilearn.masking import _apply_mask_fmri, intersect_masks
+from nilearn.masking import apply_mask_fmri, intersect_masks
 from nilearn.regions.parcellations import Parcellations
 
 
@@ -113,8 +113,8 @@ def _make_parcellation(
     """
     # check if clustering is provided
     if isinstance(clustering, nib.nifti1.Nifti1Image) or os.path.isfile(clustering):
-        _check_same_fov(masker.mask_img_, clustering)
-        labels = _apply_mask_fmri(clustering, masker.mask_img_).astype(int)
+        check_same_fov(masker.mask_img_, clustering)
+        labels = apply_mask_fmri(clustering, masker.mask_img_).astype(int)
 
     # otherwise check it's needed, if not return 1 everywhere
     elif n_pieces == 1:
@@ -146,9 +146,7 @@ def _make_parcellation(
             )
             err.args += (errmsg,)
             raise err
-        labels = _apply_mask_fmri(parcellation.labels_img_, masker.mask_img_).astype(
-            int
-        )
+        labels = apply_mask_fmri(parcellation.labels_img_, masker.mask_img_).astype(int)
 
     if verbose > 0:
         unique_labels, counts = np.unique(labels, return_counts=True)
diff --git a/fmralign/alignment_methods.py b/fmralign/alignment_methods.py
index f1f2902..c817635 100644
--- a/fmralign/alignment_methods.py
+++ b/fmralign/alignment_methods.py
@@ -478,7 +478,7 @@ def __init__(
         latent_dim=None,
         alignment_method="searchlight",
         id=None,
-        cache=True,
+        cache=False,
         n_jobs=1,
     ):
         """
@@ -506,7 +506,7 @@ def __init__(
         self.labels = None
         self.alphas = None
         self.alignment_method = alignment_method
-        if alignment_method == "parcel":
+        if alignment_method == "parcelation":
             self.parcels = None
 
         elif (
@@ -546,6 +546,7 @@ def fit(
         parcels=None,
         dists=None,
         radius=20,
+        tuning=True,
         verbose=True,
     ):
         """
@@ -564,6 +565,8 @@ def fit(
             The distances of vertices to the center of their searchlight, of shape (n_searchlights, n_vertices_sl)
         - radius (int, optional):
             The radius of the searchlight sphere, in milimeters. Defaults to 20.
+        - tuning (bool, optional):
+            Whether to compute the tuning weights. Defaults to True.
         - verbose (bool, optional):
             Whether to print progress information. Defaults to True.
         - id (str, optional):
@@ -618,15 +621,15 @@ def fit(
             self.shared_response = _stimulus_estimator(
                 full_signal, self.n_t, self.n_s, self.latent_dim
             )
-
-            self.tuning_data = Parallel(n_jobs=self.n_jobs)(
-                delayed(_tuning_estimator)(
-                    self.shared_response,
-                    self.denoised_signal[i],
-                    latent_dim=self.latent_dim,
+            if tuning:
+                self.tuning_data = Parallel(n_jobs=self.n_jobs)(
+                    delayed(_tuning_estimator)(
+                        self.shared_response,
+                        self.denoised_signal[i],
+                        latent_dim=self.latent_dim,
+                    )
+                    for i in range(self.n_s)
                 )
-                for i in range(self.n_s)
-            )
 
         return self
 
diff --git a/fmralign/hyperalignment/correlation.py b/fmralign/hyperalignment/correlation.py
new file mode 100644
index 0000000..21ba050
--- /dev/null
+++ b/fmralign/hyperalignment/correlation.py
@@ -0,0 +1,299 @@
+import numpy as np
+from scipy.stats import spearmanr
+from sklearn.metrics import pairwise_distances
+from scipy.spatial.distance import cdist
+from sklearn.manifold import MDS
+from scipy.optimize import linear_sum_assignment
+
+
+#############################################################################
+# METRICS
+#############################################################################
+
+
+def compute_correlation(X, Y, metric: str = "correlation"):
+    """Compute correlation between X and Y.
+
+    Parameters
+    ----------
+    X : ndarray of shape (n_samples, n_features)
+        The data.
+    Y : ndarray of shape (n_samples, n_features)
+        The data.
+    metric : {'correlation', 'spearman', 'euclidean'}
+        The metric to use.
+
+    Returns
+    -------
+    corr : ndarray of shape (n_samples, n_samples)
+        The correlation matrix.
+    """
+    if metric == "correlation":
+        corr = 1 - pairwise_distances(X, Y, metric="correlation")
+    elif metric == "spearman":
+        corr = np.zeros((X.shape[0], Y.shape[0]))
+        for i in range(X.shape[0]):
+            for j in range(Y.shape[0]):
+                corr[i, j] = spearmanr(X[i], Y[j])[0]
+    elif metric == "euclidean":
+        corr = -pairwise_distances(X, Y, metric="euclidean")
+    else:
+        raise ValueError("Unknown metric")
+    return corr
+
+
+def compute_similarity(X, Y, metric: str = "euclidean"):
+    """
+    Compute the similarity matrix between two sets of vectors.
+
+    Parameters:
+    X (ndarray): First set of vectors.
+    Y (ndarray): Second set of vectors.
+    metric (str, optional): The distance metric to use. Defaults to "euclidean".
+
+    Returns:
+    ndarray: The similarity matrix between X and Y.
+    """
+    assert X.shape == Y.shape
+    n = X.shape[0]
+
+    def vector_sim(u, v):
+        if metric == "euclidean":
+            return 1 - (np.linalg.norm(u - v) / np.linalg.norm(u + v)) ** 2
+        else:
+            return (1 - cdist(u, v, metric)).mean()
+
+    sim = np.zeros((n, n))
+    for i in range(n):
+        for j in range(i, n):
+            d = vector_sim(X[i], Y[j])
+            sim[i, j] = d
+            sim[j, i] = d
+    return sim
+
+
+def compute_pearson_corr(X, Y, linear_assignment: bool = True):
+    """Compute Pearson correlation between X and Y.
+    X and Y are two lists of matrices of the same shape.
+    The returned matrix will be of shape 2N x 2N, where N is the number of matrices in X and Y.
+    """
+
+    assert X.shape == Y.shape
+
+    XY = np.concatenate((X, Y), axis=0)
+    n = XY.shape[0]
+    corr_mat = np.zeros((n, n))
+    for i in range(n):
+        for j in range(n):
+            corr_i_j = pearson_corr_coeff(XY[i], XY[j])
+            corr_mat[i, j] = corr_i_j
+
+    return corr_mat
+
+
+def pearson_corr_coeff(
+    M1: np.ndarray,
+    M2: np.ndarray,
+    absolute: bool = True,
+    linear_assignment: bool = True,
+):
+    """
+    Compute Pearson correlation coefficient between matrices M1 and M2.
+
+    Parameters:
+        M1 (numpy.ndarray): First matrix.
+        M2 (numpy.ndarray): Second matrix.
+        absolute (bool, optional): Whether to compute absolute correlation coefficients. Defaults to True.
+        linear_assignment (bool, optional): Whether to perform linear assignment optimization. Defaults to True.
+
+    Returns:
+        float: Pearson correlation coefficient.
+    """
+    assert M1.shape == M2.shape
+
+    n = M1.shape[0]
+    corr = np.corrcoef(M1, M2)[:n, n:]
+
+    corr = np.abs(corr) if absolute else corr
+
+    if linear_assignment:
+        row_ind, col_ind = linear_sum_assignment(corr, maximize=True)
+
+        # permutation of columns and rows
+        corr = corr[row_ind, :]
+        corr = corr[:, col_ind]
+
+    corr_diag = np.diag(corr)
+
+    corr_coeff = np.mean(corr_diag)
+
+    return corr_coeff
+
+
+def tuning_correlation(X, Y):
+    """Compute pairwise Pearson correlation matrix between two sets of matrices."""
+    assert X.shape == Y.shape
+    n = X.shape[0]
+    corr_mat = np.zeros((n, n))
+    for i in range(n):
+        for j in range(i, n):
+            corr_i_j = pearson_corr_coeff(X[i], Y[j])
+            corr_mat[i, j] = corr_i_j
+            corr_mat[j, i] = corr_i_j
+
+    return corr_mat
+
+
+def stimulus_correlation(X, Y, linear_assignment=True, absolute=True):
+    """Compute pairwise Pearson correlation matrix between two stimulus matrices."""
+    assert X.shape == Y.shape
+    n = X.shape[0]
+    corr_mat = np.corrcoef(X, Y, dtype=np.float32)[:n, n:]
+
+    if absolute:
+        corr_mat = np.abs(corr_mat)
+
+    if linear_assignment:
+        row_ind, col_ind = linear_sum_assignment(corr_mat, maximize=True)
+        corr_mat = corr_mat[row_ind, :]
+        corr_mat = corr_mat[:, col_ind]
+
+    return corr_mat.astype(np.float16)
+
+
+def matrix_MDS(X, Y, n_components=2, dissimilarity="euclidean"):
+    """
+    Perform multidimensional scaling (MDS) on the given data matrices X and Y.
+
+    Parameters:
+        X (list): The first data matrix.
+        Y (list): The second data matrix.
+        n_components (int): The number of dimensions in the output space (default is 2).
+        dissimilarity (str or array-like): The dissimilarity measure to use. If it is a string other than "precomputed",
+                                          the dissimilarity is computed using the Euclidean distance between flattened data points.
+                                          If it is "precomputed", the dissimilarity is assumed to be a precomputed dissimilarity matrix.
+
+    Returns:
+        tuple: A tuple containing two arrays. The first array represents the transformed data points from matrix X,
+               and the second array represents the transformed data points from matrix Y.
+    """
+    assert len(X) == len(Y)
+
+    if isinstance(dissimilarity, str) and dissimilarity != "precomputed":
+        X_flat = [x.flatten() for x in X]
+        Y_flat = [y.flatten() for y in Y]
+        XY = np.array(X_flat + Y_flat)
+        mds = MDS(n_components=n_components, dissimilarity=dissimilarity)
+        transformed = mds.fit_transform(XY)
+
+    else:
+        mds = MDS(n_components=n_components, dissimilarity="precomputed")
+        transformed = mds.fit_transform(dissimilarity)
+
+    return np.array(transformed[: len(X)]), np.array(transformed[len(X) :])
+
+
+def thread_compute_correlation(X, Y, i, j):
+    """
+    Compute the correlation between two time series X_i and Y_i.
+
+    Parameters:
+       - X (ndarray): Array of shape (n_samples, n_features) representing the first time series.
+       - Y (ndarray): Array of shape (n_samples, n_features) representing the second time series.
+       - i (int): Index of the first time series.
+       - j (int): Index of the second time series.
+
+    Returns:
+    diff_TR_corr (ndarray): Array of shape (n_samples * (n_samples - 1),) containing the correlations between different time points.
+    same_TR_corr (ndarray): Array of shape (n_samples,) containing the correlations between the same time points.
+    empty_diff_TR_corr (ndarray): Empty array.
+    empty_same_TR_corr (ndarray): Empty array.
+    """
+    X_i, Y_i = X[i], Y[i]
+    corr = stimulus_correlation(X_i, Y_i, absolute=False)
+    same_TR_corr = np.diag(corr)
+    # Get all the values except the diagonal in a list
+    diff_TR_corr = corr[np.where(~np.eye(corr.shape[0], dtype=bool))]
+    diff_TR_corr = diff_TR_corr.flatten()
+    if i == j:
+        return (
+            np.array([]),
+            np.array([]),
+            [x for x in diff_TR_corr],
+            [x for x in same_TR_corr],
+        )
+
+    else:
+        return diff_TR_corr, same_TR_corr, np.array([]), np.array([])
+
+
+def multithread_compute_correlation(
+    X, Y, absolute=False, linear_assignment=True, n_jobs=40
+):
+    """
+    Compute correlations between pairs of samples in X and Y using multiple threads.
+
+    Args:
+        X (ndarray): The first set of samples, with shape (n_samples, n_features).
+        Y (ndarray): The second set of samples, with shape (n_samples, n_features).
+        absolute (bool, optional): Whether to compute absolute correlations. Defaults to False.
+        linear_assignment (bool, optional): Whether to use linear assignment for correlation computation. Defaults to True.
+        n_jobs (int, optional): The number of threads to use for parallel computation. Defaults to 40.
+
+    Returns:
+        tuple: A tuple containing four arrays:
+            - corr_same_sub_diff_TR: Correlations between different time points of the same subject.
+            - corr_same_sub_same_TR: Correlations between the same time points of the same subject.
+            - corr_diff_sub_diff_TR: Correlations between different time points of different subjects.
+            - corr_diff_sub_same_TR: Correlations between the same time points of different subjects.
+    """
+    from joblib import Parallel, delayed
+    from tqdm import tqdm
+
+    def thread_compute_correlation(X, Y, i, j, absolute=False, linear_assignment=True):
+        X_i, Y_i = X[i], Y[i]
+        corr = stimulus_correlation(X_i, Y_i, absolute=False)
+        same_TR_corr = np.diag(corr)
+        # Get all the values except the diagonal in a list
+        diff_TR_corr = corr[np.where(~np.eye(corr.shape[0], dtype=bool))]
+        diff_TR_corr = diff_TR_corr.flatten().astype(np.float16)
+        if i == j:
+            return (
+                np.array([]),
+                np.array([]),
+                [x for x in diff_TR_corr],
+                [x for x in same_TR_corr],
+            )
+
+        else:
+            return (
+                diff_TR_corr.astype(np.float16),
+                same_TR_corr.astype(np.float16),
+                np.array([]),
+                np.array([]),
+            )
+
+    from itertools import combinations
+
+    assert X.shape == Y.shape
+    n_s = X.shape[0]
+    corrdinates = list(combinations(range(n_s), 2)) + [(i, i) for i in range(n_s)]
+    results = Parallel(n_jobs=n_jobs)(
+        delayed(thread_compute_correlation)(X, Y, i, j) for (i, j) in tqdm(corrdinates)
+    )
+    results = list(zip(*results))
+    corr_same_sub_diff_TR = results[2]
+    corr_same_sub_same_TR = results[3]
+    corr_diff_sub_diff_TR = results[0]
+    corr_diff_sub_same_TR = results[1]
+
+    corr_same_sub_diff_TR = np.concatenate(corr_same_sub_diff_TR)
+    corr_same_sub_same_TR = np.concatenate(corr_same_sub_same_TR)
+    corr_diff_sub_diff_TR = np.concatenate(corr_diff_sub_diff_TR)
+    corr_diff_sub_same_TR = np.concatenate(corr_diff_sub_same_TR)
+    return (
+        corr_same_sub_diff_TR,
+        corr_same_sub_same_TR,
+        corr_diff_sub_diff_TR,
+        corr_diff_sub_same_TR,
+    )
diff --git a/fmralign/hyperalignment/individualized_neural_tuning.py b/fmralign/hyperalignment/individualized_neural_tuning.py
index b9b2839..47db245 100644
--- a/fmralign/hyperalignment/individualized_neural_tuning.py
+++ b/fmralign/hyperalignment/individualized_neural_tuning.py
@@ -1,15 +1,23 @@
 """
 A wrapper for the IndividualTuningModel class to be used in fmralign (taking Nifti1 images as input).
 """
-from .model import INTEstimator
-from .regions import compute_searchlights
+from fmralign.alignment_methods import IndividualizedNeuralTuning as BaseINT
+from .regions import compute_searchlights, compute_parcels
 from nilearn.maskers import NiftiMasker
 from nibabel import Nifti1Image
 import numpy as np
 
 
-class IndividualizedNeuralTuning(INTEstimator):
-    def __init__(self, method="searchlight", n_jobs=-1):
+class IndividualizedNeuralTuning(BaseINT):
+    def __init__(
+        self,
+        tmpl_kind="pca",
+        decomp_method=None,
+        alignment_method="searchlight",
+        n_pieces=150,
+        radius=20,
+        n_jobs=1,
+    ):
         """
         A wrapper for the IndividualTuningModel class to be used in fmralign (taking Nifti1 images as input).
         Method of alignment based on the Individualized Neural Tuning model, by Feilong Ma et al. (2023).
@@ -19,22 +27,31 @@ def __init__(self, method="searchlight", n_jobs=-1):
         Parameters:
         --------
         - method (str): The method used for hyperalignment. Can be either "searchlight" or "parcellation". Default is "searchlight".
-        - n_jobs (int): The number of parallel jobs to run. Default is -1, which uses all available processors.
+        - n_jobs (int): The number of parallel jobs to run. Default is 1.
+        - n_pieces (int): The number of pieces to divide the brain into. Default is 150.
+        - radius (int): The radius of the searchlight in millimeters. Default is 20.
 
         Returns:
         --------
         None
         """
-        super().__init__(n_jobs=n_jobs)
+        super().__init__(
+            tmpl_kind=tmpl_kind,
+            decomp_method=decomp_method,
+            alignment_method=alignment_method,
+            n_jobs=n_jobs,
+        )
+        self.n_pieces = n_pieces
+        self.radius = radius
         self.mask_img = None
         self.masker = None
-        self.method = method
 
     def fit(
         self,
         imgs,
         masker: NiftiMasker = None,
         mask_img: Nifti1Image = None,
+        tuning: bool = True,
         y=None,
         verbose=0,
     ):
@@ -45,12 +62,14 @@ def fit(
 
         Parameters
         ----------
-        imgs : list of Nifti1Image or np.ndarray
+        imgs : list of Nifti1Image
             The images to be fit.
         masker : NiftiMasker
             The masker to be used to transform the images into the common space.
         mask_img : Nifti1Image
             The mask to be used to transform the images into the common space.
+        tuning : bool
+            Whether to perform tuning or not.
         y : None
             Not used.
         verbose : int
@@ -62,21 +81,31 @@ def fit(
         self.mask_img = mask_img
         self.masker = masker
 
-        if self.method == "searchlight":
-            data, searchlights, dists = compute_searchlights(
+        X = np.array([masker.transform(img) for img in imgs])
+
+        if self.alignment_method == "searchlight":
+            _, searchlights, dists = compute_searchlights(
                 niimg=imgs[0],
                 mask_img=mask_img,
             )
+            super().fit(
+                X,
+                searchlights,
+                dists,
+                radius=self.radius,
+                tuning=tuning,
+                verbose=verbose,
+            )
 
-        elif self.method == "parcellation":
-            raise NotImplementedError("Parcellation method not implemented yet.")
-
-        if isinstance(imgs, np.ndarray):
-            X = imgs
-        else:
-            X = data
+        elif self.alignment_method == "parcellation":
+            parcels = compute_parcels(
+                niimg=imgs[0],
+                mask=masker.mask_img,
+                n_parcels=self.n_pieces,
+                n_jobs=self.n_jobs,
+            )
+            super().fit(X, parcels=parcels, tuning=tuning, verbose=verbose)
 
-        super().fit(X, searchlights, dists, verbose=verbose)
         return self
 
     def transform(self, imgs, y=None, verbose=0):
diff --git a/fmralign/hyperalignment/regions.py b/fmralign/hyperalignment/regions.py
index d012f54..8b795c1 100644
--- a/fmralign/hyperalignment/regions.py
+++ b/fmralign/hyperalignment/regions.py
@@ -12,7 +12,7 @@
 from sklearn import neighbors
 from scipy.spatial import distance_matrix
 from nilearn._utils.niimg_conversions import (
-    _safe_get_data,
+    safe_get_data,
 )
 from tqdm import tqdm
 from .linalg import procrustes
@@ -66,7 +66,9 @@ def compute_parcels(
         print("[Loading/Parcel] : Parcellating...")
 
     if isinstance(mask, Nifti1Image):
-        mask = NiftiMasker(mask_img=mask, standardize=True)
+        mask = NiftiMasker(
+            mask_img=mask, standardize=True, smoothing_fwhm=smoothing_fwhm
+        )
     # Parcellation
     indexes = [1]
     labels = _make_parcellation(
@@ -150,21 +152,21 @@ def _apply_mask_and_get_affinity(
             target_shape=niimg.shape[:3],
             interpolation="nearest",
         )
-        mask, _ = masking._load_mask_img(mask_img)
+        mask, _ = masking.load_mask_img(mask_img)
         mask_coords = list(zip(*np.where(mask != 0)))
 
-        X = masking._apply_mask_fmri(niimg, mask_img)
+        X = masking.apply_mask_fmri(niimg, mask_img)
 
     elif niimg is not None:
         affine = niimg.affine
-        if np.isnan(np.sum(_safe_get_data(niimg))):
+        if np.isnan(np.sum(safe_get_data(niimg))):
             warnings.warn(
                 "The imgs you have fed into fit_transform() contains NaN "
                 "values which will be converted to zeroes."
             )
-            X = _safe_get_data(niimg, True).reshape([-1, niimg.shape[3]]).T
+            X = safe_get_data(niimg, True).reshape([-1, niimg.shape[3]]).T
         else:
-            X = _safe_get_data(niimg).reshape([-1, niimg.shape[3]]).T
+            X = safe_get_data(niimg).reshape([-1, niimg.shape[3]]).T
 
         mask_coords = list(np.ndindex(niimg.shape[:3]))
 
@@ -290,7 +292,7 @@ def compute_searchlights(
         process_mask_img = mask_img
 
     # Compute world coordinates of the seeds
-    process_mask, process_mask_affine = masking._load_mask_img(process_mask_img)
+    process_mask, process_mask_affine = masking.load_mask_img(process_mask_img)
     process_mask_coords = np.where(process_mask != 0)
     process_mask_coords = coord_transform(
         process_mask_coords[0],
diff --git a/fmralign/hyperalignment/regions_alignment.py b/fmralign/hyperalignment/regions_alignment.py
index 082b0a5..4ba62bf 100644
--- a/fmralign/hyperalignment/regions_alignment.py
+++ b/fmralign/hyperalignment/regions_alignment.py
@@ -59,9 +59,7 @@ def __init__(
             if not os.path.exists(self.path):
                 os.makedirs(self.path)
 
-    def compute_linear_transformation(
-        self, x_i, template, i: int = 0, save=True
-    ):
+    def compute_linear_transformation(self, x_i, template, i: int = 0, save=True):
         """Compute the linear transformation W_i for a given subject.
         ----------
         Parameters
diff --git a/fmralign/hyperalignment/test_hyperalignment.py b/fmralign/hyperalignment/test_hyperalignment.py
new file mode 100644
index 0000000..14a22b8
--- /dev/null
+++ b/fmralign/hyperalignment/test_hyperalignment.py
@@ -0,0 +1,124 @@
+from fmralign.alignment_methods import IndividualizedNeuralTuning as INT
+from fmralign.generate_data import generate_dummy_signal, generate_dummy_searchlights
+import numpy as np
+
+
+def test_int_fit_predict():
+    """Test if the outputs and arguments of the INT are the correct format"""
+    # Create random data
+    X_train, X_test, S_true_first_part, S_true_second_part, Ts = generate_dummy_signal(
+        n_s=5,
+        n_t=50,
+        n_v=300,
+        S_std=1,
+        T_std=1,
+        latent_dim=6,
+        SNR=100,
+        generative_method="custom",
+        seed=0,
+    )
+    from fmralign.hyperalignment.correlation import (
+        tuning_correlation,
+        stimulus_correlation,
+    )
+
+    # Testing without searchlights
+    searchlights = [np.arange(300)]
+    dists = [np.ones((300,))]
+
+    # Test INT on the two parts of the data (ie different runs of the experiment)
+    int1 = INT(latent_dim=6)
+    int2 = INT(latent_dim=6)
+    int1.fit(
+        X_train=X_train, searchlights=searchlights, dists=dists
+    )  # S is provided if we cheat and know the ground truth
+    int2.fit(X_test, searchlights=searchlights, dists=dists)
+
+    X_pred = int1.transform(X_test)
+    # save individual components
+
+    tuning_data_run_1 = int1.tuning_data
+    tuning_data_run_2 = int2.tuning_data
+    tuning_data_run_1 = np.array(tuning_data_run_1)
+    tuning_data_run_2 = np.array(tuning_data_run_2)
+
+    stimulus_run_1 = int1.shared_response
+    S_estimated_second_part = int2.shared_response
+
+    corr1 = tuning_correlation(tuning_data_run_1, tuning_data_run_2)
+    corr2 = stimulus_correlation(stimulus_run_1.T, S_true_first_part.T)
+    corr3 = stimulus_correlation(S_estimated_second_part.T, S_true_second_part.T)
+    corr4 = tuning_correlation(X_pred, X_test)
+
+    # Check that predicted components have the same shape as original data
+
+    # Check that the correlation between the two parts of the data is high
+    corr1_out = corr1 - np.diag(corr1)
+    corr2_out = corr2 - np.diag(corr2)
+    corr3_out = corr3 - np.diag(corr3)
+    corr4_out = corr4 - np.diag(corr4)
+    assert 3 * np.mean(corr1_out) < np.mean(np.diag(corr1))
+    assert 3 * np.mean(corr2_out) < np.mean(np.diag(corr2))
+    assert 3 * np.mean(corr3_out) < np.mean(np.diag(corr3))
+    assert 3 * np.mean(corr4_out) < np.mean(np.diag(corr4))
+    assert int1.tuning_data[0].shape == (6, int1.n_v)
+    assert int2.tuning_data[0].shape == (6, int2.n_v)
+    assert int1.shared_response.shape == (int1.n_t, 6)
+    assert X_pred.shape == X_test.shape
+
+
+def test_int_with_searchlight():
+    X_train, X_test, S_true_first_part, S_true_second_part, Ts = generate_dummy_signal(
+        n_s=5,
+        n_t=50,
+        n_v=300,
+        S_std=1,
+        T_std=1,
+        latent_dim=6,
+        SNR=100,
+        generative_method="custom",
+        seed=0,
+    )
+    searchlights, dists = generate_dummy_searchlights(
+        n_searchlights=10, n_v=300, radius=5, seed=0
+    )
+    from fmralign.hyperalignment.correlation import (
+        tuning_correlation,
+        stimulus_correlation,
+    )
+
+    # Test INT on the two parts of the data (ie different runs of the experiment)
+    int1 = INT(latent_dim=6)
+    int2 = INT(latent_dim=6)
+    int1.fit(X_train=X_train, searchlights=searchlights, dists=dists, radius=5)
+    int2.fit(X_test, searchlights=searchlights, dists=dists, radius=5)
+    X_pred = int1.transform(X_test)
+
+    tuning_data_run_1 = int1.tuning_data
+    tuning_data_run_2 = int2.tuning_data
+    tuning_data_run_1 = np.array(tuning_data_run_1)
+    tuning_data_run_2 = np.array(tuning_data_run_2)
+
+    stimulus_run_1 = int1.shared_response
+    stimulus_run_2 = int2.shared_response
+
+    corr1 = tuning_correlation(tuning_data_run_1, tuning_data_run_2)
+    corr2 = stimulus_correlation(stimulus_run_1.T, S_true_first_part.T)
+    corr3 = stimulus_correlation(stimulus_run_2.T, S_true_second_part.T)
+    corr4 = tuning_correlation(X_pred, X_test)
+
+    # Check that predicted components have the same shape as original data
+
+    # Check that the correlation between the two parts of the data is high
+    corr1_out = corr1 - np.diag(corr1)
+    corr2_out = corr2 - np.diag(corr2)
+    corr3_out = corr3 - np.diag(corr3)
+    corr4_out = corr4 - np.diag(corr4)
+    assert 3 * np.mean(corr1_out) < np.mean(np.diag(corr1))
+    assert 3 * np.mean(corr2_out) < np.mean(np.diag(corr2))
+    assert 3 * np.mean(corr3_out) < np.mean(np.diag(corr3))
+    assert 3 * np.mean(corr4_out) < np.mean(np.diag(corr4))
+    assert int1.tuning_data[0].shape == (6, int1.n_v)
+    assert int2.tuning_data[0].shape == (6, int2.n_v)
+    assert int1.shared_response.shape == (int1.n_t, 6)
+    assert X_pred.shape == X_test.shape
diff --git a/fmralign/hyperalignment/toy_experiment.py b/fmralign/hyperalignment/toy_experiment.py
new file mode 100644
index 0000000..9b5350b
--- /dev/null
+++ b/fmralign/hyperalignment/toy_experiment.py
@@ -0,0 +1,177 @@
+import os
+import sys
+
+import numpy as np
+import matplotlib.pyplot as plt
+from fmralign.alignment_methods import IndividualizedNeuralTuning as INT
+from fmralign.generate_data import generate_dummy_signal, generate_dummy_searchlights
+from fmralign.hyperalignment.correlation import (
+    tuning_correlation,
+    stimulus_correlation,
+    compute_pearson_corr,
+    matrix_MDS,
+)
+
+
+#############################################################################
+# INT
+#############################################################################
+
+
+file_dir = os.path.dirname(__file__)
+sys.path.append(file_dir)
+
+
+n_s = 10
+n_t = 200
+n_v = 200
+S_std = 1
+T_std = 1
+SNR = 100
+latent_dim = 6  # if None, latent_dim = n_t
+decomposition_method = None  # if None, SVD is used
+
+
+#############################################################################
+# GENERATE DUMMY SIGNAL
+
+(
+    data_run_1,
+    data_run_2,
+    stimulus_run_1,
+    stimulus_run_2,
+    data_tuning,
+) = generate_dummy_signal(
+    n_s=n_s,
+    n_t=n_t,
+    n_v=n_v,
+    S_std=S_std,
+    T_std=T_std,
+    latent_dim=latent_dim,
+    SNR=SNR,
+    seed=0,
+)
+
+searchlights, dists = generate_dummy_searchlights(n_searchlights=12, n_v=n_v, radius=5)
+
+#############################################################################
+# Test INT on the two parts of the data (ie different runs of the experiment)
+int1 = INT(latent_dim=latent_dim, decomp_method=decomposition_method, cache=False)
+int2 = INT(latent_dim=latent_dim, decomp_method=decomposition_method, cache=False)
+int_first_part = int1.fit(
+    data_run_1,
+    searchlights=searchlights,
+    dists=dists,
+)  # S is provided if we cheat and know the ground truth
+int_second_part = int2.fit(data_run_2, searchlights=searchlights, dists=dists)
+
+print("INT 1 denoised signal", int1.denoised_signal.shape)
+
+data_pred = int1.transform(data_run_2)
+# save individual components
+
+tuning_pred_run_1 = int1.tuning_data
+tuning_pred_run_1 = np.array(tuning_pred_run_1)
+tuning_pred_run_2 = int2.tuning_data
+tuning_pred_run_2 = np.array(tuning_pred_run_2)
+
+stimulus_pred_run_1 = int1.shared_response
+stimulus_pred_run_2 = int2.shared_response
+
+
+#############################################################################
+# Plot
+plt.rc("font", size=6)
+fig, ax = plt.subplots(3, 3, figsize=(20, 10))
+
+
+# Tunning matrices
+correlation_tuning = tuning_correlation(tuning_pred_run_1, tuning_pred_run_2)
+ax[0, 0].imshow(correlation_tuning)
+ax[0, 0].set_title("Run 1 vs Run 2")
+ax[0, 0].set_xlabel("Subjects, Run 1")
+ax[0, 0].set_ylabel("Subjects, Run 2")
+fig.colorbar(ax[0, 0].imshow(correlation_tuning), ax=ax[0, 0])
+
+random_colors = np.random.rand(n_s, 3)
+# MDS of predicted images
+print("X_run_2_pred.shape", data_pred.shape)
+print("X_run_2.shape", data_run_2.shape)
+corr_tunning = compute_pearson_corr(data_pred, data_run_2)
+data_pred_reduced, data_test_reduced = matrix_MDS(
+    data_pred, data_run_2, n_components=2, dissimilarity=1 - corr_tunning
+)
+
+ax[0, 1].scatter(
+    data_pred_reduced[:, 0],
+    data_pred_reduced[:, 1],
+    label="Run 1",
+    c=random_colors,
+)
+ax[0, 1].scatter(
+    data_test_reduced[:, 0],
+    data_test_reduced[:, 1],
+    label="Run 2",
+    c=random_colors,
+)
+ax[0, 1].set_title("MDS of predicted images, dim=2")
+
+# MDS of tunning matrices
+corr_tunning = compute_pearson_corr(tuning_pred_run_1, tuning_pred_run_2)
+T_first_part_transformed, T_second_part_transformed = matrix_MDS(
+    tuning_pred_run_1, tuning_pred_run_2, n_components=2, dissimilarity=1 - corr_tunning
+)
+
+
+ax[0, 2].scatter(
+    T_first_part_transformed[:, 0],
+    T_first_part_transformed[:, 1],
+    label="Run 1",
+    c=random_colors,
+)
+ax[0, 2].scatter(
+    T_second_part_transformed[:, 0],
+    T_second_part_transformed[:, 1],
+    label="Run 2",
+    c=random_colors,
+)
+ax[0, 2].set_title("MDS of tunning matrices, dim=2")
+
+# Stimulus matrix
+
+correlation_stimulus_true_est_first_part = stimulus_correlation(
+    stimulus_pred_run_1.T, stimulus_run_1.T
+)
+ax[1, 0].imshow(correlation_stimulus_true_est_first_part)
+ax[1, 0].set_title("Stimumus Run 1 Estimated vs ground truth")
+ax[1, 0].set_xlabel("Latent components, Run 1")
+ax[1, 0].set_ylabel("Latent components, ground truth")
+fig.colorbar(ax[1, 0].imshow(correlation_stimulus_true_est_first_part), ax=ax[1, 0])
+
+correlation_stimulus_true_est_second_part = stimulus_correlation(
+    stimulus_pred_run_2.T, stimulus_run_2.T
+)
+ax[1, 1].imshow(correlation_stimulus_true_est_second_part)
+ax[1, 1].set_title("Stimulus Run 2 Estimated vs ground truth")
+ax[1, 1].set_xlabel("Latent components, Run 2")
+ax[1, 1].set_ylabel("Latent components, ground truth")
+fig.colorbar(ax[1, 1].imshow(correlation_stimulus_true_est_second_part), ax=ax[1, 1])
+
+
+# Reconstruction
+correlation_reconstruction_first_second = tuning_correlation(data_pred, data_run_2)
+ax[1, 2].imshow(correlation_reconstruction_first_second)
+ax[1, 2].set_title("Reconstruction")
+ax[1, 2].set_xlabel("Subjects, Run 2")
+ax[1, 2].set_ylabel("Subjects, Run 1")
+fig.colorbar(ax[1, 2].imshow(correlation_reconstruction_first_second), ax=ax[1, 2])
+
+
+plt.rc("font", size=10)
+# Define small font for titles
+fig.suptitle(
+    f"Tunning matrices for the two parts of the data\n ns={n_s}, nt={n_t}, nv={n_v}, S_std={S_std}, T_std={T_std}, SNR={SNR}, latent space dim={latent_dim}"
+)
+plt.tight_layout()
+
+plt.show()
diff --git a/fmralign/pairwise_alignment.py b/fmralign/pairwise_alignment.py
index f28e36e..522648b 100644
--- a/fmralign/pairwise_alignment.py
+++ b/fmralign/pairwise_alignment.py
@@ -53,9 +53,7 @@ def generate_Xi_Yi(labels, X, Y, masker, verbose):
         label = unique_labels[k]
         i = label == labels
         if (k + 1) % 25 == 0 and verbose > 0:
-            print(
-                "Fitting parcel: " + str(k + 1) + "/" + str(len(unique_labels))
-            )
+            print("Fitting parcel: " + str(k + 1) + "/" + str(len(unique_labels)))
         # should return X_i Y_i
         yield X_[:, i], Y_[:, i]
 
@@ -323,9 +321,9 @@ def fit(self, X, Y):
         else:
             self.masker_.fit()
 
-        if isinstance(
-            self.clustering, nib.nifti1.Nifti1Image
-        ) or os.path.isfile(self.clustering):
+        if isinstance(self.clustering, nib.nifti1.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(
@@ -352,9 +350,7 @@ def fit(self, X, Y):
         self.fit_, self.labels_ = [], []
         rs = ShuffleSplit(n_splits=self.n_bags, test_size=0.8, random_state=0)
 
-        outputs = Parallel(
-            n_jobs=self.n_jobs, prefer="threads", verbose=self.verbose
-        )(
+        outputs = Parallel(n_jobs=self.n_jobs, prefer="threads", verbose=self.verbose)(
             delayed(fit_one_parcellation)(
                 X_,
                 Y_,
@@ -393,9 +389,7 @@ def transform(self, X):
 
         X_transform = np.zeros_like(X_)
         for i in range(self.n_bags):
-            X_transform += piecewise_transform(
-                self.labels_[i], self.fit_[i], X_
-            )
+            X_transform += piecewise_transform(self.labels_[i], self.fit_[i], X_)
 
         X_transform /= self.n_bags
 
diff --git a/fmralign/tests/test_alignment_methods.py b/fmralign/tests/test_alignment_methods.py
index 830e479..e0c66c3 100644
--- a/fmralign/tests/test_alignment_methods.py
+++ b/fmralign/tests/test_alignment_methods.py
@@ -30,7 +30,7 @@
     generate_dummy_searchlights,
 )
 
-from hyperalignment.regions import create_parcels_from_labels
+from fmralign.hyperalignment.regions import create_parcels_from_labels
 import os
 import shutil
 
@@ -242,7 +242,6 @@ def test_searchlight_alignment_with_ridge():
     X_pred = model.transform(X_test)
 
     # assert that the saved cached files exist
-    assert_array_almost_equal(X_pred, X_test, decimal=3)
     b = os.path.exists("cache/")
     shutil.rmtree("cache")
     assert b
@@ -267,7 +266,6 @@ def test_parcel_alignment():
     X_pred = model.transform(X_test)
 
     # assert that the saved cached files exist
-    assert_array_almost_equal(X_pred, X_test, decimal=3)
     b = os.path.exists("cache/")
     shutil.rmtree("cache")
     assert b

From 45dd245b20d5f24168f53b9ed3f7976a7cc61586 Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Fri, 5 Jan 2024 14:37:41 +0100
Subject: [PATCH 19/69] fixed doc and argumets

---
 fmralign/alignment_methods.py                 | 160 +++++++++---------
 .../individualized_neural_tuning.py           |  43 ++---
 .../hyperalignment/test_hyperalignment.py     |   8 +-
 fmralign/hyperalignment/toy_experiment.py     |   4 +-
 4 files changed, 111 insertions(+), 104 deletions(-)

diff --git a/fmralign/alignment_methods.py b/fmralign/alignment_methods.py
index c817635..c264dfe 100644
--- a/fmralign/alignment_methods.py
+++ b/fmralign/alignment_methods.py
@@ -471,28 +471,32 @@ def transform(self, X):
 
 
 class IndividualizedNeuralTuning(Alignment):
+    """
+    Method of alignment based on the Individualized Neural Tuning model, by Feilong Ma et al. (2023).
+    It works on 4D fMRI data, and is based on the assumption that the neural response to a stimulus is shared across subjects.
+    It uses searchlight/parcelation alignment to denoise the data, and then computes the stimulus response matrix.
+    See article : https://doi.org/10.1162/imag_a_00032
+    """
+
     def __init__(
         self,
-        tmpl_kind="pca",
+        template="pca",
         decomp_method=None,
-        latent_dim=None,
+        n_components=None,
         alignment_method="searchlight",
         id=None,
         cache=False,
         n_jobs=1,
     ):
         """
-        Method of alignment based on the Individualized Neural Tuning model, by Feilong Ma et al. (2023).
-        It works on 4D fMRI data, and is based on the assumption that the neural response to a stimulus is shared across subjects.
-        It uses searchlight/parcelation alignment to denoise the data, and then computes the stimulus response matrix.
-        See article : https://doi.org/10.1162/imag_a_00032
-
+        Initialize the IndividualizedNeuralTuning object.
         Parameters:
         --------
 
-        - tmpl_kind (str): The type of template to use for alignment. Default is "pca".
+        - template (str): The type of template to use for alignment. Default is "pca".
         - decomp_method (str): The decomposition method to use. Default is None.
-        - latent_dim (int): The number of latent dimensions to use in the shared stimulus information matrix. Default is None.
+        - alignment_method (str): The alignment method to use. Can be either "searchlight" or "parcelation", Default is "searchlight".
+        - n_components (int): The number of latent dimensions to use in the shared stimulus information matrix. Default is None.
         - n_jobs (int): The number of parallel jobs to run. Default is -1.
 
         Returns:
@@ -521,8 +525,8 @@ def __init__(
         self.tuning_data = []
         self.denoised_signal = []
         self.decomp_method = decomp_method
-        self.tmpl_kind = tmpl_kind
-        self.latent_dim = latent_dim
+        self.tmpl_kind = template
+        self.latent_dim = n_components
         self.n_jobs = n_jobs
         self.cache = cache
 
@@ -539,6 +543,68 @@ def __init__(
 
             self.path = path
 
+    #######################################################################################
+    # Computing decomposition
+
+    @staticmethod
+    def _tuning_estimator(shared_response, target):
+        """
+        Estimate the tuning weights for individualized neural tuning.
+
+        Parameters:
+        --------
+            - shared_response (array-like):
+                The shared response matrix.
+            - target (array-like):
+                The target matrix.
+            - latent_dim (int, optional):
+                The number of latent dimensions. Defaults to None.
+
+        Returns:
+        --------
+            array-like: The estimated tuning weights.
+
+        """
+        return np.linalg.pinv(shared_response).dot(target).astype(np.float32)
+
+    @staticmethod
+    def _stimulus_estimator(full_signal, n_t, n_s, latent_dim=None):
+        """
+        Estimates the stimulus response using the given parameters.
+
+        Args:
+            full_signal (np.ndarray): The full signal data.
+            n_t (int): The number of time points.
+            n_s (int): The number of stimuli.
+            latent_dim (int, optional): The number of latent dimensions. Defaults to None.
+
+        Returns:
+            stimulus (np.ndarray): The stimulus response of shape (n_t, latent_dim) or (n_t, n_t).
+        """
+        if latent_dim is not None and latent_dim < n_t:
+            U = svd_pca(full_signal)
+            U = U[:, :latent_dim]
+        else:
+            U, _, _ = safe_svd(full_signal)
+
+        stimulus = np.sqrt(n_s) * U
+        stimulus = stimulus.astype(np.float32)
+        return stimulus
+
+    @staticmethod
+    def _reconstruct_signal(shared_response, individual_tuning):
+        """
+        Reconstructs the signal using the shared response and individual tuning.
+
+        Args:
+            shared_response (numpy.ndarray): The shared response of shape (n_t, n_t) or (n_t, latent_dim).
+            individual_tuning (numpy.ndarray): The individual tuning of shape (latent_dim, n_v) or (n_t, n_v).
+
+        Returns:
+            numpy.ndarray: The reconstructed signal of shape (n_t, n_v) (same shape as the original signal)
+        """
+        return (shared_response @ individual_tuning).astype(np.float32)
+
     def fit(
         self,
         X_train,
@@ -618,15 +684,14 @@ def fit(
         # Stimulus matrix computation
         if self.decomp_method is None:
             full_signal = np.concatenate(self.denoised_signal, axis=1)
-            self.shared_response = _stimulus_estimator(
+            self.shared_response = self._stimulus_estimator(
                 full_signal, self.n_t, self.n_s, self.latent_dim
             )
             if tuning:
                 self.tuning_data = Parallel(n_jobs=self.n_jobs)(
-                    delayed(_tuning_estimator)(
+                    delayed(self._tuning_estimator)(
                         self.shared_response,
                         self.denoised_signal[i],
-                        latent_dim=self.latent_dim,
                     )
                     for i in range(self.n_s)
                 )
@@ -655,7 +720,7 @@ def transform(self, X_test_, verbose=False):
             print("Predict : Computing stimulus matrix...")
 
         if self.decomp_method is None:
-            stimulus_ = _stimulus_estimator(
+            stimulus_ = self._stimulus_estimator(
                 full_signal, self.n_t, self.n_s, self.latent_dim
             )
 
@@ -663,7 +728,7 @@ def transform(self, X_test_, verbose=False):
             print("Predict : stimulus matrix shape: ", stimulus_.shape)
 
         reconstructed_signal = [
-            _reconstruct_signal(stimulus_, T_est) for T_est in self.tuning_data
+            self._reconstruct_signal(stimulus_, T_est) for T_est in self.tuning_data
         ]
         return np.array(reconstructed_signal, dtype=np.float32)
 
@@ -678,66 +743,3 @@ def clean_cache(self, id):
             os.remove("cache")
         except:  # noqa: E722
             print("No cache to remove")
-
-
-#######################################################################################
-# Computing decomposition
-
-
-def _tuning_estimator(shared_response, target, latent_dim=None):
-    """
-    Estimate the tuning weights for individualized neural tuning.
-
-    Parameters:
-    --------
-        - shared_response (array-like):
-            The shared response matrix.
-        - target (array-like):
-            The target matrix.
-        - latent_dim (int, optional):
-            The number of latent dimensions. Defaults to None.
-
-    Returns:
-    --------
-        array-like: The estimated tuning weights.
-
-    """
-    return np.linalg.pinv(shared_response).dot(target).astype(np.float32)
-
-
-def _stimulus_estimator(full_signal, n_t, n_s, latent_dim=None):
-    """
-    Estimates the stimulus response using the given parameters.
-
-    Args:
-        full_signal (np.ndarray): The full signal data.
-        n_t (int): The number of time points.
-        n_s (int): The number of stimuli.
-        latent_dim (int, optional): The number of latent dimensions. Defaults to None.
-
-    Returns:
-        stimulus (np.ndarray): The stimulus response of shape (n_t, latent_dim) or (n_t, n_t).
-    """
-    if latent_dim is not None and latent_dim < n_t:
-        U = svd_pca(full_signal)
-        U = U[:, :latent_dim]
-    else:
-        U, _, _ = safe_svd(full_signal)
-
-    stimulus = np.sqrt(n_s) * U
-    stimulus = stimulus.astype(np.float32)
-    return stimulus
-
-
-def _reconstruct_signal(shared_response, individual_tuning):
-    """
-    Reconstructs the signal using the shared response and individual tuning.
-
-    Args:
-        shared_response (numpy.ndarray): The shared response of shape (n_t, n_t) or (n_t, latent_dim).
-        individual_tuning (numpy.ndarray): The individual tuning of shape (latent_dim, n_v) or (n_t, n_v).
-
-    Returns:
-        numpy.ndarray: The reconstructed signal of shape (n_t, n_v) (same shape as the original signal)
-    """
-    return (shared_response @ individual_tuning).astype(np.float32)
diff --git a/fmralign/hyperalignment/individualized_neural_tuning.py b/fmralign/hyperalignment/individualized_neural_tuning.py
index 47db245..b42f12e 100644
--- a/fmralign/hyperalignment/individualized_neural_tuning.py
+++ b/fmralign/hyperalignment/individualized_neural_tuning.py
@@ -1,6 +1,3 @@
-"""
-A wrapper for the IndividualTuningModel class to be used in fmralign (taking Nifti1 images as input).
-"""
 from fmralign.alignment_methods import IndividualizedNeuralTuning as BaseINT
 from .regions import compute_searchlights, compute_parcels
 from nilearn.maskers import NiftiMasker
@@ -9,40 +6,48 @@
 
 
 class IndividualizedNeuralTuning(BaseINT):
+    """
+    Wrapper for the IndividualTuningModel class to be used in fmralign with Niimg objects.
+    Preprocessing and searchlight/parcellation alignment are done without any user input.
+
+    Method of alignment based on the Individualized Neural Tuning model, by Feilong Ma et al. (2023).
+    It uses searchlight/parcelation alignment to denoise the data, and then computes the stimulus response matrix.
+    See article : https://doi.org/10.1162/imag_a_00032
+    """
+
     def __init__(
         self,
-        tmpl_kind="pca",
+        template="pca",
         decomp_method=None,
         alignment_method="searchlight",
         n_pieces=150,
-        radius=20,
+        searchlight_radius=20,
+        n_components=None,
         n_jobs=1,
     ):
         """
-        A wrapper for the IndividualTuningModel class to be used in fmralign (taking Nifti1 images as input).
-        Method of alignment based on the Individualized Neural Tuning model, by Feilong Ma et al. (2023).
-        It uses searchlight/parcelation alignment to denoise the data, and then computes the stimulus response matrix.
-        See article : https://doi.org/10.1162/imag_a_00032
+        Initialize the IndividualizedNeuralTuning object.
 
         Parameters:
-        --------
-        - method (str): The method used for hyperalignment. Can be either "searchlight" or "parcellation". Default is "searchlight".
+        -----------
+
+        - tmpl_kind (str): The type of template used for alignment. Default is "pca".
+        - decomp_method (str): The decomposition method used for template construction. Default is None.
+        - alignment_method (str): The alignment method used. Default is "searchlight".
+        - n_pieces (int): The number of pieces to divide the data into if using parcelation. Default is 150.
+        - radius (int): The radius of the searchlight sphere in millimeters. Default is 20.
+        - latent_dim (int): The number of latent dimensions to use. Default is None.
         - n_jobs (int): The number of parallel jobs to run. Default is 1.
-        - n_pieces (int): The number of pieces to divide the brain into. Default is 150.
-        - radius (int): The radius of the searchlight in millimeters. Default is 20.
-
-        Returns:
-        --------
-        None
         """
         super().__init__(
-            tmpl_kind=tmpl_kind,
+            template=template,
             decomp_method=decomp_method,
+            n_components=n_components,
             alignment_method=alignment_method,
             n_jobs=n_jobs,
         )
         self.n_pieces = n_pieces
-        self.radius = radius
+        self.radius = searchlight_radius
         self.mask_img = None
         self.masker = None
 
diff --git a/fmralign/hyperalignment/test_hyperalignment.py b/fmralign/hyperalignment/test_hyperalignment.py
index 14a22b8..7057926 100644
--- a/fmralign/hyperalignment/test_hyperalignment.py
+++ b/fmralign/hyperalignment/test_hyperalignment.py
@@ -27,8 +27,8 @@ def test_int_fit_predict():
     dists = [np.ones((300,))]
 
     # Test INT on the two parts of the data (ie different runs of the experiment)
-    int1 = INT(latent_dim=6)
-    int2 = INT(latent_dim=6)
+    int1 = INT(n_components=6)
+    int2 = INT(n_components=6)
     int1.fit(
         X_train=X_train, searchlights=searchlights, dists=dists
     )  # S is provided if we cheat and know the ground truth
@@ -88,8 +88,8 @@ def test_int_with_searchlight():
     )
 
     # Test INT on the two parts of the data (ie different runs of the experiment)
-    int1 = INT(latent_dim=6)
-    int2 = INT(latent_dim=6)
+    int1 = INT(n_components=6)
+    int2 = INT(n_components=6)
     int1.fit(X_train=X_train, searchlights=searchlights, dists=dists, radius=5)
     int2.fit(X_test, searchlights=searchlights, dists=dists, radius=5)
     X_pred = int1.transform(X_test)
diff --git a/fmralign/hyperalignment/toy_experiment.py b/fmralign/hyperalignment/toy_experiment.py
index 9b5350b..d1f29db 100644
--- a/fmralign/hyperalignment/toy_experiment.py
+++ b/fmralign/hyperalignment/toy_experiment.py
@@ -56,8 +56,8 @@
 
 #############################################################################
 # Test INT on the two parts of the data (ie different runs of the experiment)
-int1 = INT(latent_dim=latent_dim, decomp_method=decomposition_method, cache=False)
-int2 = INT(latent_dim=latent_dim, decomp_method=decomposition_method, cache=False)
+int1 = INT(n_components=latent_dim, decomp_method=decomposition_method, cache=False)
+int2 = INT(n_components=latent_dim, decomp_method=decomposition_method, cache=False)
 int_first_part = int1.fit(
     data_run_1,
     searchlights=searchlights,

From 1d431c8e1be82dc6cf52906a694b2edde36d486f Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Fri, 5 Jan 2024 16:00:36 +0100
Subject: [PATCH 20/69] add int plot and better int testing

---
 examples/plot_int_alignment.py           | 62 +++++++++++++++++-------
 fmralign/tests/test_alignment_methods.py |  4 +-
 2 files changed, 46 insertions(+), 20 deletions(-)

diff --git a/examples/plot_int_alignment.py b/examples/plot_int_alignment.py
index b684641..898337d 100644
--- a/examples/plot_int_alignment.py
+++ b/examples/plot_int_alignment.py
@@ -77,6 +77,7 @@
 # We make a single 4D Niimg from our list of 3D filenames
 
 target_train = concat_imgs(target_train)
+target_train_data = masker.transform(target_train)
 target_test = df[df.subject == "sub-07"][df.acquisition == "pa"].path.values
 
 ###############################################################################
@@ -105,17 +106,30 @@
 #
 
 from nilearn.image import index_img
-
 from fmralign.alignment_methods import IndividualizedNeuralTuning
+from fmralign.hyperalignment.regions import compute_parcels, compute_searchlights
 
-from fmralign.hyperalignment.regions import compute_parcels
-
-parcels = compute_parcels(niimg=template_train[0], mask=masker, n_parcels=150, n_jobs=5)
 
 train_index = range(53)
-model = IndividualizedNeuralTuning(n_jobs=5, alignment_method="parcelation")
-model.fit(np.array(masked_imgs)[:, train_index, :], parcels=parcels, verbose=1)
-stimulus = model.shared_response
+model = IndividualizedNeuralTuning(n_jobs=10, alignment_method="parcelation")
+
+if False:  # Use Parcellation
+    parcels = compute_parcels(
+        niimg=template_train[0], mask=masker, n_parcels=1000, n_jobs=5
+    )
+    model.fit(np.array(masked_imgs)[:, train_index, :], parcels=parcels, verbose=False)
+else:
+    _, searchlights, dists = compute_searchlights(
+        niimg=template_train[0], mask_img=masker.mask_img, n_jobs=5
+    )
+    model.fit(
+        np.array(masked_imgs)[:, train_index, :],
+        searchlights=searchlights,
+        dists=dists,
+        verbose=False,
+    )
+
+train_stimulus = np.copy(model.shared_response)
 
 
 ###############################################################################
@@ -127,19 +141,27 @@
 # For each train subject and for the template, the AP contrasts are sorted from
 # 0, to 53, and then the PA contrasts from 53 to 106.
 #
-
-train_index = range(53)
 test_index = range(53, 106)
+if False:
+    model.fit(np.array(masked_imgs)[:, test_index, :], parcels=parcels, verbose=False)
+    test_stimulus = np.copy(model.shared_response)
+
+else:
+    model.fit(
+        np.array(masked_imgs)[:, test_index, :],
+        searchlights=searchlights,
+        dists=dists,
+        verbose=False,
+    )
+    test_stimulus = np.copy(model.shared_response)
+
 
 # We input the mapping image target_train in a list, we could have input more
 # than one subject for which we'd want to predict : [train_1, train_2 ...]
-target_train_array = np.array([masker.transform(target_train)[train_index]])
 
-prediction_from_template = model.fit(target_train_array, parcels=parcels, verbose=1)
-tuning_target = model.tuning_data
+tuning_target = np.linalg.pinv(train_stimulus[train_index, :]) @ target_train_data
 
-prediction_from_template = stimulus @ tuning_target[0]
-prediction_from_template = prediction_from_template[test_index, :]
+prediction_from_template = test_stimulus @ tuning_target
 prediction_from_template = [masker.inverse_transform(prediction_from_template)]
 
 # As a baseline prediction, let's just take the average of activations across subjects.
@@ -161,10 +183,14 @@
 # made from group average and from template with the real PA contrasts of sub-07
 
 average_score = masker.inverse_transform(
-    score_voxelwise(target_test, prediction_from_average, masker, loss="corr")
+    np.abs(score_voxelwise(target_test, prediction_from_average, masker, loss="corr"))
 )
+
+# I choose abs value in reference to the work we did with the INT
 template_score = masker.inverse_transform(
-    score_voxelwise(target_test, prediction_from_template[0], masker, loss="corr")
+    np.abs(
+        score_voxelwise(target_test, prediction_from_template[0], masker, loss="corr")
+    )
 )
 
 
@@ -177,11 +203,11 @@
 from nilearn import plotting
 
 baseline_display = plotting.plot_stat_map(
-    average_score, display_mode="z", vmax=1, cut_coords=[-15, -5]
+    average_score, display_mode="z", vmax=1, cut_coords=[-15, -5], cmap="hot"
 )
 baseline_display.title("Group average correlation wt ground truth")
 display = plotting.plot_stat_map(
-    template_score, display_mode="z", cut_coords=[-15, -5], vmax=1
+    template_score, display_mode="z", cut_coords=[-15, -5], vmax=1, cmap="hot"
 )
 display.title("Template-based prediction correlation wt ground truth")
 
diff --git a/fmralign/tests/test_alignment_methods.py b/fmralign/tests/test_alignment_methods.py
index e0c66c3..5ac2960 100644
--- a/fmralign/tests/test_alignment_methods.py
+++ b/fmralign/tests/test_alignment_methods.py
@@ -237,7 +237,7 @@ def test_searchlight_alignment_with_ridge():
         n_t=n_time_points, n_v=n_voxels, n_s=n_subjects
     )
 
-    model = INT()
+    model = INT(n_jobs=5, cache=True)
     model.fit(X_train, searchlights, dists, radius=radius)
     X_pred = model.transform(X_test)
 
@@ -261,7 +261,7 @@ def test_parcel_alignment():
         n_t=n_time_points, n_v=n_voxels, n_s=n_subjects
     )
 
-    model = INT(n_jobs=-1, alignment_method="parcel")
+    model = INT(n_jobs=5, alignment_method="parcel", cache=True)
     model.fit(X_train, parcels=parcels)
     X_pred = model.transform(X_test)
 

From a74f0a89b60e7136f0a3fc364edffeae39bdb048 Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Fri, 5 Jan 2024 16:07:29 +0100
Subject: [PATCH 21/69] citing Feilong

---
 examples/plot_int_alignment.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/examples/plot_int_alignment.py b/examples/plot_int_alignment.py
index 898337d..ccb1dff 100644
--- a/examples/plot_int_alignment.py
+++ b/examples/plot_int_alignment.py
@@ -1,7 +1,9 @@
 # -*- coding: utf-8 -*-
 
 """
-Template-based prediction.
+Hyperalignment-base prediction using Feilong Ma's IndividualNeuralTuning Model.
+See article : https://doi.org/10.1162/imag_a_00032
+
 ==========================
 
 In this tutorial, we show how to better predict new contrasts for a target
@@ -209,7 +211,7 @@
 display = plotting.plot_stat_map(
     template_score, display_mode="z", cut_coords=[-15, -5], vmax=1, cmap="hot"
 )
-display.title("Template-based prediction correlation wt ground truth")
+display.title("Hyperalignment-based prediction correlation wt ground truth")
 
 ###############################################################################
 # We observe that creating a template and aligning a new subject to it yields

From 40af9b873742bb553e4bc26dc8578c01218fccbb Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Fri, 5 Jan 2024 16:23:15 +0100
Subject: [PATCH 22/69] remove caching (might add joblib caching later)

---
 fmralign/alignment_methods.py                |  64 +++--------
 fmralign/hyperalignment/regions_alignment.py | 105 ++++---------------
 2 files changed, 34 insertions(+), 135 deletions(-)

diff --git a/fmralign/alignment_methods.py b/fmralign/alignment_methods.py
index c264dfe..1497ec8 100644
--- a/fmralign/alignment_methods.py
+++ b/fmralign/alignment_methods.py
@@ -11,7 +11,6 @@
 from sklearn.base import BaseEstimator, TransformerMixin
 from sklearn.linear_model import RidgeCV
 from sklearn.metrics.pairwise import pairwise_distances
-import os
 from .hyperalignment.regions_alignment import RegionAlignment
 from .hyperalignment.linalg import safe_svd, svd_pca
 
@@ -484,8 +483,6 @@ def __init__(
         decomp_method=None,
         n_components=None,
         alignment_method="searchlight",
-        id=None,
-        cache=False,
         n_jobs=1,
     ):
         """
@@ -528,20 +525,6 @@ def __init__(
         self.tmpl_kind = template
         self.latent_dim = n_components
         self.n_jobs = n_jobs
-        self.cache = cache
-
-        if cache:
-            if id is None:
-                self.id = "default"
-            else:
-                self.id = id
-
-            path = os.path.join(os.getcwd(), f"cache/int/{self.id}")
-            # Check if cache folder exists
-            if not os.path.exists(path):
-                os.makedirs(path)
-
-            self.path = path
 
     #######################################################################################
     # Computing decomposition
@@ -659,27 +642,20 @@ def fit(
             self.distances = dists
             self.radius = radius
 
-        # check for cached data
-        try:
-            self.denoised_signal = np.load(self.path + "/train_data_denoised.npy")
-            if verbose:
-                print("Loaded denoised data from cache")
-
-        except:  # noqa: E722
-            denoiser = RegionAlignment(
-                alignment_method=self.alignment_method,
-                n_jobs=self.n_jobs,
-                verbose=verbose,
-                path=self.path,
-            )
-            self.denoised_signal = denoiser.fit_transform(
-                X_train_,
-                regions=self.regions,
-                dists=dists,
-                radius=radius,
-            )
-            # Clear memory of the SearchlightAlignment object
-            denoiser = None
+        denoiser = RegionAlignment(
+            alignment_method=self.alignment_method,
+            n_jobs=self.n_jobs,
+            verbose=verbose,
+            path=self.path,
+        )
+        self.denoised_signal = denoiser.fit_transform(
+            X_train_,
+            regions=self.regions,
+            dists=dists,
+            radius=radius,
+        )
+        # Clear memory of the SearchlightAlignment object
+        denoiser = None
 
         # Stimulus matrix computation
         if self.decomp_method is None:
@@ -731,15 +707,3 @@ def transform(self, X_test_, verbose=False):
             self._reconstruct_signal(stimulus_, T_est) for T_est in self.tuning_data
         ]
         return np.array(reconstructed_signal, dtype=np.float32)
-
-    def clean_cache(self, id):
-        """
-        Removes the cache file associated with the given ID.
-
-        Args:
-            id (int): The ID of the cache file to be removed.
-        """
-        try:
-            os.remove("cache")
-        except:  # noqa: E722
-            print("No cache to remove")
diff --git a/fmralign/hyperalignment/regions_alignment.py b/fmralign/hyperalignment/regions_alignment.py
index 4ba62bf..f9c389b 100644
--- a/fmralign/hyperalignment/regions_alignment.py
+++ b/fmralign/hyperalignment/regions_alignment.py
@@ -5,7 +5,6 @@
     piece_ridge,
     searchlight_weights,
 )
-import os
 from joblib import Parallel, delayed
 
 
@@ -34,7 +33,6 @@ def __init__(
         alignment_method="searchlight_ridge",
         template_kind="searchlight_pca",
         verbose=True,
-        path="cache/int/",
         cache=True,
         n_jobs=-1,
     ):
@@ -51,13 +49,6 @@ def __init__(
         self.distances = None
         self.radius = None
         self.weights = None
-        self.path = path
-        self.cache = (path is not None) and (cache)
-
-        if self.cache:
-            # Check if cache folder exists
-            if not os.path.exists(self.path):
-                os.makedirs(self.path)
 
     def compute_linear_transformation(self, x_i, template, i: int = 0, save=True):
         """Compute the linear transformation W_i for a given subject.
@@ -75,28 +66,14 @@ def compute_linear_transformation(self, x_i, template, i: int = 0, save=True):
         Xhat : ndarray of shape (n_samples, n_voxels)
             The denoised estimation signal for each subject.
         """
-        try:
-            W_p = np.load(self.path + (f"/train_data_W_{i}.npy"))
-            if self.verbose:
-                print(f"Loaded W_{i} from cache")
-            x_hat = template.dot(W_p)
-            del W_p
-            return x_hat
-
-        except:  # noqa E722
-            if self.verbose:
-                print(f"No cache found, computing W_{i}")
-
-            x_hat = piece_ridge(
-                X=x_i,
-                Y=template,
-                regions=self.regions,
-                weights=self.weights,
-                verbose=self.verbose,
-            )
 
-            if self.cache:
-                np.save(self.path + (f"/train_data_W_{i}.npy"), x_hat)
+        x_hat = piece_ridge(
+            X=x_i,
+            Y=template,
+            regions=self.regions,
+            weights=self.weights,
+            verbose=self.verbose,
+        )
         return x_hat
 
     def fit_transform(
@@ -135,15 +112,6 @@ def fit_transform(
         if self.verbose:
             print(f"[{self.FUNC}] Shape of input data: ", X.shape)
 
-        try:
-            self.Xhat = np.load(self.path + "/train_data_denoised.npy")
-            if self.verbose:
-                print(f"[{self.FUNC}] Loaded denoised data from cache")
-            return self.Xhat
-        except:  # noqa E722
-            if self.verbose:
-                print(f"[{self.FUNC}] No cache found, computing denoised data")
-
         self.n_s, self.n_t, self.n_v = X.shape
         self.regions = regions
         self.FUNC = "ParcelAlignment"
@@ -157,56 +125,23 @@ def fit_transform(
         if self.verbose:
             print(f"[{self.FUNC}]Computing global template M ...")
 
-        try:
-            sl_template = np.load(self.path + ("/train_data_template.npy"))
-            if self.verbose:
-                print("Loaded template from cache")
-        except:  # noqa E722
-            if self.verbose:
-                print(f"[{self.FUNC}] No cache found, computing template")
-            if dists is None or radius is None:
-                self.weights = None
-            else:
-                self.weights = searchlight_weights(
-                    searchlights=regions, dists=dists, radius=radius
-                )
-            sl_template = template(
-                X,
-                regions=regions,
-                n_jobs=self.n_jobs,
-                template_kind=self.template_kind,
-                verbose=self.verbose,
-                weights=self.weights,
+        if dists is None or radius is None:
+            self.weights = None
+        else:
+            self.weights = searchlight_weights(
+                searchlights=regions, dists=dists, radius=radius
             )
-
-        if self.cache:
-            np.save(self.path + ("/train_data_template.npy"), sl_template)
-            if self.verbose:
-                print(f"[{self.FUNC}] Saved template to cache")
+        sl_template = template(
+            X,
+            regions=regions,
+            n_jobs=self.n_jobs,
+            template_kind=self.template_kind,
+            verbose=self.verbose,
+            weights=self.weights,
+        )
 
         self.Xhat = Parallel(n_jobs=self.n_jobs)(
             delayed(self.compute_linear_transformation)(X[i], sl_template, i)
             for i in range(self.n_s)
         )
-
-        if id is None:
-            id = np.random.randint(0, 1000000)
-
-        if self.cache:
-            np.save(self.path + ("/train_data_denoised.npy"), self.Xhat)
-
         return np.array(self.Xhat)
-
-    def get_linear_transformations(self):
-        """Return the linear transformations W_1, ... W_p for each subject.
-        ----------
-        Returns
-        -------
-        W : list of ndarray of shape (n_voxels, n_voxels)
-            The linear transformations W_1, ... W_p for each subject.
-        """
-        return np.array(self.W)
-
-    def get_denoised_estimation(self):
-        """Return the denoised estimations B_1, ... B_p for each subject."""
-        return self.Xhat

From c788ffa11e5a26537c1d03f7795331f4509589d6 Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Fri, 5 Jan 2024 16:35:19 +0100
Subject: [PATCH 23/69] better naming and doc

---
 fmralign/alignment_methods.py                 |  4 +-
 fmralign/hyperalignment/correlation.py        |  3 ++
 ...ns_alignment.py => piecewise_alignment.py} | 47 +++++++++++--------
 fmralign/hyperalignment/regions.py            |  6 ++-
 fmralign/hyperalignment/toy_experiment.py     |  4 +-
 fmralign/tests/test_alignment_methods.py      | 21 +--------
 6 files changed, 41 insertions(+), 44 deletions(-)
 rename fmralign/hyperalignment/{regions_alignment.py => piecewise_alignment.py} (73%)

diff --git a/fmralign/alignment_methods.py b/fmralign/alignment_methods.py
index 1497ec8..e8701f2 100644
--- a/fmralign/alignment_methods.py
+++ b/fmralign/alignment_methods.py
@@ -11,7 +11,7 @@
 from sklearn.base import BaseEstimator, TransformerMixin
 from sklearn.linear_model import RidgeCV
 from sklearn.metrics.pairwise import pairwise_distances
-from .hyperalignment.regions_alignment import RegionAlignment
+from .hyperalignment.piecewise_alignment import PiecewiseAlignment
 from .hyperalignment.linalg import safe_svd, svd_pca
 
 import jax
@@ -642,7 +642,7 @@ def fit(
             self.distances = dists
             self.radius = radius
 
-        denoiser = RegionAlignment(
+        denoiser = PiecewiseAlignment(
             alignment_method=self.alignment_method,
             n_jobs=self.n_jobs,
             verbose=verbose,
diff --git a/fmralign/hyperalignment/correlation.py b/fmralign/hyperalignment/correlation.py
index 21ba050..d1501ac 100644
--- a/fmralign/hyperalignment/correlation.py
+++ b/fmralign/hyperalignment/correlation.py
@@ -1,3 +1,6 @@
+"""Some tools to compute correlation matrices. Functions in this module are
+meant to be used as a test for the hyperalignment algorithm only."""
+
 import numpy as np
 from scipy.stats import spearmanr
 from sklearn.metrics import pairwise_distances
diff --git a/fmralign/hyperalignment/regions_alignment.py b/fmralign/hyperalignment/piecewise_alignment.py
similarity index 73%
rename from fmralign/hyperalignment/regions_alignment.py
rename to fmralign/hyperalignment/piecewise_alignment.py
index f9c389b..b350c2d 100644
--- a/fmralign/hyperalignment/regions_alignment.py
+++ b/fmralign/hyperalignment/piecewise_alignment.py
@@ -1,3 +1,9 @@
+"""Piecewise alignment model. This model decomposes the data into regions (pieces).
+Those can either be searchlights or parcels (computed with standard parcellation algorithms).
+See the ```nilearn``` documentation for more details:
+- https://nilearn.github.io/modules/generated/nilearn.regions.Parcellations.html
+- https://nilearn.github.io/dev/modules/generated/nilearn.decoding.SearchLight.html
+"""
 import numpy as np
 from sklearn.base import BaseEstimator, TransformerMixin
 from .regions import (
@@ -8,24 +14,13 @@
 from joblib import Parallel, delayed
 
 
-class RegionAlignment(BaseEstimator, TransformerMixin):
+class PiecewiseAlignment(BaseEstimator, TransformerMixin):
     """Searchlight alignment model. This model decomposes the data into a
     global template and a linear transformation for each subject.
-    The global template is computed using a searchlight approach.
+    The global template is computed using a searchlight/parcellation approach.
     The linear transformation is computed using a ridge regression.
     This step is enssential to the hyperalignment model, as it is
-    used as a denoiser for the data.
-    Parameters
-    ----------
-    alignment_method : str, default="ridge"
-        The alignment method to use. Can be "ridge" or "ensemble_ridge".
-    template_kind : str, default="pca"
-        The kind of template to use. Can be "pca" or "mean".
-    demean : bool, default=False
-        Whether to demean the data before alignment.
-    verbose : bool, default=True
-        Whether to display progress bar.
-    n_jobs : int, default=-1
+    used to remove noise from the raw data.
     """
 
     def __init__(
@@ -33,9 +28,21 @@ def __init__(
         alignment_method="searchlight_ridge",
         template_kind="searchlight_pca",
         verbose=True,
-        cache=True,
         n_jobs=-1,
     ):
+        """
+        Parameters
+        ----------
+        alignment_method : str, default="ridge"
+            The alignment method to use. Can be "ridge" or "ensemble_ridge".
+        template_kind : str, default="pca"
+            The kind of template to use. Can be "pca" or "mean".
+        demean : bool, default=False
+            Whether to demean the data before alignment.
+        verbose : bool, default=True
+            Whether to display progress bar.
+        n_jobs : int, default=-1
+        """
         self.W = []
         self.Xhat = []
         self.n_s = None
@@ -50,12 +57,12 @@ def __init__(
         self.radius = None
         self.weights = None
 
-    def compute_linear_transformation(self, x_i, template, i: int = 0, save=True):
-        """Compute the linear transformation W_i for a given subject.
-        ----------
+    def compute_linear_transformation(self, data, template, i: int = 0, save=True):
+        """Compute the linear transformation for a given subject provided the global template.
+
         Parameters
         ----------
-        x_i : ndarray of shape (n_samples, n_voxels)
+        data : ndarray of shape (n_samples, n_voxels)
             The brain images for one subject.
             Those are the B_1, ..., B_n in the paper.
         template : ndarray of shape (n_samples, n_voxels)
@@ -68,7 +75,7 @@ def compute_linear_transformation(self, x_i, template, i: int = 0, save=True):
         """
 
         x_hat = piece_ridge(
-            X=x_i,
+            X=data,
             Y=template,
             regions=self.regions,
             weights=self.weights,
diff --git a/fmralign/hyperalignment/regions.py b/fmralign/hyperalignment/regions.py
index 8b795c1..03f1061 100644
--- a/fmralign/hyperalignment/regions.py
+++ b/fmralign/hyperalignment/regions.py
@@ -1,4 +1,8 @@
-"""Utilities for computing searchlights. Adapted from nilearn.\n
+"""Utilities for computing searchlights. Adapted from ```nilearn```.\n
+See the ```nilearn``` documentation for more details:
+- https://nilearn.github.io/modules/generated/nilearn.regions.Parcellations.html
+- https://nilearn.github.io/dev/modules/generated/nilearn.decoding.SearchLight.html
+
 Author: Denis Fouchard, INRIA Saclay, MIND, 2023.
 """
 
diff --git a/fmralign/hyperalignment/toy_experiment.py b/fmralign/hyperalignment/toy_experiment.py
index d1f29db..982993a 100644
--- a/fmralign/hyperalignment/toy_experiment.py
+++ b/fmralign/hyperalignment/toy_experiment.py
@@ -56,8 +56,8 @@
 
 #############################################################################
 # Test INT on the two parts of the data (ie different runs of the experiment)
-int1 = INT(n_components=latent_dim, decomp_method=decomposition_method, cache=False)
-int2 = INT(n_components=latent_dim, decomp_method=decomposition_method, cache=False)
+int1 = INT(n_components=latent_dim, decomp_method=decomposition_method)
+int2 = INT(n_components=latent_dim, decomp_method=decomposition_method)
 int_first_part = int1.fit(
     data_run_1,
     searchlights=searchlights,
diff --git a/fmralign/tests/test_alignment_methods.py b/fmralign/tests/test_alignment_methods.py
index 5ac2960..2d18654 100644
--- a/fmralign/tests/test_alignment_methods.py
+++ b/fmralign/tests/test_alignment_methods.py
@@ -31,8 +31,6 @@
 )
 
 from fmralign.hyperalignment.regions import create_parcels_from_labels
-import os
-import shutil
 
 
 def test_scaled_procrustes_algorithmic():
@@ -237,14 +235,9 @@ def test_searchlight_alignment_with_ridge():
         n_t=n_time_points, n_v=n_voxels, n_s=n_subjects
     )
 
-    model = INT(n_jobs=5, cache=True)
+    model = INT(n_jobs=5)
     model.fit(X_train, searchlights, dists, radius=radius)
     X_pred = model.transform(X_test)
-
-    # assert that the saved cached files exist
-    b = os.path.exists("cache/")
-    shutil.rmtree("cache")
-    assert b
     assert X_pred.shape == X_test.shape
 
 
@@ -261,17 +254,7 @@ def test_parcel_alignment():
         n_t=n_time_points, n_v=n_voxels, n_s=n_subjects
     )
 
-    model = INT(n_jobs=5, alignment_method="parcel", cache=True)
+    model = INT(n_jobs=5, alignment_method="parcel")
     model.fit(X_train, parcels=parcels)
     X_pred = model.transform(X_test)
-
-    # assert that the saved cached files exist
-    b = os.path.exists("cache/")
-    shutil.rmtree("cache")
-    assert b
     assert X_pred.shape == X_test.shape
-
-
-# Cleaning up the cache folder
-if os.path.exists("cache/"):
-    shutil.rmtree("cache")

From 9cb5e50cc94e9c7f8cd2d031d730a1c6a30839be Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Fri, 5 Jan 2024 16:43:06 +0100
Subject: [PATCH 24/69] improve doc

---
 fmralign/hyperalignment/correlation.py        | 12 +++++++----
 .../individualized_neural_tuning.py           | 21 ++++++++++++-------
 .../hyperalignment/piecewise_alignment.py     |  6 +++---
 3 files changed, 25 insertions(+), 14 deletions(-)

diff --git a/fmralign/hyperalignment/correlation.py b/fmralign/hyperalignment/correlation.py
index d1501ac..c34cdfe 100644
--- a/fmralign/hyperalignment/correlation.py
+++ b/fmralign/hyperalignment/correlation.py
@@ -201,10 +201,14 @@ def thread_compute_correlation(X, Y, i, j):
     Compute the correlation between two time series X_i and Y_i.
 
     Parameters:
-       - X (ndarray): Array of shape (n_samples, n_features) representing the first time series.
-       - Y (ndarray): Array of shape (n_samples, n_features) representing the second time series.
-       - i (int): Index of the first time series.
-       - j (int): Index of the second time series.
+    - X (ndarray):
+        ndrray of shape (n_samples, n_features) representing the first time series.
+    - Y (ndarray):
+        ndrray of shape (n_samples, n_features) representing the second time series.
+    - i (int):
+        Index of the first time series.
+    - j (int):
+        Index of the second time series.
 
     Returns:
     diff_TR_corr (ndarray): Array of shape (n_samples * (n_samples - 1),) containing the correlations between different time points.
diff --git a/fmralign/hyperalignment/individualized_neural_tuning.py b/fmralign/hyperalignment/individualized_neural_tuning.py
index b42f12e..273b0ed 100644
--- a/fmralign/hyperalignment/individualized_neural_tuning.py
+++ b/fmralign/hyperalignment/individualized_neural_tuning.py
@@ -30,13 +30,20 @@ def __init__(
 
         Parameters:
         -----------
-
-        - tmpl_kind (str): The type of template used for alignment. Default is "pca".
-        - decomp_method (str): The decomposition method used for template construction. Default is None.
-        - alignment_method (str): The alignment method used. Default is "searchlight".
-        - n_pieces (int): The number of pieces to divide the data into if using parcelation. Default is 150.
-        - radius (int): The radius of the searchlight sphere in millimeters. Default is 20.
-        - latent_dim (int): The number of latent dimensions to use. Default is None.
+        - tmpl_kind (str):
+            The type of template used for alignment. Default is "pca".
+        - decomp_method (str):
+            The decomposition method used for template construction. Default is None.
+        - alignment_method (str) in ["searchlight", "parcellation"]:
+            The alignment method used. Default is "searchlight".
+        - n_pieces (int): The number of pieces to divide the data into if using parcellation. Default is 150.
+        - radius (int):
+            The radius of the searchlight sphere in millimeters.
+            Only used if alignment_method is "searchlight".
+            Default is 20.
+        - n_components (int):
+            The number of latent dimensions to use. If None, all the components are used.
+            Default is None.
         - n_jobs (int): The number of parallel jobs to run. Default is 1.
         """
         super().__init__(
diff --git a/fmralign/hyperalignment/piecewise_alignment.py b/fmralign/hyperalignment/piecewise_alignment.py
index b350c2d..0c8cb67 100644
--- a/fmralign/hyperalignment/piecewise_alignment.py
+++ b/fmralign/hyperalignment/piecewise_alignment.py
@@ -28,7 +28,7 @@ def __init__(
         alignment_method="searchlight_ridge",
         template_kind="searchlight_pca",
         verbose=True,
-        n_jobs=-1,
+        n_jobs=1,
     ):
         """
         Parameters
@@ -92,8 +92,8 @@ def fit_transform(
         weights=None,
         id=None,
     ):
-        """From brain imgs compute the INT model (M, Ws, S)
-          with the given parameters)
+        """From given fmri data, compute the global template and the linear transformation.
+        This provides denoised signal estimations using template alignment.
 
         Parameters
         ----------

From f6a3950fe132913ed27eb9f4e8af0c2b73aa1b51 Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Fri, 5 Jan 2024 16:49:15 +0100
Subject: [PATCH 25/69] fix tests

---
 fmralign/alignment_methods.py            | 5 +----
 fmralign/tests/test_alignment_methods.py | 2 +-
 2 files changed, 2 insertions(+), 5 deletions(-)

diff --git a/fmralign/alignment_methods.py b/fmralign/alignment_methods.py
index e8701f2..34d123f 100644
--- a/fmralign/alignment_methods.py
+++ b/fmralign/alignment_methods.py
@@ -643,10 +643,7 @@ def fit(
             self.radius = radius
 
         denoiser = PiecewiseAlignment(
-            alignment_method=self.alignment_method,
-            n_jobs=self.n_jobs,
-            verbose=verbose,
-            path=self.path,
+            alignment_method=self.alignment_method, n_jobs=self.n_jobs, verbose=verbose
         )
         self.denoised_signal = denoiser.fit_transform(
             X_train_,
diff --git a/fmralign/tests/test_alignment_methods.py b/fmralign/tests/test_alignment_methods.py
index 2d18654..1e936b9 100644
--- a/fmralign/tests/test_alignment_methods.py
+++ b/fmralign/tests/test_alignment_methods.py
@@ -236,7 +236,7 @@ def test_searchlight_alignment_with_ridge():
     )
 
     model = INT(n_jobs=5)
-    model.fit(X_train, searchlights, dists, radius=radius)
+    model.fit(X_train, searchlights=searchlights, dists=dists, radius=radius)
     X_pred = model.transform(X_test)
     assert X_pred.shape == X_test.shape
 

From 2fd46d38910eaa021df2216bdc7f1847a17e696f Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Tue, 9 Jan 2024 11:21:34 +0100
Subject: [PATCH 26/69] Delete any reference to Feilong Ma #streisandeffect

---
 .gitignore                                    |  1 +
 examples/plot_int_alignment.py                | 22 +++++++++----------
 fmralign/alignment_methods.py                 |  2 +-
 fmralign/generate_data.py                     |  9 +++++++-
 .../individualized_neural_tuning.py           |  2 +-
 fmralign/hyperalignment/local_template.py     |  1 -
 fmralign/hyperalignment/toy_experiment.py     | 18 ++++++++++++---
 7 files changed, 37 insertions(+), 18 deletions(-)

diff --git a/.gitignore b/.gitignore
index e479904..d023765 100644
--- a/.gitignore
+++ b/.gitignore
@@ -93,3 +93,4 @@ venv.bak/
 
 # Rope project settings
 .ropeproject
+examples/plot_srm_alignment.py
diff --git a/examples/plot_int_alignment.py b/examples/plot_int_alignment.py
index ccb1dff..eab04a9 100644
--- a/examples/plot_int_alignment.py
+++ b/examples/plot_int_alignment.py
@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 
 """
-Hyperalignment-base prediction using Feilong Ma's IndividualNeuralTuning Model.
+Hyperalignment-base prediction using the IndividualNeuralTuning Model.
 See article : https://doi.org/10.1162/imag_a_00032
 
 ==========================
@@ -113,11 +113,13 @@
 
 
 train_index = range(53)
-model = IndividualizedNeuralTuning(n_jobs=10, alignment_method="parcelation")
+model = IndividualizedNeuralTuning(
+    n_jobs=10, alignment_method="parcelation", n_components=20
+)
 
-if False:  # Use Parcellation
+if True:  # Use Parcellation
     parcels = compute_parcels(
-        niimg=template_train[0], mask=masker, n_parcels=1000, n_jobs=5
+        niimg=template_train[0], mask=masker, n_parcels=200, n_jobs=5
     )
     model.fit(np.array(masked_imgs)[:, train_index, :], parcels=parcels, verbose=False)
 else:
@@ -144,7 +146,7 @@
 # 0, to 53, and then the PA contrasts from 53 to 106.
 #
 test_index = range(53, 106)
-if False:
+if True:
     model.fit(np.array(masked_imgs)[:, test_index, :], parcels=parcels, verbose=False)
     test_stimulus = np.copy(model.shared_response)
 
@@ -185,14 +187,12 @@
 # made from group average and from template with the real PA contrasts of sub-07
 
 average_score = masker.inverse_transform(
-    np.abs(score_voxelwise(target_test, prediction_from_average, masker, loss="corr"))
+    score_voxelwise(target_test, prediction_from_average, masker, loss="corr")
 )
 
 # I choose abs value in reference to the work we did with the INT
 template_score = masker.inverse_transform(
-    np.abs(
-        score_voxelwise(target_test, prediction_from_template[0], masker, loss="corr")
-    )
+    score_voxelwise(target_test, prediction_from_template[0], masker, loss="corr")
 )
 
 
@@ -205,11 +205,11 @@
 from nilearn import plotting
 
 baseline_display = plotting.plot_stat_map(
-    average_score, display_mode="z", vmax=1, cut_coords=[-15, -5], cmap="hot"
+    average_score, display_mode="z", vmax=1, cut_coords=[-15, -5]
 )
 baseline_display.title("Group average correlation wt ground truth")
 display = plotting.plot_stat_map(
-    template_score, display_mode="z", cut_coords=[-15, -5], vmax=1, cmap="hot"
+    template_score, display_mode="z", cut_coords=[-15, -5], vmax=1
 )
 display.title("Hyperalignment-based prediction correlation wt ground truth")
 
diff --git a/fmralign/alignment_methods.py b/fmralign/alignment_methods.py
index 34d123f..87f9803 100644
--- a/fmralign/alignment_methods.py
+++ b/fmralign/alignment_methods.py
@@ -548,7 +548,7 @@ def _tuning_estimator(shared_response, target):
             array-like: The estimated tuning weights.
 
         """
-        return np.linalg.pinv(shared_response).dot(target).astype(np.float32)
+        return np.linalg.pinv(shared_response) @ target
 
     @staticmethod
     def _stimulus_estimator(full_signal, n_t, n_s, latent_dim=None):
diff --git a/fmralign/generate_data.py b/fmralign/generate_data.py
index d85f26b..4a6c9a4 100644
--- a/fmralign/generate_data.py
+++ b/fmralign/generate_data.py
@@ -77,11 +77,18 @@ def generate_dummy_signal(
         S_train = np.sqrt(Sigma)[:, None] * rng.randn(n_t, latent_dim)
         S_test = np.sqrt(Sigma)[:, None] * rng.randn(n_t, latent_dim)
 
+    elif generative_method == "multiviewica":
+        S_train = np.random.laplace(size=(n_t, latent_dim))
+        S_test = np.random.laplace(size=(n_t, latent_dim))
+
+    else:
+        raise ValueError("Unknown generative method")
+
     # Generate indiivdual spatial components
     data_train, data_test = [], []
     Ts = []
     for _ in range(n_s):
-        if generative_method == "custom":
+        if generative_method == "custom" or generative_method == "multiviewica":
             W = T_mean + T_std * np.random.randn(latent_dim, n_v)
         else:
             W = projection(rng.randn(latent_dim, n_v))
diff --git a/fmralign/hyperalignment/individualized_neural_tuning.py b/fmralign/hyperalignment/individualized_neural_tuning.py
index 273b0ed..6215db8 100644
--- a/fmralign/hyperalignment/individualized_neural_tuning.py
+++ b/fmralign/hyperalignment/individualized_neural_tuning.py
@@ -10,7 +10,7 @@ class IndividualizedNeuralTuning(BaseINT):
     Wrapper for the IndividualTuningModel class to be used in fmralign with Niimg objects.
     Preprocessing and searchlight/parcellation alignment are done without any user input.
 
-    Method of alignment based on the Individualized Neural Tuning model, by Feilong Ma et al. (2023).
+    Method of alignment based on the Individualized Neural Tuning model.
     It uses searchlight/parcelation alignment to denoise the data, and then computes the stimulus response matrix.
     See article : https://doi.org/10.1162/imag_a_00032
     """
diff --git a/fmralign/hyperalignment/local_template.py b/fmralign/hyperalignment/local_template.py
index c16a92c..5b1aef2 100644
--- a/fmralign/hyperalignment/local_template.py
+++ b/fmralign/hyperalignment/local_template.py
@@ -1,7 +1,6 @@
 """ Local template computation functions. Those functions are part of the warp hyperalignment
 introducted by Feilong Ma et al. 2023. The functions are adapted from the original code and
 adapted for more general regionations approches.
-Authors: Feilong Ma (Haxby lab, Dartmouth College), Denis Fouchard (MIND, INRIA Saclay).
 """
 
 
diff --git a/fmralign/hyperalignment/toy_experiment.py b/fmralign/hyperalignment/toy_experiment.py
index 982993a..c8f331f 100644
--- a/fmralign/hyperalignment/toy_experiment.py
+++ b/fmralign/hyperalignment/toy_experiment.py
@@ -23,12 +23,12 @@
 
 
 n_s = 10
-n_t = 200
+n_t = 100
 n_v = 200
 S_std = 1
 T_std = 1
 SNR = 100
-latent_dim = 6  # if None, latent_dim = n_t
+latent_dim = None  # if None, latent_dim = n_t
 decomposition_method = None  # if None, SVD is used
 
 
@@ -50,9 +50,18 @@
     latent_dim=latent_dim,
     SNR=SNR,
     seed=0,
+    generative_method="fastsrm",
 )
 
-searchlights, dists = generate_dummy_searchlights(n_searchlights=12, n_v=n_v, radius=5)
+SEARCHLIGHT = False
+
+if SEARCHLIGHT:
+    searchlights, dists = generate_dummy_searchlights(
+        n_searchlights=12, n_v=n_v, radius=5
+    )
+else:
+    searchlights = [np.arange(n_v)]
+    dists = [np.zeros((n_v,))]
 
 #############################################################################
 # Test INT on the two parts of the data (ie different runs of the experiment)
@@ -139,6 +148,9 @@
 
 # Stimulus matrix
 
+print("stimulus_pred_run_1.shape", stimulus_pred_run_1.shape)
+print("stimulus_run_1.shape", stimulus_run_1.shape)
+
 correlation_stimulus_true_est_first_part = stimulus_correlation(
     stimulus_pred_run_1.T, stimulus_run_1.T
 )

From e35c01b152b6fff8f154f4b2d44af45fdd95273c Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Wed, 10 Jan 2024 11:31:33 +0100
Subject: [PATCH 27/69] bug fixes + rm tqdm dependency

---
 examples/plot_int_alignment.py                | 42 ++++++------
 fmralign/alignment_methods.py                 |  2 +-
 fmralign/hyperalignment/correlation.py        | 15 +++--
 fmralign/hyperalignment/local_template.py     | 14 +---
 .../hyperalignment/piecewise_alignment.py     |  9 ++-
 fmralign/hyperalignment/regions.py            | 13 +---
 fmralign/hyperalignment/toy_experiment.py     | 65 +++++++++----------
 7 files changed, 68 insertions(+), 92 deletions(-)

diff --git a/examples/plot_int_alignment.py b/examples/plot_int_alignment.py
index eab04a9..9c013f6 100644
--- a/examples/plot_int_alignment.py
+++ b/examples/plot_int_alignment.py
@@ -112,30 +112,10 @@
 from fmralign.hyperalignment.regions import compute_parcels, compute_searchlights
 
 
-train_index = range(53)
 model = IndividualizedNeuralTuning(
-    n_jobs=10, alignment_method="parcelation", n_components=20
+    n_jobs=8, alignment_method="searchlight", n_components=30
 )
 
-if True:  # Use Parcellation
-    parcels = compute_parcels(
-        niimg=template_train[0], mask=masker, n_parcels=200, n_jobs=5
-    )
-    model.fit(np.array(masked_imgs)[:, train_index, :], parcels=parcels, verbose=False)
-else:
-    _, searchlights, dists = compute_searchlights(
-        niimg=template_train[0], mask_img=masker.mask_img, n_jobs=5
-    )
-    model.fit(
-        np.array(masked_imgs)[:, train_index, :],
-        searchlights=searchlights,
-        dists=dists,
-        verbose=False,
-    )
-
-train_stimulus = np.copy(model.shared_response)
-
-
 ###############################################################################
 # Predict new data for left-out subject
 # -------------------------------------
@@ -145,12 +125,30 @@
 # For each train subject and for the template, the AP contrasts are sorted from
 # 0, to 53, and then the PA contrasts from 53 to 106.
 #
+
+train_index = range(53)
 test_index = range(53, 106)
-if True:
+
+if False:
+    parcels = compute_parcels(
+        niimg=template_train[0], mask=masker, n_parcels=1000, n_jobs=5
+    )
+    model.fit(np.array(masked_imgs)[:, train_index, :], parcels=parcels, verbose=False)
+    train_stimulus = np.copy(model.shared_response)
     model.fit(np.array(masked_imgs)[:, test_index, :], parcels=parcels, verbose=False)
     test_stimulus = np.copy(model.shared_response)
 
 else:
+    _, searchlights, dists = compute_searchlights(
+        niimg=template_train[0], mask_img=masker.mask_img, n_jobs=5
+    )
+    model.fit(
+        np.array(masked_imgs)[:, train_index, :],
+        searchlights=searchlights,
+        dists=dists,
+        verbose=False,
+    )
+    train_stimulus = np.copy(model.shared_response)
     model.fit(
         np.array(masked_imgs)[:, test_index, :],
         searchlights=searchlights,
diff --git a/fmralign/alignment_methods.py b/fmralign/alignment_methods.py
index 87f9803..125278d 100644
--- a/fmralign/alignment_methods.py
+++ b/fmralign/alignment_methods.py
@@ -548,7 +548,7 @@ def _tuning_estimator(shared_response, target):
             array-like: The estimated tuning weights.
 
         """
-        return np.linalg.pinv(shared_response) @ target
+        return np.linalg.pinv(shared_response).dot(target)
 
     @staticmethod
     def _stimulus_estimator(full_signal, n_t, n_s, latent_dim=None):
diff --git a/fmralign/hyperalignment/correlation.py b/fmralign/hyperalignment/correlation.py
index c34cdfe..89c7011 100644
--- a/fmralign/hyperalignment/correlation.py
+++ b/fmralign/hyperalignment/correlation.py
@@ -75,7 +75,7 @@ def vector_sim(u, v):
     return sim
 
 
-def compute_pearson_corr(X, Y, linear_assignment: bool = True):
+def compute_pearson_corr(X, Y, linear_assignment: bool = False):
     """Compute Pearson correlation between X and Y.
     X and Y are two lists of matrices of the same shape.
     The returned matrix will be of shape 2N x 2N, where N is the number of matrices in X and Y.
@@ -88,7 +88,9 @@ def compute_pearson_corr(X, Y, linear_assignment: bool = True):
     corr_mat = np.zeros((n, n))
     for i in range(n):
         for j in range(n):
-            corr_i_j = pearson_corr_coeff(XY[i], XY[j])
+            corr_i_j = pearson_corr_coeff(
+                XY[i], XY[j], linear_assignment=linear_assignment
+            )
             corr_mat[i, j] = corr_i_j
 
     return corr_mat
@@ -151,7 +153,7 @@ def stimulus_correlation(X, Y, linear_assignment=True, absolute=True):
     """Compute pairwise Pearson correlation matrix between two stimulus matrices."""
     assert X.shape == Y.shape
     n = X.shape[0]
-    corr_mat = np.corrcoef(X, Y, dtype=np.float32)[:n, n:]
+    corr_mat = np.corrcoef(X, Y)[:n, n:]
 
     if absolute:
         corr_mat = np.abs(corr_mat)
@@ -161,7 +163,7 @@ def stimulus_correlation(X, Y, linear_assignment=True, absolute=True):
         corr_mat = corr_mat[row_ind, :]
         corr_mat = corr_mat[:, col_ind]
 
-    return corr_mat.astype(np.float16)
+    return corr_mat
 
 
 def matrix_MDS(X, Y, n_components=2, dissimilarity="euclidean"):
@@ -255,7 +257,6 @@ def multithread_compute_correlation(
             - corr_diff_sub_same_TR: Correlations between the same time points of different subjects.
     """
     from joblib import Parallel, delayed
-    from tqdm import tqdm
 
     def thread_compute_correlation(X, Y, i, j, absolute=False, linear_assignment=True):
         X_i, Y_i = X[i], Y[i]
@@ -284,9 +285,9 @@ def thread_compute_correlation(X, Y, i, j, absolute=False, linear_assignment=Tru
 
     assert X.shape == Y.shape
     n_s = X.shape[0]
-    corrdinates = list(combinations(range(n_s), 2)) + [(i, i) for i in range(n_s)]
+    coordinates = list(combinations(range(n_s), 2)) + [(i, i) for i in range(n_s)]
     results = Parallel(n_jobs=n_jobs)(
-        delayed(thread_compute_correlation)(X, Y, i, j) for (i, j) in tqdm(corrdinates)
+        delayed(thread_compute_correlation)(X, Y, i, j) for (i, j) in coordinates
     )
     results = list(zip(*results))
     corr_same_sub_diff_TR = results[2]
diff --git a/fmralign/hyperalignment/local_template.py b/fmralign/hyperalignment/local_template.py
index 5b1aef2..de20986 100644
--- a/fmralign/hyperalignment/local_template.py
+++ b/fmralign/hyperalignment/local_template.py
@@ -8,7 +8,6 @@
 from scipy.stats import zscore
 from sklearn.decomposition import PCA
 from sklearn.utils.extmath import randomized_svd
-from tqdm import tqdm
 
 from .linalg import safe_svd
 from .linalg import procrustes
@@ -148,11 +147,7 @@ def compute_procrustes_template(
         X = X[:, :, region]
     common_space = np.copy(X[0])
     aligned_X = [X[0]]
-    if debug:
-        iter_X = tqdm(X[1:])
-        iter_X.set_description("Computing procrustes alignment (level 1)...")
-    else:
-        iter_X = X[1:]
+    iter_X = X[1:]
     for x in iter_X:
         T = procrustes(x, common_space, reflection=reflection, scaling=scaling)
         aligned_x = x.dot(T)
@@ -164,12 +159,7 @@ def compute_procrustes_template(
             common_space = np.nan_to_num(zscore(common_space, axis=0))
 
     aligned_X2 = []
-
-    if debug:
-        iter2 = tqdm(range(level2_iter))
-        iter2.set_description("Computing procrustes alignment (level 2)...")
-    else:
-        iter2 = range(level2_iter)
+    iter2 = range(level2_iter)
 
     for level2 in iter2:
         common_space = np.zeros_like(X[0])
diff --git a/fmralign/hyperalignment/piecewise_alignment.py b/fmralign/hyperalignment/piecewise_alignment.py
index 0c8cb67..542694a 100644
--- a/fmralign/hyperalignment/piecewise_alignment.py
+++ b/fmralign/hyperalignment/piecewise_alignment.py
@@ -57,7 +57,7 @@ def __init__(
         self.radius = None
         self.weights = None
 
-    def compute_linear_transformation(self, data, template, i: int = 0, save=True):
+    def compute_linear_transformation(self, data, template):
         """Compute the linear transformation for a given subject provided the global template.
 
         Parameters
@@ -75,8 +75,8 @@ def compute_linear_transformation(self, data, template, i: int = 0, save=True):
         """
 
         x_hat = piece_ridge(
-            X=data,
-            Y=template,
+            X=template,
+            Y=data,
             regions=self.regions,
             weights=self.weights,
             verbose=self.verbose,
@@ -90,7 +90,6 @@ def fit_transform(
         dists=None,
         radius=None,
         weights=None,
-        id=None,
     ):
         """From given fmri data, compute the global template and the linear transformation.
         This provides denoised signal estimations using template alignment.
@@ -148,7 +147,7 @@ def fit_transform(
         )
 
         self.Xhat = Parallel(n_jobs=self.n_jobs)(
-            delayed(self.compute_linear_transformation)(X[i], sl_template, i)
+            delayed(self.compute_linear_transformation)(X[i], sl_template)
             for i in range(self.n_s)
         )
         return np.array(self.Xhat)
diff --git a/fmralign/hyperalignment/regions.py b/fmralign/hyperalignment/regions.py
index 03f1061..9016146 100644
--- a/fmralign/hyperalignment/regions.py
+++ b/fmralign/hyperalignment/regions.py
@@ -18,7 +18,6 @@
 from nilearn._utils.niimg_conversions import (
     safe_get_data,
 )
-from tqdm import tqdm
 from .linalg import procrustes
 from .linalg import ridge
 from .local_template import compute_template
@@ -377,8 +376,6 @@ def iter_hyperalignment(
 
     if weights is not None:
         zip_iter = zip(searchlights, weights)
-        if verbose:
-            zip_iter = tqdm(zip_iter, leave=False)
         for sl, w in zip_iter:
             x, y = X[:, sl], Y[:, sl]
             t = sl_func(x, y)
@@ -386,8 +383,6 @@ def iter_hyperalignment(
             del t
     else:
         searchlights_iter = searchlights
-        if verbose:
-            searchlights_iter = tqdm(searchlights_iter, leave=False)
         for sl in searchlights_iter:
             x, y = X[:, sl], Y[:, sl]
             t = sl_func(x, y)
@@ -495,12 +490,6 @@ def template(
         numpy.ndarray: The computed template of shape (n_features,).
 
     """
-
-    if verbose:
-        iterator = tqdm(regions, leave=False)
-        iterator.set_description("Computing local templates")
-    else:
-        iterator = regions
     with Parallel(n_jobs=n_jobs, batch_size=1, verbose=1) as parallel:
         local_templates = parallel(
             delayed(compute_template)(
@@ -510,7 +499,7 @@ def template(
                 max_npc=None,
                 common_topography=True,
             )
-            for region in iterator
+            for region in regions
         )
 
     template = np.zeros_like(X[0])
diff --git a/fmralign/hyperalignment/toy_experiment.py b/fmralign/hyperalignment/toy_experiment.py
index c8f331f..7ba81cc 100644
--- a/fmralign/hyperalignment/toy_experiment.py
+++ b/fmralign/hyperalignment/toy_experiment.py
@@ -23,12 +23,12 @@
 
 
 n_s = 10
-n_t = 100
-n_v = 200
-S_std = 1
+n_t = 200
+n_v = 500
+S_std = 5
 T_std = 1
 SNR = 100
-latent_dim = None  # if None, latent_dim = n_t
+latent_dim = 15  # if None, latent_dim = n_t
 decomposition_method = None  # if None, SVD is used
 
 
@@ -49,8 +49,7 @@
     T_std=T_std,
     latent_dim=latent_dim,
     SNR=SNR,
-    seed=0,
-    generative_method="fastsrm",
+    seed=42,
 )
 
 SEARCHLIGHT = False
@@ -60,25 +59,27 @@
         n_searchlights=12, n_v=n_v, radius=5
     )
 else:
-    searchlights = [np.arange(n_v)]
-    dists = [np.zeros((n_v,))]
+    parcels = [range(n_v)]
 
 #############################################################################
 # Test INT on the two parts of the data (ie different runs of the experiment)
-int1 = INT(n_components=latent_dim, decomp_method=decomposition_method)
-int2 = INT(n_components=latent_dim, decomp_method=decomposition_method)
+int1 = INT(
+    n_components=latent_dim,
+    decomp_method=decomposition_method,
+    alignment_method="parcelation",
+)
+int2 = INT(
+    n_components=latent_dim,
+    decomp_method=decomposition_method,
+    alignment_method="parcelation",
+)
 int_first_part = int1.fit(
-    data_run_1,
-    searchlights=searchlights,
-    dists=dists,
+    data_run_1, parcels=parcels, verbose=False
 )  # S is provided if we cheat and know the ground truth
-int_second_part = int2.fit(data_run_2, searchlights=searchlights, dists=dists)
+int_second_part = int2.fit(data_run_2, parcels=parcels, verbose=False)
 
-print("INT 1 denoised signal", int1.denoised_signal.shape)
 
-data_pred = int1.transform(data_run_2)
 # save individual components
-
 tuning_pred_run_1 = int1.tuning_data
 tuning_pred_run_1 = np.array(tuning_pred_run_1)
 tuning_pred_run_2 = int2.tuning_data
@@ -87,25 +88,27 @@
 stimulus_pred_run_1 = int1.shared_response
 stimulus_pred_run_2 = int2.shared_response
 
+data_pred = int1.transform(data_run_2)
+
 
 #############################################################################
 # Plot
+#############################################################################
+
 plt.rc("font", size=6)
-fig, ax = plt.subplots(3, 3, figsize=(20, 10))
+fig, ax = plt.subplots(2, 3, figsize=(10, 5))
 
 
 # Tunning matrices
 correlation_tuning = tuning_correlation(tuning_pred_run_1, tuning_pred_run_2)
 ax[0, 0].imshow(correlation_tuning)
-ax[0, 0].set_title("Run 1 vs Run 2")
+ax[0, 0].set_title("Correlation Tuning Run 1 vs Run 2")
 ax[0, 0].set_xlabel("Subjects, Run 1")
 ax[0, 0].set_ylabel("Subjects, Run 2")
 fig.colorbar(ax[0, 0].imshow(correlation_tuning), ax=ax[0, 0])
 
 random_colors = np.random.rand(n_s, 3)
 # MDS of predicted images
-print("X_run_2_pred.shape", data_pred.shape)
-print("X_run_2.shape", data_run_2.shape)
 corr_tunning = compute_pearson_corr(data_pred, data_run_2)
 data_pred_reduced, data_test_reduced = matrix_MDS(
     data_pred, data_run_2, n_components=2, dissimilarity=1 - corr_tunning
@@ -146,16 +149,12 @@
 )
 ax[0, 2].set_title("MDS of tunning matrices, dim=2")
 
-# Stimulus matrix
-
-print("stimulus_pred_run_1.shape", stimulus_pred_run_1.shape)
-print("stimulus_run_1.shape", stimulus_run_1.shape)
-
+# Stimulus matrix correlation
 correlation_stimulus_true_est_first_part = stimulus_correlation(
     stimulus_pred_run_1.T, stimulus_run_1.T
 )
 ax[1, 0].imshow(correlation_stimulus_true_est_first_part)
-ax[1, 0].set_title("Stimumus Run 1 Estimated vs ground truth")
+ax[1, 0].set_title("Stimumus Estimated vs ground truth (Run 1)")
 ax[1, 0].set_xlabel("Latent components, Run 1")
 ax[1, 0].set_ylabel("Latent components, ground truth")
 fig.colorbar(ax[1, 0].imshow(correlation_stimulus_true_est_first_part), ax=ax[1, 0])
@@ -164,25 +163,25 @@
     stimulus_pred_run_2.T, stimulus_run_2.T
 )
 ax[1, 1].imshow(correlation_stimulus_true_est_second_part)
-ax[1, 1].set_title("Stimulus Run 2 Estimated vs ground truth")
+ax[1, 1].set_title("Stimulus Estimated vs ground truth (Run 2)")
 ax[1, 1].set_xlabel("Latent components, Run 2")
 ax[1, 1].set_ylabel("Latent components, ground truth")
 fig.colorbar(ax[1, 1].imshow(correlation_stimulus_true_est_second_part), ax=ax[1, 1])
 
 
 # Reconstruction
-correlation_reconstruction_first_second = tuning_correlation(data_pred, data_run_2)
-ax[1, 2].imshow(correlation_reconstruction_first_second)
-ax[1, 2].set_title("Reconstruction")
+corr_reconstruction = tuning_correlation(data_pred, data_run_2)
+ax[1, 2].imshow(corr_reconstruction)
+ax[1, 2].set_title("Reconstruction correlation")
 ax[1, 2].set_xlabel("Subjects, Run 2")
 ax[1, 2].set_ylabel("Subjects, Run 1")
-fig.colorbar(ax[1, 2].imshow(correlation_reconstruction_first_second), ax=ax[1, 2])
+fig.colorbar(ax[1, 2].imshow(corr_reconstruction), ax=ax[1, 2])
 
 
 plt.rc("font", size=10)
 # Define small font for titles
 fig.suptitle(
-    f"Tunning matrices for the two parts of the data\n ns={n_s}, nt={n_t}, nv={n_v}, S_std={S_std}, T_std={T_std}, SNR={SNR}, latent space dim={latent_dim}"
+    f"Correlation Run 1/2\n ns={n_s}, nt={n_t}, nv={n_v}, S_std={S_std}, T_std={T_std}, SNR={SNR}, latent space dim={latent_dim}"
 )
 plt.tight_layout()
 

From 4a8ddaa1b9e3f8e979c8081302a74da806efee76 Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Wed, 10 Jan 2024 17:20:12 +0100
Subject: [PATCH 28/69] bug fixes + interesting results

---
 examples/plot_int_alignment.py                | 79 +++++++++++++------
 fmralign/hyperalignment/correlation.py        |  2 +-
 fmralign/hyperalignment/local_template.py     | 36 ++++++---
 .../hyperalignment/piecewise_alignment.py     |  2 +-
 fmralign/hyperalignment/regions.py            |  4 +-
 5 files changed, 82 insertions(+), 41 deletions(-)

diff --git a/examples/plot_int_alignment.py b/examples/plot_int_alignment.py
index 9c013f6..e661861 100644
--- a/examples/plot_int_alignment.py
+++ b/examples/plot_int_alignment.py
@@ -23,7 +23,10 @@
     :depth: 1
 
 """
+# %%
+import warnings
 
+warnings.filterwarnings("ignore")
 ###############################################################################
 # Retrieve the data
 # -----------------
@@ -39,7 +42,21 @@
 from fmralign.fetch_example_data import fetch_ibc_subjects_contrasts
 
 imgs, df, mask_img = fetch_ibc_subjects_contrasts(
-    ["sub-01", "sub-02", "sub-04", "sub-05", "sub-06", "sub-07"]
+    [
+        "sub-01",
+        "sub-02",
+        "sub-04",
+        "sub-05",
+        "sub-06",
+        "sub-07",
+        "sub-08",
+        "sub-09",
+        "sub-11",
+        "sub-12",
+        "sub-13",
+        "sub-14",
+        "sub-15",
+    ],
 )
 
 ###############################################################################
@@ -69,7 +86,7 @@
 from nilearn.image import concat_imgs
 
 template_train = []
-for i in range(5):
+for i in range(6):
     template_train.append(concat_imgs(imgs[i]))
 target_train = df[df.subject == "sub-07"][df.acquisition == "ap"].path.values
 
@@ -82,6 +99,7 @@
 target_train_data = masker.transform(target_train)
 target_test = df[df.subject == "sub-07"][df.acquisition == "pa"].path.values
 
+
 ###############################################################################
 # Compute a baseline (average of subjects)
 # ----------------------------------------
@@ -92,7 +110,7 @@
 import numpy as np
 
 masked_imgs = [masker.transform(img) for img in template_train]
-average_img = np.mean(masked_imgs, axis=0)
+average_img = np.mean(masked_imgs[:-1], axis=0)
 average_subject = masker.inverse_transform(average_img)
 
 ###############################################################################
@@ -111,11 +129,6 @@
 from fmralign.alignment_methods import IndividualizedNeuralTuning
 from fmralign.hyperalignment.regions import compute_parcels, compute_searchlights
 
-
-model = IndividualizedNeuralTuning(
-    n_jobs=8, alignment_method="searchlight", n_components=30
-)
-
 ###############################################################################
 # Predict new data for left-out subject
 # -------------------------------------
@@ -129,42 +142,52 @@
 train_index = range(53)
 test_index = range(53, 106)
 
-if False:
+train_data = np.array(masked_imgs)[:, train_index, :]
+test_data = np.array(masked_imgs)[:, train_index, :][:-1]
+
+if True:
     parcels = compute_parcels(
         niimg=template_train[0], mask=masker, n_parcels=1000, n_jobs=5
     )
-    model.fit(np.array(masked_imgs)[:, train_index, :], parcels=parcels, verbose=False)
+    model = IndividualizedNeuralTuning(
+        n_jobs=8, alignment_method="searchlight", n_components=None
+    )
+    model.fit(train_data, parcels=parcels, verbose=False)
     train_stimulus = np.copy(model.shared_response)
-    model.fit(np.array(masked_imgs)[:, test_index, :], parcels=parcels, verbose=False)
-    test_stimulus = np.copy(model.shared_response)
+    train_tuning = model.tuning_data[-1]
+    model_bis = IndividualizedNeuralTuning(
+        n_jobs=8, alignment_method="searchlight", n_components=None
+    )
+    model_bis.fit(test_data, parcels=parcels, verbose=False)
+    test_stimulus = np.copy(model_bis.shared_response)
 
 else:
     _, searchlights, dists = compute_searchlights(
         niimg=template_train[0], mask_img=masker.mask_img, n_jobs=5
     )
-    model.fit(
-        np.array(masked_imgs)[:, train_index, :],
+    model_bis.fit(
+        train_data,
         searchlights=searchlights,
         dists=dists,
         verbose=False,
     )
-    train_stimulus = np.copy(model.shared_response)
-    model.fit(
-        np.array(masked_imgs)[:, test_index, :],
+    train_stimulus = np.copy(model_bis.shared_response)
+    model_bis.fit(
+        test_data,
         searchlights=searchlights,
         dists=dists,
         verbose=False,
     )
-    test_stimulus = np.copy(model.shared_response)
+    test_stimulus = np.copy(model_bis.shared_response)
 
 
+# %%
 # We input the mapping image target_train in a list, we could have input more
 # than one subject for which we'd want to predict : [train_1, train_2 ...]
 
-tuning_target = np.linalg.pinv(train_stimulus[train_index, :]) @ target_train_data
+prediction_from_template = test_stimulus @ train_tuning
+prediction_from_template = masker.inverse_transform(prediction_from_template)
 
-prediction_from_template = test_stimulus @ tuning_target
-prediction_from_template = [masker.inverse_transform(prediction_from_template)]
 
 # As a baseline prediction, let's just take the average of activations across subjects.
 
@@ -178,21 +201,26 @@
 # measure the correlation between its profile of activation without and with
 # alignment, to see if alignment was able to predict a signal more alike the ground truth.
 #
-
+# %%
 from fmralign.metrics import score_voxelwise
 
 # Now we use this scoring function to compare the correlation of predictions
 # made from group average and from template with the real PA contrasts of sub-07
 
+
 average_score = masker.inverse_transform(
-    score_voxelwise(target_test, prediction_from_average, masker, loss="corr")
+    C_avg := score_voxelwise(target_test, prediction_from_average, masker, loss="corr")
 )
 
-# I choose abs value in reference to the work we did with the INT
 template_score = masker.inverse_transform(
-    score_voxelwise(target_test, prediction_from_template[0], masker, loss="corr")
+    C_temp := score_voxelwise(
+        target_test, prediction_from_template, masker, loss="corr"
+    )
 )
 
+print("============Mean correlation============")
+print(f"Baseline : {np.mean(C_avg)}")
+print(f"Template : {np.mean(C_temp)}")
 
 ###############################################################################
 # Plotting the measures
@@ -200,6 +228,7 @@
 # Finally we plot both scores
 #
 
+# %%
 from nilearn import plotting
 
 baseline_display = plotting.plot_stat_map(
diff --git a/fmralign/hyperalignment/correlation.py b/fmralign/hyperalignment/correlation.py
index 89c7011..ec74cca 100644
--- a/fmralign/hyperalignment/correlation.py
+++ b/fmralign/hyperalignment/correlation.py
@@ -237,7 +237,7 @@ def thread_compute_correlation(X, Y, i, j):
 
 
 def multithread_compute_correlation(
-    X, Y, absolute=False, linear_assignment=True, n_jobs=40
+    X, Y, absolute=False, linear_assignment=True, n_jobs=1
 ):
     """
     Compute correlations between pairs of samples in X and Y using multiple threads.
diff --git a/fmralign/hyperalignment/local_template.py b/fmralign/hyperalignment/local_template.py
index de20986..1d8c11b 100644
--- a/fmralign/hyperalignment/local_template.py
+++ b/fmralign/hyperalignment/local_template.py
@@ -81,7 +81,7 @@ def PCA_decomposition(
         raise NotImplementedError
 
 
-def compute_PCA_template(X, sl=None, max_npc=None, flavor="sklearn", demean=False):
+def compute_PCA_template(X, sl=None, n_components=None, flavor="sklearn", demean=False):
     """
     Compute the PCA template from the input data.
 
@@ -92,7 +92,7 @@ def compute_PCA_template(X, sl=None, max_npc=None, flavor="sklearn", demean=Fals
         The input data array of shape (n_samples, n_features, n_timepoints).
     - sl (slice, optional):
         The slice of timepoints to consider. Defaults to None.
-    - max_npc (int, optional):
+    - n_components (int, optional):
         The maximum number of principal components to keep. Defaults to None.
     - flavor (str, optional):
         The flavor of PCA algorithm to use. Defaults to "sklearn".
@@ -109,13 +109,26 @@ def compute_PCA_template(X, sl=None, max_npc=None, flavor="sklearn", demean=Fals
         X = X[:, :, sl]
     else:
         X = X
-    max_npc = min(X.shape[1], X.shape[2])
+    n_components = min(X.shape[1], X.shape[2])
     XX, cc = PCA_decomposition(
-        X, n_components=max_npc, flavor=flavor, adjust_ns=True, demean=demean
+        X, n_components=n_components, flavor=flavor, adjust_ns=True, demean=demean
     )
     return XX.astype(np.float32)
 
 
+def compute_PCA_var1_template(
+    dss, sl=None, n_components=None, flavor="sklearn", demean=True
+):
+    if sl is not None:
+        dss = dss[:, :, sl]
+    XX, cc = PCA_decomposition(
+        dss, n_components=n_components, flavor=flavor, adjust_ns=False, demean=demean
+    )
+    w = np.sqrt(np.sum(cc**2, axis=2)).mean(axis=1)
+    XX *= w[np.newaxis]
+    return XX
+
+
 def compute_procrustes_template(
     X,
     region=None,
@@ -124,7 +137,6 @@ def compute_procrustes_template(
     zscore_common=True,
     level2_iter=1,
     X2=None,
-    debug=False,
 ):
     """
     Compute the Procrustes template for a given set of data.
@@ -195,8 +207,8 @@ def compute_procrustes_template(
 def compute_template(
     X,
     region,
-    kind="searchlight_pca",
-    max_npc=None,
+    kind="pca",
+    n_components=150,
     common_topography=True,
     demean=True,
 ):
@@ -211,7 +223,7 @@ def compute_template(
        - sl (ndarray or None): The searchlight indices for searchlight-based template computation.
        - region (int or None, optional): The index of the region to consider. If None, all regions are considered (or searchlights). Defaults to None.
        - kind (str): The type of template computation algorithm to use. Can be "pca", "pcav1", "pcav2", or "cls".
-       - max_npc (int or None): The maximum number of principal components to use for PCA-based template computation.
+       - n_components (int or None): The maximum number of principal components to use for PCA-based template computation.
        - common_topography (bool): Whether to enforce common topography across datasets.
        - demean (bool): Whether to demean the datasets before template computation.
 
@@ -220,9 +232,9 @@ def compute_template(
         tmpl : The computed template on all parcels (or searchlights).
     """
     mapping = {
-        "searchlight_pca": compute_PCA_template,
-        "parcels_pca": compute_PCA_template,
-        "cls": compute_procrustes_template,
+        "pca": compute_PCA_template,
+        "pcav1": compute_PCA_var1_template,
+        "procrustes": compute_procrustes_template,
     }
 
     if kind == "procrustes":
@@ -234,7 +246,7 @@ def compute_template(
             zscore_common=True,
         )
     elif kind in mapping:
-        tmpl = mapping[kind](X, sl=region, max_npc=max_npc, demean=demean)
+        tmpl = mapping[kind](X, sl=region, n_components=n_components, demean=demean)
     else:
         raise ValueError("Unknown template kind")
 
diff --git a/fmralign/hyperalignment/piecewise_alignment.py b/fmralign/hyperalignment/piecewise_alignment.py
index 542694a..3ac8afb 100644
--- a/fmralign/hyperalignment/piecewise_alignment.py
+++ b/fmralign/hyperalignment/piecewise_alignment.py
@@ -26,7 +26,7 @@ class PiecewiseAlignment(BaseEstimator, TransformerMixin):
     def __init__(
         self,
         alignment_method="searchlight_ridge",
-        template_kind="searchlight_pca",
+        template_kind="pcav1",
         verbose=True,
         n_jobs=1,
     ):
diff --git a/fmralign/hyperalignment/regions.py b/fmralign/hyperalignment/regions.py
index 9016146..9bddfc7 100644
--- a/fmralign/hyperalignment/regions.py
+++ b/fmralign/hyperalignment/regions.py
@@ -472,7 +472,7 @@ def template(
     X,
     regions,
     n_jobs=1,
-    template_kind="searchlight_pca",
+    template_kind="pca",
     verbose=False,
     weights=None,
 ):
@@ -496,7 +496,7 @@ def template(
                 X,
                 region=region,
                 kind=template_kind,
-                max_npc=None,
+                n_components=None,
                 common_topography=True,
             )
             for region in regions

From dbf694ecd8e5cb5f0ba876491cd7db9d4443a340 Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Fri, 12 Jan 2024 13:34:02 +0100
Subject: [PATCH 29/69] rm ref to ha in template alignment

---
 examples/plot_int_alignment.py | 20 ++++++++++----------
 fmralign/template_alignment.py |  4 ----
 2 files changed, 10 insertions(+), 14 deletions(-)

diff --git a/examples/plot_int_alignment.py b/examples/plot_int_alignment.py
index e661861..9851e6c 100644
--- a/examples/plot_int_alignment.py
+++ b/examples/plot_int_alignment.py
@@ -49,13 +49,6 @@
         "sub-05",
         "sub-06",
         "sub-07",
-        "sub-08",
-        "sub-09",
-        "sub-11",
-        "sub-12",
-        "sub-13",
-        "sub-14",
-        "sub-15",
     ],
 )
 
@@ -145,7 +138,7 @@
 train_data = np.array(masked_imgs)[:, train_index, :]
 test_data = np.array(masked_imgs)[:, train_index, :][:-1]
 
-if True:
+if False:
     parcels = compute_parcels(
         niimg=template_train[0], mask=masker, n_parcels=1000, n_jobs=5
     )
@@ -165,13 +158,20 @@
     _, searchlights, dists = compute_searchlights(
         niimg=template_train[0], mask_img=masker.mask_img, n_jobs=5
     )
-    model_bis.fit(
+    model = IndividualizedNeuralTuning(
+        n_jobs=8, alignment_method="searchlight", n_components=None
+    )
+    model.fit(
         train_data,
         searchlights=searchlights,
         dists=dists,
         verbose=False,
     )
-    train_stimulus = np.copy(model_bis.shared_response)
+    train_stimulus = np.copy(model.shared_response)
+    train_tuning = model.tuning_data[-1]
+    model_bis = IndividualizedNeuralTuning(
+        n_jobs=8, alignment_method="searchlight", n_components=None
+    )
     model_bis.fit(
         test_data,
         searchlights=searchlights,
diff --git a/fmralign/template_alignment.py b/fmralign/template_alignment.py
index 90bc57c..65cdb37 100644
--- a/fmralign/template_alignment.py
+++ b/fmralign/template_alignment.py
@@ -376,8 +376,6 @@ def fit(self, imgs):
 
         """
 
-        # Alignment method is hyperalignment (too different from other methods)
-
         # Check if the input is a list, if list of lists, concatenate each subjects
         # data into one unique image.
         if not isinstance(imgs, (list, np.ndarray)) or len(imgs) < 2:
@@ -443,8 +441,6 @@ def transform(self, imgs, train_index, test_index):
             Each Niimg has the same length as the list test_index
 
         """
-        if self.alignment_method == "hyperalignment":
-            return self.model.transform(imgs, verbose=self.verbose)
 
         if not isinstance(imgs, (list, np.ndarray)):
             raise ValueError(

From e78747aa84539d6d2f8b548c1d7a816f71c89773 Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Mon, 15 Jan 2024 12:24:40 +0100
Subject: [PATCH 30/69] fix searchlight toy

---
 fmralign/hyperalignment/toy_experiment.py | 48 ++++++++++++++++-------
 1 file changed, 33 insertions(+), 15 deletions(-)

diff --git a/fmralign/hyperalignment/toy_experiment.py b/fmralign/hyperalignment/toy_experiment.py
index 7ba81cc..81899d6 100644
--- a/fmralign/hyperalignment/toy_experiment.py
+++ b/fmralign/hyperalignment/toy_experiment.py
@@ -56,27 +56,45 @@
 
 if SEARCHLIGHT:
     searchlights, dists = generate_dummy_searchlights(
-        n_searchlights=12, n_v=n_v, radius=5
+        n_searchlights=n_v, n_v=n_v, radius=5
     )
+    int1 = INT(
+        n_components=latent_dim,
+        decomp_method=decomposition_method,
+        alignment_method="searchlight",
+    )
+    int2 = INT(
+        n_components=latent_dim,
+        decomp_method=decomposition_method,
+        alignment_method="searchlight",
+    )
+    int_first_part = int1.fit(
+        data_run_1, searchlights=searchlights, dists=dists, radius=5, verbose=False
+    )  # S is provided if we cheat and know the ground truth
+    int_second_part = int2.fit(
+        data_run_2, searchlights=searchlights, dists=dists, radius=5, verbose=False
+    )
+
+
 else:
     parcels = [range(n_v)]
+    int1 = INT(
+        n_components=latent_dim,
+        decomp_method=decomposition_method,
+        alignment_method="parcelation",
+    )
+    int2 = INT(
+        n_components=latent_dim,
+        decomp_method=decomposition_method,
+        alignment_method="parcelation",
+    )
+    int_first_part = int1.fit(
+        data_run_1, parcels=parcels, verbose=False
+    )  # S is provided if we cheat and know the ground truth
+    int_second_part = int2.fit(data_run_2, parcels=parcels, verbose=False)
 
 #############################################################################
 # Test INT on the two parts of the data (ie different runs of the experiment)
-int1 = INT(
-    n_components=latent_dim,
-    decomp_method=decomposition_method,
-    alignment_method="parcelation",
-)
-int2 = INT(
-    n_components=latent_dim,
-    decomp_method=decomposition_method,
-    alignment_method="parcelation",
-)
-int_first_part = int1.fit(
-    data_run_1, parcels=parcels, verbose=False
-)  # S is provided if we cheat and know the ground truth
-int_second_part = int2.fit(data_run_2, parcels=parcels, verbose=False)
 
 
 # save individual components

From 027154b92d8f22227a3a0199892f212212064572 Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Fri, 19 Jan 2024 13:30:27 +0100
Subject: [PATCH 31/69] better doc and parameters

---
 examples/plot_hyperalignment.py               | 257 ++++++++++++++++++
 fmralign/alignment_methods.py                 |  25 +-
 fmralign/hyperalignment/correlation.py        |   7 +-
 .../hyperalignment/piecewise_alignment.py     |   9 +-
 4 files changed, 280 insertions(+), 18 deletions(-)
 create mode 100644 examples/plot_hyperalignment.py

diff --git a/examples/plot_hyperalignment.py b/examples/plot_hyperalignment.py
new file mode 100644
index 0000000..e3c711b
--- /dev/null
+++ b/examples/plot_hyperalignment.py
@@ -0,0 +1,257 @@
+# -*- coding: utf-8 -*-
+
+"""
+Hyperalignment-base prediction using the IndividualNeuralTuning Model.
+See article : https://doi.org/10.1162/imag_a_00032
+
+==========================
+
+In this tutorial, we show how to better predict new contrasts for a target
+subject using many source subjects corresponding contrasts. For this purpose,
+we create a template to which we align the target subject, using shared information.
+We then predict new images for the target and compare them to a baseline.
+
+We mostly rely on Python common packages and on nilearn to handle
+functional data in a clean fashion.
+
+
+To run this example, you must launch IPython via ``ipython
+--matplotlib`` in a terminal, or use ``jupyter-notebook``.
+
+.. contents:: **Contents**
+    :local:
+    :depth: 1
+
+"""
+# %%
+import warnings
+
+warnings.filterwarnings("ignore")
+###############################################################################
+# Retrieve the data
+# -----------------
+# In this example we use the IBC dataset, which includes a large number of
+# different contrasts maps for 12 subjects.
+# We download the images for subjects sub-01, sub-02, sub-04, sub-05, sub-06
+# and sub-07 (or retrieve them if they were already downloaded).
+# imgs is the list of paths to available statistical images for each subjects.
+# df is a dataframe with metadata about each of them.
+# mask is a binary image used to extract grey matter regions.
+#
+
+from fmralign.fetch_example_data import fetch_ibc_subjects_contrasts
+
+imgs, df, mask_img = fetch_ibc_subjects_contrasts(
+    [
+        "sub-01",
+        "sub-02",
+        "sub-04",
+        "sub-05",
+        "sub-06",
+        "sub-07",
+    ],
+)
+
+SEARCHLIGHT = False
+
+###############################################################################
+# Definine a masker
+# -----------------
+# We define a nilearn masker that will be used to handle relevant data.
+#   For more information, visit :
+#   'http://nilearn.github.io/manipulating_images/masker_objects.html'
+#
+
+from nilearn.maskers import NiftiMasker
+
+masker = NiftiMasker(mask_img=mask_img).fit()
+
+###############################################################################
+# Prepare the data
+# ----------------
+# For each subject, we will use two series of contrasts acquired during
+# two independent sessions with a different phase encoding:
+# Antero-posterior(AP) or Postero-anterior(PA).
+#
+
+
+# To infer a template for subjects sub-01 to sub-06 for both AP and PA data,
+# we make a list of 4D niimgs from our list of list of files containing 3D images
+
+from nilearn.image import concat_imgs
+
+template_train = []
+for i in range(6):
+    template_train.append(concat_imgs(imgs[i]))
+target_train = df[df.subject == "sub-07"][df.acquisition == "ap"].path.values
+
+# For subject sub-07, we split it in two folds:
+#   - target train: sub-07 AP contrasts, used to learn alignment to template
+#   - target test: sub-07 PA contrasts, used as a ground truth to score predictions
+# We make a single 4D Niimg from our list of 3D filenames
+
+target_train = concat_imgs(target_train)
+target_train_data = masker.transform(target_train)
+target_test = df[df.subject == "sub-07"][df.acquisition == "pa"].path.values
+
+
+###############################################################################
+# Compute a baseline (average of subjects)
+# ----------------------------------------
+# We create an image with as many contrasts as any subject representing for
+# each contrast the average of all train subjects maps.
+#
+
+import numpy as np
+
+masked_imgs = [masker.transform(img) for img in template_train]
+average_img = np.mean(masked_imgs[:-1], axis=0)
+average_subject = masker.inverse_transform(average_img)
+
+###############################################################################
+# Create a template from the training subjects.
+# ---------------------------------------------
+# We define an estimator using the class TemplateAlignment:
+#   * We align the whole brain through 'multiple' local alignments.
+#   * These alignments are calculated on a parcellation of the brain in 150 pieces,
+#     this parcellation creates group of functionnally similar voxels.
+#   * The template is created iteratively, aligning all subjects data into a
+#     common space, from which the template is inferred and aligning again to this
+#     new template space.
+#
+
+from fmralign.hyperalignment.regions import compute_parcels, compute_searchlights
+
+###############################################################################
+# Predict new data for left-out subject
+# -------------------------------------
+# We use target_train data to fit the transform, indicating it corresponds to
+# the contrasts indexed by train_index and predict from this learnt alignment
+# contrasts corresponding to template test_index numbers.
+# For each train subject and for the template, the AP contrasts are sorted from
+# 0, to 53, and then the PA contrasts from 53 to 106.
+#
+
+train_index = range(53)
+test_index = range(53, 106)
+
+train_data = np.array(masked_imgs)[:, train_index, :]
+test_data = np.array(masked_imgs)[:, test_index, :][:-1]
+
+if not SEARCHLIGHT:
+    parcels = compute_parcels(
+        niimg=template_train[0], mask=masker, n_parcels=1000, n_jobs=5
+    )
+
+else:
+    _, searchlights, dists = compute_searchlights(
+        niimg=template_train[0], mask_img=masker.mask_img, radius=5, n_jobs=5
+    )
+
+# %%
+
+from fmralign.hyperalignment.regions import template, piece_ridge, searchlight_weights
+
+if not SEARCHLIGHT:
+    sl_template = template(
+        X=train_data, regions=parcels, n_jobs=5, template_kind="procrustes"
+    )
+    train_tuning = piece_ridge(X=sl_template, Y=train_data[-1], regions=parcels)
+    test_template = template(X=test_data, regions=parcels, n_jobs=5)
+
+    train_tuning = piece_ridge(
+        X=sl_template, Y=train_data[-1], regions=parcels, return_betas=True
+    )
+else:
+    weights = searchlight_weights(searchlights=searchlights, dists=dists, radius=20)
+    sl_template = template(
+        X=train_data, regions=searchlights, n_jobs=5, weights=weights
+    )
+    train_tuning = piece_ridge(
+        X=sl_template,
+        Y=train_data[-1],
+        regions=searchlights,
+        weights=weights,
+        alpha=10,
+        return_betas=True,
+    )
+    test_template = template(
+        X=test_data, regions=searchlights, n_jobs=5, weights=weights
+    )
+
+# %%
+
+target_pred = test_template @ train_tuning
+
+target_pred = masker.inverse_transform(target_pred)
+
+# %%
+from fmralign.metrics import score_voxelwise
+
+# Now we use this scoring function to compare the correlation of predictions
+# made from group average and from template with the real PA contrasts of sub-07
+
+
+template_score = masker.inverse_transform(
+    C_temp := score_voxelwise(target_test, target_pred, masker=masker, loss="corr")
+)
+
+print("============Global correlation============")
+print(f"Template avg : {np.mean(C_temp)}")
+print(f"Template std : {np.std(C_temp)}")
+
+###############################################################################
+# Plotting the measures
+# ---------------------
+# Finally we plot both scores
+#
+
+# %%
+from nilearn import plotting
+
+display = plotting.plot_stat_map(
+    template_score, display_mode="z", cut_coords=[-15, -5], vmax=1
+)
+# display.title("Hyperalignment-based prediction correlation wt ground truth")
+
+###############################################################################
+# We observe that creating a template and aligning a new subject to it yields
+# a prediction that is better correlated with the ground truth than just using
+# the average activations of subjects.
+#
+
+
+# %%
+from nilearn.image import load_img
+from nilearn.maskers import NiftiMasker
+
+moviewatching_mask_img = load_img("/Users/df/nilearn_data/movie_watching.nii.gz")
+moviewatching_mask_img = masker.transform(moviewatching_mask_img)
+moviewatching_mask_img -= np.mean(moviewatching_mask_img, axis=0)
+moviewatching_mask_img /= np.std(moviewatching_mask_img, axis=0)
+moviewatching_mask_img = moviewatching_mask_img > 0
+moviewatching_mask_img = masker.inverse_transform(moviewatching_mask_img)
+moviewatching_masker = NiftiMasker(mask_img=moviewatching_mask_img).fit()
+
+
+target_pred = moviewatching_masker.transform(target_pred)
+target_test = moviewatching_masker.transform(target_test)
+
+target_pred = target_pred - np.mean(target_pred, axis=0)
+
+C_temp_localized = score_voxelwise(
+    target_test, target_pred, masker=moviewatching_masker, loss="corr"
+)
+
+template_localized_score = moviewatching_masker.inverse_transform(C_temp)
+
+print("============Global correlation============")
+print(f"Template avg : {np.mean(C_temp)}")
+print(f"Template std : {np.std(C_temp)}")
+
+localized_display = plotting.plot_stat_map(
+    template_score, display_mode="z", cut_coords=[-15, -5], vmax=1
+)
+
+
+plotting.show()
diff --git a/fmralign/alignment_methods.py b/fmralign/alignment_methods.py
index 125278d..8631a36 100644
--- a/fmralign/alignment_methods.py
+++ b/fmralign/alignment_methods.py
@@ -518,7 +518,6 @@ def __init__(
             self.distances = None
             self.radius = None
 
-        self.path = None
         self.tuning_data = []
         self.denoised_signal = []
         self.decomp_method = decomp_method
@@ -590,7 +589,7 @@ def _reconstruct_signal(shared_response, individual_tuning):
 
     def fit(
         self,
-        X_train,
+        X,
         searchlights=None,
         parcels=None,
         dists=None,
@@ -604,7 +603,7 @@ def fit(
         Parameters:
         --------
 
-        - X_train (array-like):
+        - X (array-like):
             The training data of shape (n_subjects, n_samples, n_voxels).
         - searchlights (array-like):
             The searchlight indices for each subject, of shape (n_s, n_searchlights).
@@ -628,15 +627,17 @@ def fit(
             The fitted model.
         """
 
-        X_train_ = np.array(X_train, copy=True, dtype=np.float32)
+        X_ = np.array(X, copy=True, dtype=np.float32)
 
-        self.n_s, self.n_t, self.n_v = X_train_.shape
+        self.n_s, self.n_t, self.n_v = X_.shape
 
         self.tuning_data = np.empty(self.n_s, dtype=np.float32)
         self.denoised_signal = np.empty(self.n_s, dtype=np.float32)
 
         if searchlights is None:
             self.regions = parcels
+            self.distances = None
+            self.radius = None
         else:
             self.regions = searchlights
             self.distances = dists
@@ -646,13 +647,11 @@ def fit(
             alignment_method=self.alignment_method, n_jobs=self.n_jobs, verbose=verbose
         )
         self.denoised_signal = denoiser.fit_transform(
-            X_train_,
+            X_,
             regions=self.regions,
-            dists=dists,
-            radius=radius,
+            dists=self.distances,
+            radius=self.radius,
         )
-        # Clear memory of the SearchlightAlignment object
-        denoiser = None
 
         # Stimulus matrix computation
         if self.decomp_method is None:
@@ -671,12 +670,12 @@ def fit(
 
         return self
 
-    def transform(self, X_test_, verbose=False):
+    def transform(self, X, verbose=False):
         """
         Transforms the input test data using the hyperalignment model.
 
         Args:
-            X_test_ (list of arrays):
+            X (list of arrays):
                 The input test data.
             verbose (bool, optional):
                 Whether to print verbose output. Defaults to False.
@@ -687,7 +686,7 @@ def transform(self, X_test_, verbose=False):
             numpy.ndarray: The transformed test data.
         """
 
-        full_signal = np.concatenate(X_test_, axis=1, dtype=np.float32)
+        full_signal = np.concatenate(X, axis=1, dtype=np.float32)
 
         if verbose:
             print("Predict : Computing stimulus matrix...")
diff --git a/fmralign/hyperalignment/correlation.py b/fmralign/hyperalignment/correlation.py
index ec74cca..e03ad12 100644
--- a/fmralign/hyperalignment/correlation.py
+++ b/fmralign/hyperalignment/correlation.py
@@ -258,9 +258,12 @@ def multithread_compute_correlation(
     """
     from joblib import Parallel, delayed
 
-    def thread_compute_correlation(X, Y, i, j, absolute=False, linear_assignment=True):
+    def thread_compute_correlation(X, Y, i, j):
         X_i, Y_i = X[i], Y[i]
-        corr = stimulus_correlation(X_i, Y_i, absolute=False)
+        corr = np.corrcoef(X_i, Y_i)[X.shape[1] :, : X.shape[1]]
+        row_ind, col_ind = linear_sum_assignment(corr, maximize=True)
+        corr = corr[row_ind, :]
+        corr = corr[:, col_ind]
         same_TR_corr = np.diag(corr)
         # Get all the values except the diagonal in a list
         diff_TR_corr = corr[np.where(~np.eye(corr.shape[0], dtype=bool))]
diff --git a/fmralign/hyperalignment/piecewise_alignment.py b/fmralign/hyperalignment/piecewise_alignment.py
index 3ac8afb..ed3153c 100644
--- a/fmralign/hyperalignment/piecewise_alignment.py
+++ b/fmralign/hyperalignment/piecewise_alignment.py
@@ -26,7 +26,7 @@ class PiecewiseAlignment(BaseEstimator, TransformerMixin):
     def __init__(
         self,
         alignment_method="searchlight_ridge",
-        template_kind="pcav1",
+        template_kind="pca",
         verbose=True,
         n_jobs=1,
     ):
@@ -132,11 +132,14 @@ def fit_transform(
             print(f"[{self.FUNC}]Computing global template M ...")
 
         if dists is None or radius is None:
-            self.weights = None
-        else:
+            self.weights = weights
+        elif weights is None:
             self.weights = searchlight_weights(
                 searchlights=regions, dists=dists, radius=radius
             )
+        else:
+            self.weights = weights
+
         sl_template = template(
             X,
             regions=regions,

From de8cf24436047a9494d180579474c13e2ecb4358 Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Fri, 19 Jan 2024 13:33:13 +0100
Subject: [PATCH 32/69] fix nilearn depreciation

---
 fmralign/_utils.py                 | 2 +-
 fmralign/alignment_methods.py      | 6 ++++--
 fmralign/hyperalignment/regions.py | 2 +-
 3 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/fmralign/_utils.py b/fmralign/_utils.py
index f31641f..9201c31 100644
--- a/fmralign/_utils.py
+++ b/fmralign/_utils.py
@@ -146,7 +146,7 @@ def _make_parcellation(
             )
             err.args += (errmsg,)
             raise err
-        labels = apply_mask_fmri(parcellation.labels_img_, masker.mask_img_).astype(int)
+        labels = apply_mask_fmri(parcellation.labels_img_, masker.mask_img).astype(int)
 
     if verbose > 0:
         unique_labels, counts = np.unique(labels, return_counts=True)
diff --git a/fmralign/alignment_methods.py b/fmralign/alignment_methods.py
index 125278d..fa93c05 100644
--- a/fmralign/alignment_methods.py
+++ b/fmralign/alignment_methods.py
@@ -548,7 +548,9 @@ def _tuning_estimator(shared_response, target):
             array-like: The estimated tuning weights.
 
         """
-        return np.linalg.pinv(shared_response).dot(target)
+        if shared_response.shape[1] < shared_response.shape[0]:
+            return (np.linalg.pinv(shared_response)).dot(target)
+        return np.linalg.inv(shared_response).dot(target)
 
     @staticmethod
     def _stimulus_estimator(full_signal, n_t, n_s, latent_dim=None):
@@ -568,7 +570,7 @@ def _stimulus_estimator(full_signal, n_t, n_s, latent_dim=None):
             U = svd_pca(full_signal)
             U = U[:, :latent_dim]
         else:
-            U, _, _ = safe_svd(full_signal)
+            U, _, _ = safe_svd(full_signal, remove_mean=False)
 
         stimulus = np.sqrt(n_s) * U
         stimulus = stimulus.astype(np.float32)
diff --git a/fmralign/hyperalignment/regions.py b/fmralign/hyperalignment/regions.py
index 9bddfc7..609cc8e 100644
--- a/fmralign/hyperalignment/regions.py
+++ b/fmralign/hyperalignment/regions.py
@@ -496,7 +496,7 @@ def template(
                 X,
                 region=region,
                 kind=template_kind,
-                n_components=None,
+                n_components=150,
                 common_topography=True,
             )
             for region in regions

From 0e71995d679ee8a012bf1ea9447653620708f02f Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Fri, 19 Jan 2024 13:35:07 +0100
Subject: [PATCH 33/69] fix experiments + add return betas to ridge

---
 examples/plot_int_alignment.py     | 16 +++++----
 fmralign/hyperalignment/regions.py | 56 +++++++++++++++++++++---------
 2 files changed, 49 insertions(+), 23 deletions(-)

diff --git a/examples/plot_int_alignment.py b/examples/plot_int_alignment.py
index 9851e6c..5bcf3cd 100644
--- a/examples/plot_int_alignment.py
+++ b/examples/plot_int_alignment.py
@@ -136,20 +136,20 @@
 test_index = range(53, 106)
 
 train_data = np.array(masked_imgs)[:, train_index, :]
-test_data = np.array(masked_imgs)[:, train_index, :][:-1]
+test_data = np.array(masked_imgs)[:, test_index, :][:-1]
 
-if False:
+if True:
     parcels = compute_parcels(
-        niimg=template_train[0], mask=masker, n_parcels=1000, n_jobs=5
+        niimg=template_train[0], mask=masker, n_parcels=100, n_jobs=5
     )
     model = IndividualizedNeuralTuning(
-        n_jobs=8, alignment_method="searchlight", n_components=None
+        n_jobs=8, alignment_method="parcellation", n_components=None
     )
     model.fit(train_data, parcels=parcels, verbose=False)
     train_stimulus = np.copy(model.shared_response)
-    train_tuning = model.tuning_data[-1]
+    train_tuning = np.linalg.pinv(train_stimulus) @ model.denoised_signal[-1]
     model_bis = IndividualizedNeuralTuning(
-        n_jobs=8, alignment_method="searchlight", n_components=None
+        n_jobs=8, alignment_method="parcellation", n_components=None
     )
     model_bis.fit(test_data, parcels=parcels, verbose=False)
     test_stimulus = np.copy(model_bis.shared_response)
@@ -185,7 +185,7 @@
 # We input the mapping image target_train in a list, we could have input more
 # than one subject for which we'd want to predict : [train_1, train_2 ...]
 
-prediction_from_template = test_stimulus @ train_tuning
+prediction_from_template = -test_stimulus @ train_tuning
 prediction_from_template = masker.inverse_transform(prediction_from_template)
 
 
@@ -247,3 +247,5 @@
 #
 
 plotting.show()
+
+# %%
diff --git a/fmralign/hyperalignment/regions.py b/fmralign/hyperalignment/regions.py
index 9bddfc7..f713edd 100644
--- a/fmralign/hyperalignment/regions.py
+++ b/fmralign/hyperalignment/regions.py
@@ -355,41 +355,63 @@ def iter_hyperalignment(
     searchlights,
     sl_func,
     weights=None,
+    return_betas=False,
     verbose=False,
 ):
     """
-    Perform searchlight hyperalignment on the given data.
+    Tool function to iterate hyperalignment over pieces of data.
 
-    Args:
-        X (ndarray): The source data matrix of shape (n_samples, n_features).
-        Y (ndarray): The target data matrix of shape (n_samples, n_features).
-        searchlights (list): List of searchlight indices.
-        dists (ndarray): distances of vertices to the center of their searchlight, of shape (n_searchlights, n_vertices_sl)
-        radius (float): Radius of the searchlight.
-        sl_func (function): Function to compute the searchlight transformation.
-        weighted (bool, optional): Whether to use weighted searchlights. Defaults to True.
+    Parameters
+    ----------
+    X : array-like of shape (n_samples, n_features)
+        The source data matrix.
+    Y : array-like of shape (n_samples, n_features)
+        The target data matrix.
+    searchlights : array-like
+        The indices of the searchlight regions.
+    sl_func : function
+        The function to use for hyperalignment.
+    weights : array-like, optional
+        The weights to use for weighted hyperalignment. Defaults to None.
+    return_betas : bool, optional
+        Whether to return the coefficients of regression instead of the prediciton.
+        Defaults to False.
+    verbose : bool, optional
+        Whether to display progress. Defaults to False.
+
+    Returns
+    -------
+    res : array-like
+        The transformed data matrix.
 
-    Returns:
-        ndarray: The transformed matrix T of shape (n_features, n_features).
     """
-    Yhat = np.zeros_like(X, dtype=np.float32)
+    if return_betas:
+        T = np.zeros((X.shape[1], Y.shape[1]), dtype=np.float32)
+    else:
+        Yhat = np.zeros_like(X, dtype=np.float32)
 
     if weights is not None:
         zip_iter = zip(searchlights, weights)
         for sl, w in zip_iter:
             x, y = X[:, sl], Y[:, sl]
             t = sl_func(x, y)
+        if return_betas:
+            T[np.ix_(sl, sl)] += t * w[np.newaxis]
+        else:
             Yhat[:, sl] += x @ t * w[np.newaxis]
-            del t
+
     else:
         searchlights_iter = searchlights
         for sl in searchlights_iter:
             x, y = X[:, sl], Y[:, sl]
             t = sl_func(x, y)
-            Yhat[:, sl] += x @ t
-            del t
+            if return_betas:
+                T[np.ix_(sl, sl)] += t
+            else:
+                Yhat[:, sl] += x @ t
 
-    return Yhat
+    res = T if return_betas else Yhat
+    return res
 
 
 def piece_procrustes(
@@ -438,6 +460,7 @@ def piece_ridge(
     alpha=1e3,
     weights=None,
     verbose=False,
+    return_betas=False,
 ):
     """
     Perform searchlight ridge regression for hyperalignment.
@@ -464,6 +487,7 @@ def piece_ridge(
         sl_func=sl_func,
         weights=weights,
         verbose=verbose,
+        return_betas=return_betas,
     )
     return T
 

From 9b38655087f0c99d87e91343abc8b6703fc02944 Mon Sep 17 00:00:00 2001
From: FOUCHARD Denis <denis.fouchard@inria.fr>
Date: Fri, 19 Jan 2024 15:12:38 +0100
Subject: [PATCH 34/69] Add _ridge function for ridge regression

---
 fmralign/hyperalignment/linalg.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/fmralign/hyperalignment/linalg.py b/fmralign/hyperalignment/linalg.py
index 71c30ed..dc02847 100644
--- a/fmralign/hyperalignment/linalg.py
+++ b/fmralign/hyperalignment/linalg.py
@@ -108,6 +108,14 @@ def ridge(X, Y, alpha):
     return betas
 
 
+def _ridge(X, Y, alpha):
+    from sklearn.linear_model import Ridge
+
+    ridge = Ridge(alpha=alpha, fit_intercept=False)
+    ridge.fit(X, Y)
+    return ridge.coef_.T
+
+
 def procrustes(X, Y, reflection=True, scaling=False):
     r"""
     The orthogonal Procrustes algorithm.

From 1ee469c2a1dd75c999ad8b1302cda21792dd1212 Mon Sep 17 00:00:00 2001
From: FOUCHARD Denis <denis.fouchard@inria.fr>
Date: Fri, 19 Jan 2024 15:13:11 +0100
Subject: [PATCH 35/69] Fix indentation in piece_ridge function

---
 fmralign/hyperalignment/regions.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/fmralign/hyperalignment/regions.py b/fmralign/hyperalignment/regions.py
index eaf51b3..0f7e333 100644
--- a/fmralign/hyperalignment/regions.py
+++ b/fmralign/hyperalignment/regions.py
@@ -480,6 +480,7 @@ def piece_ridge(
 
     """
     sl_func = functools.partial(ridge, alpha=alpha)
+
     T = iter_hyperalignment(
         X,
         Y,

From 29faee8814ec496fb3f41e978b7ecf057635789e Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Fri, 2 Feb 2024 19:05:03 +0100
Subject: [PATCH 36/69] some tweaks

---
 fmralign/alignment_methods.py                 | 12 +++----
 fmralign/hyperalignment/linalg.py             |  7 ++--
 fmralign/hyperalignment/local_template.py     | 11 +++---
 .../hyperalignment/piecewise_alignment.py     |  4 +++
 fmralign/hyperalignment/regions.py            |  3 +-
 .../hyperalignment/test_hyperalignment.py     | 34 +++++++++----------
 6 files changed, 38 insertions(+), 33 deletions(-)

diff --git a/fmralign/alignment_methods.py b/fmralign/alignment_methods.py
index 2f95198..9e7a6d3 100644
--- a/fmralign/alignment_methods.py
+++ b/fmralign/alignment_methods.py
@@ -521,8 +521,8 @@ def __init__(
         self.tuning_data = []
         self.denoised_signal = []
         self.decomp_method = decomp_method
-        self.tmpl_kind = template
-        self.latent_dim = n_components
+        self.template = template
+        self.n_components = n_components
         self.n_jobs = n_jobs
 
     #######################################################################################
@@ -565,11 +565,9 @@ def _stimulus_estimator(full_signal, n_t, n_s, latent_dim=None):
         Returns:
             stimulus (np.ndarray): The stimulus response of shape (n_t, latent_dim) or (n_t, n_t).
         """
+        U = svd_pca(full_signal)
         if latent_dim is not None and latent_dim < n_t:
-            U = svd_pca(full_signal)
             U = U[:, :latent_dim]
-        else:
-            U, _, _ = safe_svd(full_signal, remove_mean=False)
 
         stimulus = np.sqrt(n_s) * U
         stimulus = stimulus.astype(np.float32)
@@ -659,7 +657,7 @@ def fit(
         if self.decomp_method is None:
             full_signal = np.concatenate(self.denoised_signal, axis=1)
             self.shared_response = self._stimulus_estimator(
-                full_signal, self.n_t, self.n_s, self.latent_dim
+                full_signal, self.n_t, self.n_s, self.n_components
             )
             if tuning:
                 self.tuning_data = Parallel(n_jobs=self.n_jobs)(
@@ -695,7 +693,7 @@ def transform(self, X, verbose=False):
 
         if self.decomp_method is None:
             stimulus_ = self._stimulus_estimator(
-                full_signal, self.n_t, self.n_s, self.latent_dim
+                full_signal, self.n_t, self.n_s, self.n_components
             )
 
         if verbose:
diff --git a/fmralign/hyperalignment/linalg.py b/fmralign/hyperalignment/linalg.py
index dc02847..c0ae5ff 100644
--- a/fmralign/hyperalignment/linalg.py
+++ b/fmralign/hyperalignment/linalg.py
@@ -8,7 +8,8 @@
 """
 
 import numpy as np
-from scipy.linalg import svd, LinAlgError
+from scipy.linalg import LinAlgError
+from scipy.linalg import svd
 
 __all__ = ["safe_svd", "svd_pca", "ridge"]
 
@@ -86,7 +87,7 @@ def svd_pca(X, remove_mean=True):
     return X_new
 
 
-def ridge(X, Y, alpha):
+def ridge(X, Y, alpha=10):
     """Solve ridge regression problem for matrix target using SVD.
     Parameters
     ----------
@@ -101,7 +102,7 @@ def ridge(X, Y, alpha):
     betas : ndarray of shape (n_features, n_targets)
         The solution to the ridge regression problem.
     """
-    U, s, Vt = safe_svd(X, remove_mean=False)
+    U, s, Vt = safe_svd(X, remove_mean=True)
     d = s / (alpha + s**2)
     d_UT_Y = d[:, np.newaxis] * (U.T @ Y)
     betas = Vt.T @ d_UT_Y
diff --git a/fmralign/hyperalignment/local_template.py b/fmralign/hyperalignment/local_template.py
index 1d8c11b..34f7943 100644
--- a/fmralign/hyperalignment/local_template.py
+++ b/fmralign/hyperalignment/local_template.py
@@ -14,7 +14,7 @@
 
 
 def PCA_decomposition(
-    X, n_components=None, flavor="sklearn", adjust_ns=False, demean=True
+    X, n_components=None, flavor="sklearn", adjust_ns=True, demean=True
 ):
     """Decompose concatenated data matrices using PCA/SVD.
 
@@ -106,12 +106,13 @@ def compute_PCA_template(X, sl=None, n_components=None, flavor="sklearn", demean
         The PCA template array of shape (n_samples, n_features, n_components).
     """
     if sl is not None:
-        X = X[:, :, sl]
+        X_ = X[:, :, sl]
     else:
-        X = X
-    n_components = min(X.shape[1], X.shape[2])
+        X_ = X
+    n = min(X_.shape[1], X_.shape[2])
+    n_components = min(n, n_components)
     XX, cc = PCA_decomposition(
-        X, n_components=n_components, flavor=flavor, adjust_ns=True, demean=demean
+        X_, n_components=n_components, flavor=flavor, adjust_ns=True, demean=demean
     )
     return XX.astype(np.float32)
 
diff --git a/fmralign/hyperalignment/piecewise_alignment.py b/fmralign/hyperalignment/piecewise_alignment.py
index ed3153c..35fa6c3 100644
--- a/fmralign/hyperalignment/piecewise_alignment.py
+++ b/fmralign/hyperalignment/piecewise_alignment.py
@@ -27,6 +27,7 @@ def __init__(
         self,
         alignment_method="searchlight_ridge",
         template_kind="pca",
+        common_topography=True,
         verbose=True,
         n_jobs=1,
     ):
@@ -51,6 +52,7 @@ def __init__(
         self.warp_alignment_method = alignment_method
         self.template_kind = template_kind
         self.verbose = verbose
+        self.common_topography = common_topography
         self.n_jobs = n_jobs
         self.regions = None
         self.distances = None
@@ -77,6 +79,7 @@ def compute_linear_transformation(self, data, template):
         x_hat = piece_ridge(
             X=template,
             Y=data,
+            alpha=10,
             regions=self.regions,
             weights=self.weights,
             verbose=self.verbose,
@@ -146,6 +149,7 @@ def fit_transform(
             n_jobs=self.n_jobs,
             template_kind=self.template_kind,
             verbose=self.verbose,
+            common_topography=self.common_topography,
             weights=self.weights,
         )
 
diff --git a/fmralign/hyperalignment/regions.py b/fmralign/hyperalignment/regions.py
index 0f7e333..81f432b 100644
--- a/fmralign/hyperalignment/regions.py
+++ b/fmralign/hyperalignment/regions.py
@@ -499,6 +499,7 @@ def template(
     n_jobs=1,
     template_kind="pca",
     verbose=False,
+    common_topography=True,
     weights=None,
 ):
     """
@@ -522,7 +523,7 @@ def template(
                 region=region,
                 kind=template_kind,
                 n_components=150,
-                common_topography=True,
+                common_topography=common_topography,
             )
             for region in regions
         )
diff --git a/fmralign/hyperalignment/test_hyperalignment.py b/fmralign/hyperalignment/test_hyperalignment.py
index 7057926..aee883b 100644
--- a/fmralign/hyperalignment/test_hyperalignment.py
+++ b/fmralign/hyperalignment/test_hyperalignment.py
@@ -30,7 +30,7 @@ def test_int_fit_predict():
     int1 = INT(n_components=6)
     int2 = INT(n_components=6)
     int1.fit(
-        X_train=X_train, searchlights=searchlights, dists=dists
+        X_train, searchlights=searchlights, dists=dists
     )  # S is provided if we cheat and know the ground truth
     int2.fit(X_test, searchlights=searchlights, dists=dists)
 
@@ -68,7 +68,7 @@ def test_int_fit_predict():
 
 
 def test_int_with_searchlight():
-    X_train, X_test, S_true_first_part, S_true_second_part, Ts = generate_dummy_signal(
+    X_train, X_test, stimulus_train, stimulus_test, _ = generate_dummy_signal(
         n_s=5,
         n_t=50,
         n_v=300,
@@ -88,23 +88,23 @@ def test_int_with_searchlight():
     )
 
     # Test INT on the two parts of the data (ie different runs of the experiment)
-    int1 = INT(n_components=6)
-    int2 = INT(n_components=6)
-    int1.fit(X_train=X_train, searchlights=searchlights, dists=dists, radius=5)
-    int2.fit(X_test, searchlights=searchlights, dists=dists, radius=5)
-    X_pred = int1.transform(X_test)
-
-    tuning_data_run_1 = int1.tuning_data
-    tuning_data_run_2 = int2.tuning_data
+    model1 = INT(n_components=6)
+    model2 = INT(n_components=6)
+    model1.fit(X_train, searchlights=searchlights, dists=dists, radius=5)
+    model2.fit(X_test, searchlights=searchlights, dists=dists, radius=5)
+    X_pred = model1.transform(X_test)
+
+    tuning_data_run_1 = model1.tuning_data
+    tuning_data_run_2 = model2.tuning_data
     tuning_data_run_1 = np.array(tuning_data_run_1)
     tuning_data_run_2 = np.array(tuning_data_run_2)
 
-    stimulus_run_1 = int1.shared_response
-    stimulus_run_2 = int2.shared_response
+    stimulus_run_1 = model1.shared_response
+    stimulus_run_2 = model2.shared_response
 
     corr1 = tuning_correlation(tuning_data_run_1, tuning_data_run_2)
-    corr2 = stimulus_correlation(stimulus_run_1.T, S_true_first_part.T)
-    corr3 = stimulus_correlation(stimulus_run_2.T, S_true_second_part.T)
+    corr2 = stimulus_correlation(stimulus_run_1.T, stimulus_train.T)
+    corr3 = stimulus_correlation(stimulus_run_2.T, stimulus_test.T)
     corr4 = tuning_correlation(X_pred, X_test)
 
     # Check that predicted components have the same shape as original data
@@ -118,7 +118,7 @@ def test_int_with_searchlight():
     assert 3 * np.mean(corr2_out) < np.mean(np.diag(corr2))
     assert 3 * np.mean(corr3_out) < np.mean(np.diag(corr3))
     assert 3 * np.mean(corr4_out) < np.mean(np.diag(corr4))
-    assert int1.tuning_data[0].shape == (6, int1.n_v)
-    assert int2.tuning_data[0].shape == (6, int2.n_v)
-    assert int1.shared_response.shape == (int1.n_t, 6)
+    assert model1.tuning_data[0].shape == (6, model1.n_v)
+    assert model2.tuning_data[0].shape == (6, model2.n_v)
+    assert model1.shared_response.shape == (model1.n_t, 6)
     assert X_pred.shape == X_test.shape

From 890e4a12e7f679fb0a55efc1d81b33c402d0c90c Mon Sep 17 00:00:00 2001
From: FOUCHARD Denis <denis.fouchard@inria.fr>
Date: Tue, 6 Feb 2024 16:30:57 +0100
Subject: [PATCH 37/69] fix INT plot and last tweaks

---
 examples/plot_hyperalignment.py               | 257 ------------------
 examples/plot_int_alignment.py                |  21 +-
 .../hyperalignment/piecewise_alignment.py     |   2 +-
 fmralign/hyperalignment/regions.py            |  27 +-
 4 files changed, 20 insertions(+), 287 deletions(-)
 delete mode 100644 examples/plot_hyperalignment.py

diff --git a/examples/plot_hyperalignment.py b/examples/plot_hyperalignment.py
deleted file mode 100644
index e3c711b..0000000
--- a/examples/plot_hyperalignment.py
+++ /dev/null
@@ -1,257 +0,0 @@
-# -*- coding: utf-8 -*-
-
-"""
-Hyperalignment-base prediction using the IndividualNeuralTuning Model.
-See article : https://doi.org/10.1162/imag_a_00032
-
-==========================
-
-In this tutorial, we show how to better predict new contrasts for a target
-subject using many source subjects corresponding contrasts. For this purpose,
-we create a template to which we align the target subject, using shared information.
-We then predict new images for the target and compare them to a baseline.
-
-We mostly rely on Python common packages and on nilearn to handle
-functional data in a clean fashion.
-
-
-To run this example, you must launch IPython via ``ipython
---matplotlib`` in a terminal, or use ``jupyter-notebook``.
-
-.. contents:: **Contents**
-    :local:
-    :depth: 1
-
-"""
-# %%
-import warnings
-
-warnings.filterwarnings("ignore")
-###############################################################################
-# Retrieve the data
-# -----------------
-# In this example we use the IBC dataset, which includes a large number of
-# different contrasts maps for 12 subjects.
-# We download the images for subjects sub-01, sub-02, sub-04, sub-05, sub-06
-# and sub-07 (or retrieve them if they were already downloaded).
-# imgs is the list of paths to available statistical images for each subjects.
-# df is a dataframe with metadata about each of them.
-# mask is a binary image used to extract grey matter regions.
-#
-
-from fmralign.fetch_example_data import fetch_ibc_subjects_contrasts
-
-imgs, df, mask_img = fetch_ibc_subjects_contrasts(
-    [
-        "sub-01",
-        "sub-02",
-        "sub-04",
-        "sub-05",
-        "sub-06",
-        "sub-07",
-    ],
-)
-
-SEARCHLIGHT = False
-
-###############################################################################
-# Definine a masker
-# -----------------
-# We define a nilearn masker that will be used to handle relevant data.
-#   For more information, visit :
-#   'http://nilearn.github.io/manipulating_images/masker_objects.html'
-#
-
-from nilearn.maskers import NiftiMasker
-
-masker = NiftiMasker(mask_img=mask_img).fit()
-
-###############################################################################
-# Prepare the data
-# ----------------
-# For each subject, we will use two series of contrasts acquired during
-# two independent sessions with a different phase encoding:
-# Antero-posterior(AP) or Postero-anterior(PA).
-#
-
-
-# To infer a template for subjects sub-01 to sub-06 for both AP and PA data,
-# we make a list of 4D niimgs from our list of list of files containing 3D images
-
-from nilearn.image import concat_imgs
-
-template_train = []
-for i in range(6):
-    template_train.append(concat_imgs(imgs[i]))
-target_train = df[df.subject == "sub-07"][df.acquisition == "ap"].path.values
-
-# For subject sub-07, we split it in two folds:
-#   - target train: sub-07 AP contrasts, used to learn alignment to template
-#   - target test: sub-07 PA contrasts, used as a ground truth to score predictions
-# We make a single 4D Niimg from our list of 3D filenames
-
-target_train = concat_imgs(target_train)
-target_train_data = masker.transform(target_train)
-target_test = df[df.subject == "sub-07"][df.acquisition == "pa"].path.values
-
-
-###############################################################################
-# Compute a baseline (average of subjects)
-# ----------------------------------------
-# We create an image with as many contrasts as any subject representing for
-# each contrast the average of all train subjects maps.
-#
-
-import numpy as np
-
-masked_imgs = [masker.transform(img) for img in template_train]
-average_img = np.mean(masked_imgs[:-1], axis=0)
-average_subject = masker.inverse_transform(average_img)
-
-###############################################################################
-# Create a template from the training subjects.
-# ---------------------------------------------
-# We define an estimator using the class TemplateAlignment:
-#   * We align the whole brain through 'multiple' local alignments.
-#   * These alignments are calculated on a parcellation of the brain in 150 pieces,
-#     this parcellation creates group of functionnally similar voxels.
-#   * The template is created iteratively, aligning all subjects data into a
-#     common space, from which the template is inferred and aligning again to this
-#     new template space.
-#
-
-from fmralign.hyperalignment.regions import compute_parcels, compute_searchlights
-
-###############################################################################
-# Predict new data for left-out subject
-# -------------------------------------
-# We use target_train data to fit the transform, indicating it corresponds to
-# the contrasts indexed by train_index and predict from this learnt alignment
-# contrasts corresponding to template test_index numbers.
-# For each train subject and for the template, the AP contrasts are sorted from
-# 0, to 53, and then the PA contrasts from 53 to 106.
-#
-
-train_index = range(53)
-test_index = range(53, 106)
-
-train_data = np.array(masked_imgs)[:, train_index, :]
-test_data = np.array(masked_imgs)[:, test_index, :][:-1]
-
-if not SEARCHLIGHT:
-    parcels = compute_parcels(
-        niimg=template_train[0], mask=masker, n_parcels=1000, n_jobs=5
-    )
-
-else:
-    _, searchlights, dists = compute_searchlights(
-        niimg=template_train[0], mask_img=masker.mask_img, radius=5, n_jobs=5
-    )
-
-# %%
-
-from fmralign.hyperalignment.regions import template, piece_ridge, searchlight_weights
-
-if not SEARCHLIGHT:
-    sl_template = template(
-        X=train_data, regions=parcels, n_jobs=5, template_kind="procrustes"
-    )
-    train_tuning = piece_ridge(X=sl_template, Y=train_data[-1], regions=parcels)
-    test_template = template(X=test_data, regions=parcels, n_jobs=5)
-
-    train_tuning = piece_ridge(
-        X=sl_template, Y=train_data[-1], regions=parcels, return_betas=True
-    )
-else:
-    weights = searchlight_weights(searchlights=searchlights, dists=dists, radius=20)
-    sl_template = template(
-        X=train_data, regions=searchlights, n_jobs=5, weights=weights
-    )
-    train_tuning = piece_ridge(
-        X=sl_template,
-        Y=train_data[-1],
-        regions=searchlights,
-        weights=weights,
-        alpha=10,
-        return_betas=True,
-    )
-    test_template = template(
-        X=test_data, regions=searchlights, n_jobs=5, weights=weights
-    )
-
-# %%
-
-target_pred = test_template @ train_tuning
-
-target_pred = masker.inverse_transform(target_pred)
-
-# %%
-from fmralign.metrics import score_voxelwise
-
-# Now we use this scoring function to compare the correlation of predictions
-# made from group average and from template with the real PA contrasts of sub-07
-
-
-template_score = masker.inverse_transform(
-    C_temp := score_voxelwise(target_test, target_pred, masker=masker, loss="corr")
-)
-
-print("============Global correlation============")
-print(f"Template avg : {np.mean(C_temp)}")
-print(f"Template std : {np.std(C_temp)}")
-
-###############################################################################
-# Plotting the measures
-# ---------------------
-# Finally we plot both scores
-#
-
-# %%
-from nilearn import plotting
-
-display = plotting.plot_stat_map(
-    template_score, display_mode="z", cut_coords=[-15, -5], vmax=1
-)
-# display.title("Hyperalignment-based prediction correlation wt ground truth")
-
-###############################################################################
-# We observe that creating a template and aligning a new subject to it yields
-# a prediction that is better correlated with the ground truth than just using
-# the average activations of subjects.
-#
-
-
-# %%
-from nilearn.image import load_img
-from nilearn.maskers import NiftiMasker
-
-moviewatching_mask_img = load_img("/Users/df/nilearn_data/movie_watching.nii.gz")
-moviewatching_mask_img = masker.transform(moviewatching_mask_img)
-moviewatching_mask_img -= np.mean(moviewatching_mask_img, axis=0)
-moviewatching_mask_img /= np.std(moviewatching_mask_img, axis=0)
-moviewatching_mask_img = moviewatching_mask_img > 0
-moviewatching_mask_img = masker.inverse_transform(moviewatching_mask_img)
-moviewatching_masker = NiftiMasker(mask_img=moviewatching_mask_img).fit()
-
-
-target_pred = moviewatching_masker.transform(target_pred)
-target_test = moviewatching_masker.transform(target_test)
-
-target_pred = target_pred - np.mean(target_pred, axis=0)
-
-C_temp_localized = score_voxelwise(
-    target_test, target_pred, masker=moviewatching_masker, loss="corr"
-)
-
-template_localized_score = moviewatching_masker.inverse_transform(C_temp)
-
-print("============Global correlation============")
-print(f"Template avg : {np.mean(C_temp)}")
-print(f"Template std : {np.std(C_temp)}")
-
-localized_display = plotting.plot_stat_map(
-    template_score, display_mode="z", cut_coords=[-15, -5], vmax=1
-)
-
-
-plotting.show()
diff --git a/examples/plot_int_alignment.py b/examples/plot_int_alignment.py
index 5bcf3cd..4a6b0f7 100644
--- a/examples/plot_int_alignment.py
+++ b/examples/plot_int_alignment.py
@@ -134,11 +134,13 @@
 
 train_index = range(53)
 test_index = range(53, 106)
+full_index = range(106)
 
 train_data = np.array(masked_imgs)[:, train_index, :]
-test_data = np.array(masked_imgs)[:, test_index, :][:-1]
+test_data = np.array(masked_imgs)[:, full_index, :]
+target_test_masked = np.array(masked_imgs)[:, test_index, :]
 
-if True:
+if False:
     parcels = compute_parcels(
         niimg=template_train[0], mask=masker, n_parcels=100, n_jobs=5
     )
@@ -156,7 +158,7 @@
 
 else:
     _, searchlights, dists = compute_searchlights(
-        niimg=template_train[0], mask_img=masker.mask_img, n_jobs=5
+        niimg=template_train[0], mask_img=masker.mask_img, radius=5, n_jobs=5
     )
     model = IndividualizedNeuralTuning(
         n_jobs=8, alignment_method="searchlight", n_components=None
@@ -167,8 +169,8 @@
         dists=dists,
         verbose=False,
     )
-    train_stimulus = np.copy(model.shared_response)
-    train_tuning = model.tuning_data[-1]
+    train_target_denoised = model.denoised_signal[-1]
+
     model_bis = IndividualizedNeuralTuning(
         n_jobs=8, alignment_method="searchlight", n_components=None
     )
@@ -178,14 +180,17 @@
         dists=dists,
         verbose=False,
     )
-    test_stimulus = np.copy(model_bis.shared_response)
+    train_tuning = (
+        np.linalg.pinv(model_bis.shared_response[train_index]) @ train_target_denoised
+    )
 
 
 # %%
 # We input the mapping image target_train in a list, we could have input more
 # than one subject for which we'd want to predict : [train_1, train_2 ...]
 
-prediction_from_template = -test_stimulus @ train_tuning
+stimulus_ = np.copy(model_bis.shared_response)
+prediction_from_template = stimulus_[test_index] @ train_tuning
 prediction_from_template = masker.inverse_transform(prediction_from_template)
 
 
@@ -238,7 +243,7 @@
 display = plotting.plot_stat_map(
     template_score, display_mode="z", cut_coords=[-15, -5], vmax=1
 )
-display.title("Hyperalignment-based prediction correlation wt ground truth")
+display.title("INT prediction correlation wt ground truth")
 
 ###############################################################################
 # We observe that creating a template and aligning a new subject to it yields
diff --git a/fmralign/hyperalignment/piecewise_alignment.py b/fmralign/hyperalignment/piecewise_alignment.py
index 35fa6c3..a51dd7f 100644
--- a/fmralign/hyperalignment/piecewise_alignment.py
+++ b/fmralign/hyperalignment/piecewise_alignment.py
@@ -4,6 +4,7 @@
 - https://nilearn.github.io/modules/generated/nilearn.regions.Parcellations.html
 - https://nilearn.github.io/dev/modules/generated/nilearn.decoding.SearchLight.html
 """
+
 import numpy as np
 from sklearn.base import BaseEstimator, TransformerMixin
 from .regions import (
@@ -81,7 +82,6 @@ def compute_linear_transformation(self, data, template):
             Y=data,
             alpha=10,
             regions=self.regions,
-            weights=self.weights,
             verbose=self.verbose,
         )
         return x_hat
diff --git a/fmralign/hyperalignment/regions.py b/fmralign/hyperalignment/regions.py
index 81f432b..2bec23c 100644
--- a/fmralign/hyperalignment/regions.py
+++ b/fmralign/hyperalignment/regions.py
@@ -354,7 +354,6 @@ def iter_hyperalignment(
     Y,
     searchlights,
     sl_func,
-    weights=None,
     return_betas=False,
     verbose=False,
 ):
@@ -390,25 +389,14 @@ def iter_hyperalignment(
     else:
         Yhat = np.zeros_like(X, dtype=np.float32)
 
-    if weights is not None:
-        zip_iter = zip(searchlights, weights)
-        for sl, w in zip_iter:
-            x, y = X[:, sl], Y[:, sl]
-            t = sl_func(x, y)
+    searchlights_iter = searchlights
+    for sl in searchlights_iter:
+        x, y = X[:, sl], Y[:, sl]
+        t = sl_func(x, y)
         if return_betas:
-            T[np.ix_(sl, sl)] += t * w[np.newaxis]
+            T[np.ix_(sl, sl)] += t
         else:
-            Yhat[:, sl] += x @ t * w[np.newaxis]
-
-    else:
-        searchlights_iter = searchlights
-        for sl in searchlights_iter:
-            x, y = X[:, sl], Y[:, sl]
-            t = sl_func(x, y)
-            if return_betas:
-                T[np.ix_(sl, sl)] += t
-            else:
-                Yhat[:, sl] += x @ t
+            Yhat[:, sl] += x @ t
 
     res = T if return_betas else Yhat
     return res
@@ -418,7 +406,6 @@ def piece_procrustes(
     X,
     Y,
     regions,
-    weights=None,
     T0=None,
     reflection=True,
     scaling=False,
@@ -447,7 +434,6 @@ def piece_procrustes(
         Y,
         regions,
         T0=T0,
-        weights=weights,
         sl_func=sl_func,
     )
     return T
@@ -486,7 +472,6 @@ def piece_ridge(
         Y,
         regions,
         sl_func=sl_func,
-        weights=weights,
         verbose=verbose,
         return_betas=return_betas,
     )

From e38a144a7557071d9dcde7e29cfafc41db9a7a9c Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Tue, 13 Feb 2024 09:59:25 +0100
Subject: [PATCH 38/69] Merging and correcting PR remarks

---
 fmralign/alignment_methods.py          | 52 +++++++++++++-------------
 fmralign/hyperalignment/correlation.py | 18 +++++----
 2 files changed, 37 insertions(+), 33 deletions(-)

diff --git a/fmralign/alignment_methods.py b/fmralign/alignment_methods.py
index 9e7a6d3..a8b964d 100644
--- a/fmralign/alignment_methods.py
+++ b/fmralign/alignment_methods.py
@@ -501,9 +501,8 @@ def __init__(
         None
         """
 
-        self.n_s = None
-        self.n_t = None
-        self.n_v = None
+        self.n_subjects = None
+        self.n_time_points = None
         self.labels = None
         self.alphas = None
         self.alignment_method = alignment_method
@@ -521,7 +520,6 @@ def __init__(
         self.tuning_data = []
         self.denoised_signal = []
         self.decomp_method = decomp_method
-        self.template = template
         self.n_components = n_components
         self.n_jobs = n_jobs
 
@@ -531,20 +529,17 @@ def __init__(
     @staticmethod
     def _tuning_estimator(shared_response, target):
         """
-        Estimate the tuning weights for individualized neural tuning.
+        Estimate the tuning matrix for individualized neural tuning.
 
         Parameters:
         --------
-            - shared_response (array-like):
-                The shared response matrix.
-            - target (array-like):
-                The target matrix.
-            - latent_dim (int, optional):
-                The number of latent dimensions. Defaults to None.
+        - shared_response (array-like): The shared response matrix.
+        - target (array-like): The target matrix.
+        - latent_dim (int, optional): The number of latent dimensions (if PCA is used). Defaults to None.
 
         Returns:
         --------
-            array-like: The estimated tuning weights.
+        array-like: The estimated tuning matrix for the given target.
 
         """
         if shared_response.shape[1] < shared_response.shape[0]:
@@ -552,20 +547,22 @@ def _tuning_estimator(shared_response, target):
         return np.linalg.inv(shared_response).dot(target)
 
     @staticmethod
-    def _stimulus_estimator(full_signal, n_t, n_s, latent_dim=None):
+    def _stimulus_estimator(full_signal, n_t, n_s, latent_dim=None, scaling=True):
         """
-        Estimates the stimulus response using the given parameters.
+        Estimates the stimulus matrix for the Individualized Neural Tuning model.
 
-        Args:
-            full_signal (np.ndarray): The full signal data.
-            n_t (int): The number of time points.
-            n_s (int): The number of stimuli.
-            latent_dim (int, optional): The number of latent dimensions. Defaults to None.
-
-        Returns:
-            stimulus (np.ndarray): The stimulus response of shape (n_t, latent_dim) or (n_t, n_t).
+        Parameters:
+        --------
+        - full_signal (numpy.ndarray): The full signal of shape (n_t, n_s).
+        - n_t (int): The number of time points.
+        - n_s (int): The number of subjects.
+        - latent_dim (int, optional): The number of latent dimensions to use. Defaults to None.
+        - scaling (bool, optional): Whether to scale the stimulus matrix sources. Defaults to True.
         """
-        U = svd_pca(full_signal)
+        if scaling:
+            U = svd_pca(full_signal)
+        else:
+            U, _, _ = safe_svd(full_signal)
         if latent_dim is not None and latent_dim < n_t:
             U = U[:, :latent_dim]
 
@@ -699,7 +696,10 @@ def transform(self, X, verbose=False):
         if verbose:
             print("Predict : stimulus matrix shape: ", stimulus_.shape)
 
-        reconstructed_signal = [
-            self._reconstruct_signal(stimulus_, T_est) for T_est in self.tuning_data
-        ]
+        reconstructed_signal = Parallel(n_jobs=self.n_jobs)(
+            delayed(
+                self._reconstruct_signal(stimulus_, T_est) for T_est in self.tuning_data
+            )
+        )
+
         return np.array(reconstructed_signal, dtype=np.float32)
diff --git a/fmralign/hyperalignment/correlation.py b/fmralign/hyperalignment/correlation.py
index e03ad12..b7e6cdf 100644
--- a/fmralign/hyperalignment/correlation.py
+++ b/fmralign/hyperalignment/correlation.py
@@ -142,7 +142,7 @@ def tuning_correlation(X, Y):
     corr_mat = np.zeros((n, n))
     for i in range(n):
         for j in range(i, n):
-            corr_i_j = pearson_corr_coeff(X[i], Y[j])
+            corr_i_j = pearson_corr_coeff(X[i], Y[j], absolute=False)
             corr_mat[i, j] = corr_i_j
             corr_mat[j, i] = corr_i_j
 
@@ -260,14 +260,12 @@ def multithread_compute_correlation(
 
     def thread_compute_correlation(X, Y, i, j):
         X_i, Y_i = X[i], Y[i]
-        corr = np.corrcoef(X_i, Y_i)[X.shape[1] :, : X.shape[1]]
-        row_ind, col_ind = linear_sum_assignment(corr, maximize=True)
-        corr = corr[row_ind, :]
-        corr = corr[:, col_ind]
+        corr = stimulus_correlation(
+            X_i, Y_i, absolute=absolute, linear_assignment=linear_assignment
+        )
         same_TR_corr = np.diag(corr)
         # Get all the values except the diagonal in a list
         diff_TR_corr = corr[np.where(~np.eye(corr.shape[0], dtype=bool))]
-        diff_TR_corr = diff_TR_corr.flatten().astype(np.float16)
         if i == j:
             return (
                 np.array([]),
@@ -290,7 +288,13 @@ def thread_compute_correlation(X, Y, i, j):
     n_s = X.shape[0]
     coordinates = list(combinations(range(n_s), 2)) + [(i, i) for i in range(n_s)]
     results = Parallel(n_jobs=n_jobs)(
-        delayed(thread_compute_correlation)(X, Y, i, j) for (i, j) in coordinates
+        delayed(thread_compute_correlation)(
+            X,
+            Y,
+            i,
+            j,
+        )
+        for (i, j) in coordinates
     )
     results = list(zip(*results))
     corr_same_sub_diff_TR = results[2]

From b207f612b9c6ca4ec49f243de31660672c78acd4 Mon Sep 17 00:00:00 2001
From: DF <95576336+denisfouchard@users.noreply.github.com>
Date: Thu, 15 Feb 2024 15:11:26 +0100
Subject: [PATCH 39/69] Update examples/plot_int_alignment.py

Co-authored-by: Pierre-Louis Barbarant <104081777+pbarbarant@users.noreply.github.com>
---
 examples/plot_int_alignment.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/plot_int_alignment.py b/examples/plot_int_alignment.py
index 4a6b0f7..80a9389 100644
--- a/examples/plot_int_alignment.py
+++ b/examples/plot_int_alignment.py
@@ -53,7 +53,7 @@
 )
 
 ###############################################################################
-# Definine a masker
+# Define a masker
 # -----------------
 # We define a nilearn masker that will be used to handle relevant data.
 #   For more information, visit :

From cbbd37eda38a0863ccee3d4f24fc5598d36b6404 Mon Sep 17 00:00:00 2001
From: Denis Fouchard <denis.fouchard@inria.fr>
Date: Thu, 15 Feb 2024 15:12:20 +0100
Subject: [PATCH 40/69] some fixes

---
 fmralign/_utils.py                     |  5 +--
 fmralign/alignment_methods.py          | 60 ++++++++++++--------------
 fmralign/hyperalignment/correlation.py |  2 +-
 fmralign/hyperalignment/regions.py     |  2 +-
 pyproject.toml                         |  4 +-
 5 files changed, 33 insertions(+), 40 deletions(-)

diff --git a/fmralign/_utils.py b/fmralign/_utils.py
index 6c48098..fd99e48 100644
--- a/fmralign/_utils.py
+++ b/fmralign/_utils.py
@@ -147,10 +147,7 @@ def _make_parcellation(
             err.args += (errmsg,)
             raise err
 
-        labels = _apply_mask_fmri(parcellation.labels_img_, masker.mask_img_).astype(
-            int
-        )
-
+        labels = _apply_mask_fmri(parcellation.labels_img_, masker.mask_img).astype(int)
 
     if verbose > 0:
         unique_labels, counts = np.unique(labels, return_counts=True)
diff --git a/fmralign/alignment_methods.py b/fmralign/alignment_methods.py
index a8b964d..ac7b2aa 100644
--- a/fmralign/alignment_methods.py
+++ b/fmralign/alignment_methods.py
@@ -479,7 +479,6 @@ class IndividualizedNeuralTuning(Alignment):
 
     def __init__(
         self,
-        template="pca",
         decomp_method=None,
         n_components=None,
         alignment_method="searchlight",
@@ -490,7 +489,6 @@ def __init__(
         Parameters:
         --------
 
-        - template (str): The type of template to use for alignment. Default is "pca".
         - decomp_method (str): The decomposition method to use. Default is None.
         - alignment_method (str): The alignment method to use. Can be either "searchlight" or "parcelation", Default is "searchlight".
         - n_components (int): The number of latent dimensions to use in the shared stimulus information matrix. Default is None.
@@ -575,11 +573,13 @@ def _reconstruct_signal(shared_response, individual_tuning):
         """
         Reconstructs the signal using the shared response and individual tuning.
 
-        Args:
-            shared_response (numpy.ndarray): The shared response of shape (n_t, n_t) or (n_t, latent_dim).
-            individual_tuning (numpy.ndarray): The individual tuning of shape (latent_dim, n_v) or (n_t, n_v).
+        Parameters:
+        --------
+        - shared_response (numpy.ndarray): The shared response of shape (n_t, n_t) or (n_t, latent_dim).
+        - individual_tuning (numpy.ndarray): The individual tuning of shape (latent_dim, n_v) or (n_t, n_v).
 
         Returns:
+        --------
             numpy.ndarray: The reconstructed signal of shape (n_t, n_v) (same shape as the original signal)
         """
         return (shared_response @ individual_tuning).astype(np.float32)
@@ -651,19 +651,19 @@ def fit(
         )
 
         # Stimulus matrix computation
-        if self.decomp_method is None:
-            full_signal = np.concatenate(self.denoised_signal, axis=1)
-            self.shared_response = self._stimulus_estimator(
-                full_signal, self.n_t, self.n_s, self.n_components
-            )
-            if tuning:
-                self.tuning_data = Parallel(n_jobs=self.n_jobs)(
-                    delayed(self._tuning_estimator)(
-                        self.shared_response,
-                        self.denoised_signal[i],
-                    )
-                    for i in range(self.n_s)
+
+        full_signal = np.concatenate(self.denoised_signal, axis=1)
+        self.shared_response = self._stimulus_estimator(
+            full_signal, self.n_t, self.n_s, self.n_components
+        )
+        if tuning:
+            self.tuning_data = Parallel(n_jobs=self.n_jobs)(
+                delayed(self._tuning_estimator)(
+                    self.shared_response,
+                    self.denoised_signal[i],
                 )
+                for i in range(self.n_s)
+            )
 
         return self
 
@@ -671,16 +671,14 @@ def transform(self, X, verbose=False):
         """
         Transforms the input test data using the hyperalignment model.
 
-        Args:
-            X (list of arrays):
-                The input test data.
-            verbose (bool, optional):
-                Whether to print verbose output. Defaults to False.
-            id (int, optional):
-                Identifier for the transformation. Defaults to None.
+        Parameters:
+        --------
+        - X (array-like): The test data of shape (n_subjects, n_samples, n_voxels).
+        - verbose (bool, optional): Whether to print progress information. Defaults to False.
 
         Returns:
-            numpy.ndarray: The transformed test data.
+        --------
+        - array-like: The transformed data of shape (n_subjects, n_samples, n_voxels).
         """
 
         full_signal = np.concatenate(X, axis=1, dtype=np.float32)
@@ -688,18 +686,16 @@ def transform(self, X, verbose=False):
         if verbose:
             print("Predict : Computing stimulus matrix...")
 
-        if self.decomp_method is None:
-            stimulus_ = self._stimulus_estimator(
-                full_signal, self.n_t, self.n_s, self.n_components
-            )
+        stimulus_ = self._stimulus_estimator(
+            full_signal, self.n_t, self.n_s, self.n_components
+        )
 
         if verbose:
             print("Predict : stimulus matrix shape: ", stimulus_.shape)
 
         reconstructed_signal = Parallel(n_jobs=self.n_jobs)(
-            delayed(
-                self._reconstruct_signal(stimulus_, T_est) for T_est in self.tuning_data
-            )
+            delayed(self._reconstruct_signal)(stimulus_, T_est)
+            for T_est in self.tuning_data
         )
 
         return np.array(reconstructed_signal, dtype=np.float32)
diff --git a/fmralign/hyperalignment/correlation.py b/fmralign/hyperalignment/correlation.py
index b7e6cdf..cb73465 100644
--- a/fmralign/hyperalignment/correlation.py
+++ b/fmralign/hyperalignment/correlation.py
@@ -142,7 +142,7 @@ def tuning_correlation(X, Y):
     corr_mat = np.zeros((n, n))
     for i in range(n):
         for j in range(i, n):
-            corr_i_j = pearson_corr_coeff(X[i], Y[j], absolute=False)
+            corr_i_j = pearson_corr_coeff(X[i], Y[j], absolute=True)
             corr_mat[i, j] = corr_i_j
             corr_mat[j, i] = corr_i_j
 
diff --git a/fmralign/hyperalignment/regions.py b/fmralign/hyperalignment/regions.py
index 2bec23c..a4e9679 100644
--- a/fmralign/hyperalignment/regions.py
+++ b/fmralign/hyperalignment/regions.py
@@ -16,7 +16,7 @@
 from sklearn import neighbors
 from scipy.spatial import distance_matrix
 from nilearn._utils.niimg_conversions import (
-    safe_get_data,
+    _safe_get_data as safe_get_data,
 )
 from .linalg import procrustes
 from .linalg import ridge
diff --git a/pyproject.toml b/pyproject.toml
index 1623dd5..33cca82 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -23,11 +23,11 @@ dependencies = [
     "scikit-learn",
     "joblib",
     "scipy",
-    "nibabel"
+    "nibabel",
     "nilearn>=0.10.2",
     "POT",
     "fastsrm"
-
+]
 dynamic = ["version"]
 
 [project.optional-dependencies]

From 41d243afe41aa2cff15e95152305c4f0bda605ef Mon Sep 17 00:00:00 2001
From: Elizabeth DuPre <emd222@cornell.edu>
Date: Thu, 31 Aug 2023 14:30:20 -0700
Subject: [PATCH 41/69] Adressing latest PR comments

---
 .gitignore                              |   3 +-
 examples/plot_int_alignment.py          |  71 ++------
 examples/plot_pairwise_roi_alignment.py |   4 +-
 examples/plot_srm_alignment.py          | 205 ++++++++++++++++++++++++
 fmralign/__init__.py                    |   2 +-
 fmralign/_utils.py                      |   9 +-
 fmralign/alignment_methods.py           |  96 ++++++-----
 7 files changed, 284 insertions(+), 106 deletions(-)
 create mode 100644 examples/plot_srm_alignment.py

diff --git a/.gitignore b/.gitignore
index d023765..5c4fd70 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,5 @@
+.DS_Store
+
 # hatchling
 fmralign/_version.py
 fmarlign/__pycache__
@@ -93,4 +95,3 @@ venv.bak/
 
 # Rope project settings
 .ropeproject
-examples/plot_srm_alignment.py
diff --git a/examples/plot_int_alignment.py b/examples/plot_int_alignment.py
index 80a9389..e12a2da 100644
--- a/examples/plot_int_alignment.py
+++ b/examples/plot_int_alignment.py
@@ -110,8 +110,8 @@
 # Create a template from the training subjects.
 # ---------------------------------------------
 # We define an estimator using the class TemplateAlignment:
-#   * We align the whole brain through 'multiple' local alignments.
-#   * These alignments are calculated on a parcellation of the brain in 150 pieces,
+#   * We align the whole brain through multiple local alignments.
+#   * These alignments are calculated on a parcellation of the brain in 100 pieces,
 #     this parcellation creates group of functionnally similar voxels.
 #   * The template is created iteratively, aligning all subjects data into a
 #     common space, from which the template is inferred and aligning again to this
@@ -120,7 +120,7 @@
 
 from nilearn.image import index_img
 from fmralign.alignment_methods import IndividualizedNeuralTuning
-from fmralign.hyperalignment.regions import compute_parcels, compute_searchlights
+from fmralign.hyperalignment.regions import compute_parcels
 
 ###############################################################################
 # Predict new data for left-out subject
@@ -140,51 +140,19 @@
 test_data = np.array(masked_imgs)[:, full_index, :]
 target_test_masked = np.array(masked_imgs)[:, test_index, :]
 
-if False:
-    parcels = compute_parcels(
-        niimg=template_train[0], mask=masker, n_parcels=100, n_jobs=5
-    )
-    model = IndividualizedNeuralTuning(
-        n_jobs=8, alignment_method="parcellation", n_components=None
-    )
-    model.fit(train_data, parcels=parcels, verbose=False)
-    train_stimulus = np.copy(model.shared_response)
-    train_tuning = np.linalg.pinv(train_stimulus) @ model.denoised_signal[-1]
-    model_bis = IndividualizedNeuralTuning(
-        n_jobs=8, alignment_method="parcellation", n_components=None
-    )
-    model_bis.fit(test_data, parcels=parcels, verbose=False)
-    test_stimulus = np.copy(model_bis.shared_response)
-
-else:
-    _, searchlights, dists = compute_searchlights(
-        niimg=template_train[0], mask_img=masker.mask_img, radius=5, n_jobs=5
-    )
-    model = IndividualizedNeuralTuning(
-        n_jobs=8, alignment_method="searchlight", n_components=None
-    )
-    model.fit(
-        train_data,
-        searchlights=searchlights,
-        dists=dists,
-        verbose=False,
-    )
-    train_target_denoised = model.denoised_signal[-1]
-
-    model_bis = IndividualizedNeuralTuning(
-        n_jobs=8, alignment_method="searchlight", n_components=None
-    )
-    model_bis.fit(
-        test_data,
-        searchlights=searchlights,
-        dists=dists,
-        verbose=False,
-    )
-    train_tuning = (
-        np.linalg.pinv(model_bis.shared_response[train_index]) @ train_target_denoised
-    )
-
 
+parcels = compute_parcels(niimg=template_train[0], mask=masker, n_parcels=100, n_jobs=5)
+model = IndividualizedNeuralTuning(
+    n_jobs=8, alignment_method="parcellation", n_components=None
+)
+model.fit(train_data, parcels=parcels, verbose=False)
+train_stimulus = np.copy(model.shared_response)
+train_tuning = np.linalg.pinv(train_stimulus) @ model.denoised_signal[-1]
+model_bis = IndividualizedNeuralTuning(
+    n_jobs=8, alignment_method="parcellation", n_components=None
+)
+model_bis.fit(test_data, parcels=parcels, verbose=False)
+test_stimulus = np.copy(model_bis.shared_response)
 # %%
 # We input the mapping image target_train in a list, we could have input more
 # than one subject for which we'd want to predict : [train_1, train_2 ...]
@@ -214,18 +182,13 @@
 
 
 average_score = masker.inverse_transform(
-    C_avg := score_voxelwise(target_test, prediction_from_average, masker, loss="corr")
+    score_voxelwise(target_test, prediction_from_average, masker, loss="corr")
 )
 
 template_score = masker.inverse_transform(
-    C_temp := score_voxelwise(
-        target_test, prediction_from_template, masker, loss="corr"
-    )
+    score_voxelwise(target_test, prediction_from_template, masker, loss="corr")
 )
 
-print("============Mean correlation============")
-print(f"Baseline : {np.mean(C_avg)}")
-print(f"Template : {np.mean(C_temp)}")
 
 ###############################################################################
 # Plotting the measures
diff --git a/examples/plot_pairwise_roi_alignment.py b/examples/plot_pairwise_roi_alignment.py
index 0baa6ea..5c6f365 100644
--- a/examples/plot_pairwise_roi_alignment.py
+++ b/examples/plot_pairwise_roi_alignment.py
@@ -54,9 +54,7 @@
 # Select visual cortex, create a mask and resample it to the right resolution
 
 mask_visual = new_img_like(atlas, atlas.get_fdata() == 1)
-resampled_mask_visual = resample_to_img(
-    mask_visual, mask, interpolation="nearest"
-)
+resampled_mask_visual = resample_to_img(mask_visual, mask, interpolation="nearest")
 
 # Plot the mask we will use
 plot_roi(
diff --git a/examples/plot_srm_alignment.py b/examples/plot_srm_alignment.py
new file mode 100644
index 0000000..643e8d5
--- /dev/null
+++ b/examples/plot_srm_alignment.py
@@ -0,0 +1,205 @@
+# -*- coding: utf-8 -*-
+
+"""
+Hyperalignment-base prediction using Feilong Ma's IndividualNeuralTuning Model.
+See article : https://doi.org/10.1162/imag_a_00032
+
+==========================
+
+In this tutorial, we show how to better predict new contrasts for a target
+subject using many source subjects corresponding contrasts. For this purpose,
+we create a template to which we align the target subject, using shared information.
+We then predict new images for the target and compare them to a baseline.
+
+We mostly rely on Python common packages and on nilearn to handle
+functional data in a clean fashion.
+
+
+To run this example, you must launch IPython via ``ipython
+--matplotlib`` in a terminal, or use ``jupyter-notebook``.
+
+.. contents:: **Contents**
+    :local:
+    :depth: 1
+
+"""
+
+###############################################################################
+# Retrieve the data
+# -----------------
+# In this example we use the IBC dataset, which includes a large number of
+# different contrasts maps for 12 subjects.
+# We download the images for subjects sub-01, sub-02, sub-04, sub-05, sub-06
+# and sub-07 (or retrieve them if they were already downloaded).
+# imgs is the list of paths to available statistical images for each subjects.
+# df is a dataframe with metadata about each of them.
+# mask is a binary image used to extract grey matter regions.
+#
+
+from fmralign.fetch_example_data import fetch_ibc_subjects_contrasts
+
+imgs, df, mask_img = fetch_ibc_subjects_contrasts(
+    ["sub-01", "sub-02", "sub-04", "sub-05", "sub-06", "sub-07"]
+)
+
+###############################################################################
+# Definine a masker
+# -----------------
+# We define a nilearn masker that will be used to handle relevant data.
+#   For more information, visit :
+#   'http://nilearn.github.io/manipulating_images/masker_objects.html'
+#
+
+from nilearn.maskers import NiftiMasker
+
+masker = NiftiMasker(mask_img=mask_img).fit()
+
+###############################################################################
+# Prepare the data
+# ----------------
+# For each subject, we will use two series of contrasts acquired during
+# two independent sessions with a different phase encoding:
+# Antero-posterior(AP) or Postero-anterior(PA).
+#
+
+
+# To infer a template for subjects sub-01 to sub-06 for both AP and PA data,
+# we make a list of 4D niimgs from our list of list of files containing 3D images
+
+from nilearn.image import concat_imgs
+
+template_train = []
+for i in range(6):
+    template_train.append(concat_imgs(imgs[i]))
+target_train = df[df.subject == "sub-07"][df.acquisition == "ap"].path.values
+
+# For subject sub-07, we split it in two folds:
+#   - target train: sub-07 AP contrasts, used to learn alignment to template
+#   - target test: sub-07 PA contrasts, used as a ground truth to score predictions
+# We make a single 4D Niimg from our list of 3D filenames
+
+target_train = concat_imgs(target_train)
+target_train_data = masker.transform(target_train)
+target_test = df[df.subject == "sub-07"][df.acquisition == "pa"].path.values
+
+###############################################################################
+# Compute a baseline (average of subjects)
+# ----------------------------------------
+# We create an image with as many contrasts as any subject representing for
+# each contrast the average of all train subjects maps.
+#
+
+import numpy as np
+
+masked_imgs = [masker.transform(img) for img in template_train]
+average_img = np.mean(masked_imgs[:-1], axis=0)
+average_subject = masker.inverse_transform(average_img)
+
+###############################################################################
+# Create a template from the training subjects.
+# ---------------------------------------------
+# We define an estimator using the class TemplateAlignment:
+#   * We align the whole brain through 'multiple' local alignments.
+#   * These alignments are calculated on a parcellation of the brain in 150 pieces,
+#     this parcellation creates group of functionnally similar voxels.
+#   * The template is created iteratively, aligning all subjects data into a
+#     common space, from which the template is inferred and aligning again to this
+#     new template space.
+#
+
+from nilearn.image import index_img
+from fastsrm.identifiable_srm import IdentifiableFastSRM as SRM
+
+
+train_index = range(53)
+model = SRM(n_components=20, n_iter=500)
+
+train_imgs = np.array(masked_imgs)[:, train_index, :]
+train_imgs = [x.T for x in train_imgs]
+model.fit(train_imgs)
+
+
+###############################################################################
+# Predict new data for left-out subject
+# -------------------------------------
+# We use target_train data to fit the transform, indicating it corresponds to
+# the contrasts indexed by train_index and predict from this learnt alignment
+# contrasts corresponding to template test_index numbers.
+# For each train subject and for the template, the AP contrasts are sorted from
+# 0, to 53, and then the PA contrasts from 53 to 106.
+#
+test_index = range(53, 106)
+
+test_imgs = np.array(masked_imgs)[:, test_index, :][:-1]
+test_imgs = [x.T for x in test_imgs]
+
+model2 = SRM(n_components=20, n_iter=500)
+shared_response = model2.fit_transform(test_imgs)
+
+# We input the mapping image target_train in a list, we could have input more
+# than one subject for which we'd want to predict : [train_1, train_2 ...]
+
+prediction_from_template = model.inverse_transform(
+    shared_response, subjects_indexes=[5]
+)
+
+prediction_from_template = prediction_from_template[0].T
+prediction_from_template = [masker.inverse_transform(prediction_from_template)]
+
+# As a baseline prediction, let's just take the average of activations across subjects.
+
+prediction_from_average = index_img(average_subject, test_index)
+
+###############################################################################
+# Score the baseline and the prediction
+# -------------------------------------
+# We use a utility scoring function to measure the voxelwise correlation
+# between the prediction and the ground truth. That is, for each voxel, we
+# measure the correlation between its profile of activation without and with
+# alignment, to see if alignment was able to predict a signal more alike the ground truth.
+#
+
+from fmralign.metrics import score_voxelwise
+
+# Now we use this scoring function to compare the correlation of predictions
+# made from group average and from template with the real PA contrasts of sub-07
+
+average_score = masker.inverse_transform(
+    C_avg := score_voxelwise(target_test, prediction_from_average, masker, loss="corr")
+)
+
+# I choose abs value in reference to the work we did with the INT
+template_score = masker.inverse_transform(
+    C_temp := score_voxelwise(
+        target_test, prediction_from_template[0], masker, loss="corr"
+    )
+)
+
+print("============Mean correlation============")
+print(f"Baseline : {np.mean(C_avg)}")
+print(f"Template : {np.mean(C_temp)}")
+
+###############################################################################
+# Plotting the measures
+# ---------------------
+# Finally we plot both scores
+#
+
+from nilearn import plotting
+
+baseline_display = plotting.plot_stat_map(
+    average_score, display_mode="z", vmax=1, cut_coords=[-15, -5]
+)
+baseline_display.title("Group average correlation wt ground truth")
+display = plotting.plot_stat_map(
+    template_score, display_mode="z", cut_coords=[-15, -5], vmax=1
+)
+display.title("SRM-based prediction correlation wt ground truth")
+
+###############################################################################
+# We observe that creating a template and aligning a new subject to it yields
+# a prediction that is better correlated with the ground truth than just using
+# the average activations of subjects.
+#
+
+plotting.show()
diff --git a/fmralign/__init__.py b/fmralign/__init__.py
index 8e6a9c1..3a8d6d5 100644
--- a/fmralign/__init__.py
+++ b/fmralign/__init__.py
@@ -1 +1 @@
-# from ._version import __version__  # noqa: F401
+from ._version import __version__  # noqa: F401
diff --git a/fmralign/_utils.py b/fmralign/_utils.py
index 78b8895..f6130da 100644
--- a/fmralign/_utils.py
+++ b/fmralign/_utils.py
@@ -71,13 +71,7 @@ def _check_labels(labels, threshold=1000):
 
 
 def _make_parcellation(
-    imgs,
-    clustering_index,
-    clustering,
-    n_pieces,
-    masker,
-    smoothing_fwhm=5,
-    verbose=0,
+    imgs, clustering_index, clustering, n_pieces, masker, smoothing_fwhm=5, verbose=0
 ):
     """
     Use nilearn Parcellation class in our pipeline.
@@ -146,7 +140,6 @@ def _make_parcellation(
             )
             err.args += (errmsg,)
             raise err
-
         labels = apply_mask_fmri(parcellation.labels_img_, masker.mask_img_).astype(int)
 
     if verbose > 0:
diff --git a/fmralign/alignment_methods.py b/fmralign/alignment_methods.py
index ac7b2aa..9c4ff80 100644
--- a/fmralign/alignment_methods.py
+++ b/fmralign/alignment_methods.py
@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 """Module implementing alignment estimators on ndarrays."""
 import warnings
+
 import numpy as np
 import scipy
 from joblib import Parallel, delayed
@@ -11,13 +12,10 @@
 from sklearn.base import BaseEstimator, TransformerMixin
 from sklearn.linear_model import RidgeCV
 from sklearn.metrics.pairwise import pairwise_distances
-from .hyperalignment.piecewise_alignment import PiecewiseAlignment
-from .hyperalignment.linalg import safe_svd, svd_pca
 
-import jax
-from ott.geometry import costs, geometry
-from ott.problems.linear import linear_problem
-from ott.solvers.linear import sinkhorn
+# Fast implementation for parallelized computing
+from fmralign.hyperalignment.linalg import safe_svd, svd_pca
+from fmralign.hyperalignment.piecewise_alignment import PiecewiseAlignment
 
 
 def scaled_procrustes(X, Y, scaling=False, primal=None):
@@ -448,6 +446,10 @@ def fit(self, X, Y):
         Y: (n_samples, n_features) nd array
             target data
         """
+        import jax
+        from ott.geometry import costs, geometry
+        from ott.problems.linear import linear_problem
+        from ott.solvers.linear import sinkhorn
 
         if self.metric == "euclidean":
             cost_matrix = costs.Euclidean().all_pairs(x=X.T, y=Y.T)
@@ -471,7 +473,7 @@ def transform(self, X):
 
 class IndividualizedNeuralTuning(Alignment):
     """
-    Method of alignment based on the Individualized Neural Tuning model, by Feilong Ma et al. (2023).
+    Method of alignment based on the Individualized Neural Tuning model.
     It works on 4D fMRI data, and is based on the assumption that the neural response to a stimulus is shared across subjects.
     It uses searchlight/parcelation alignment to denoise the data, and then computes the stimulus response matrix.
     See article : https://doi.org/10.1162/imag_a_00032
@@ -489,10 +491,14 @@ def __init__(
         Parameters:
         --------
 
-        - decomp_method (str): The decomposition method to use. Default is None.
-        - alignment_method (str): The alignment method to use. Can be either "searchlight" or "parcelation", Default is "searchlight".
-        - n_components (int): The number of latent dimensions to use in the shared stimulus information matrix. Default is None.
-        - n_jobs (int): The number of parallel jobs to run. Default is -1.
+        decomp_method : str
+             The decomposition method to use. If None, "pca" will be used. Can be ["pca", "procrustes"] Default is None.
+        alignment_method : str
+             The alignment method to use. Can be either "searchlight" or "parcelation", Default is "searchlight".
+        n_components : int
+             The number of latent dimensions to use in the shared stimulus information matrix. Default is None.
+        n_jobs : int
+             The number of parallel jobs to run. Default is -1.
 
         Returns:
         --------
@@ -507,10 +513,7 @@ def __init__(
         if alignment_method == "parcelation":
             self.parcels = None
 
-        elif (
-            alignment_method == "searchlight"
-            or alignment_method == "ensemble_searchlight"
-        ):
+        elif alignment_method == "searchlight":
             self.searchlights = None
             self.distances = None
             self.radius = None
@@ -531,9 +534,12 @@ def _tuning_estimator(shared_response, target):
 
         Parameters:
         --------
-        - shared_response (array-like): The shared response matrix.
-        - target (array-like): The target matrix.
-        - latent_dim (int, optional): The number of latent dimensions (if PCA is used). Defaults to None.
+        shared_response : array-like
+             The shared response matrix.
+        target : array-like
+             The target matrix.
+        latent_dim : int, optional
+             The number of latent dimensions (if PCA is used). Defaults to None.
 
         Returns:
         --------
@@ -551,11 +557,16 @@ def _stimulus_estimator(full_signal, n_t, n_s, latent_dim=None, scaling=True):
 
         Parameters:
         --------
-        - full_signal (numpy.ndarray): The full signal of shape (n_t, n_s).
-        - n_t (int): The number of time points.
-        - n_s (int): The number of subjects.
-        - latent_dim (int, optional): The number of latent dimensions to use. Defaults to None.
-        - scaling (bool, optional): Whether to scale the stimulus matrix sources. Defaults to True.
+        full_signal : numpy.ndarray
+             The full signal of shape (n_t, n_s).
+        n_t : int
+             The number of time points.
+        n_s : int
+             The number of subjects.
+        latent_dim : int, optional
+             The number of latent dimensions to use. Defaults to None.
+        scaling : bool, optional
+             Whether to scale the stimulus matrix sources. Defaults to True.
         """
         if scaling:
             U = svd_pca(full_signal)
@@ -575,12 +586,14 @@ def _reconstruct_signal(shared_response, individual_tuning):
 
         Parameters:
         --------
-        - shared_response (numpy.ndarray): The shared response of shape (n_t, n_t) or (n_t, latent_dim).
-        - individual_tuning (numpy.ndarray): The individual tuning of shape (latent_dim, n_v) or (n_t, n_v).
+        shared_response : numpy.ndarray
+             The shared response of shape (n_t, n_t) or (n_t, latent_dim).
+        individual_tuning : numpy.ndarray
+             The individual tuning of shape (latent_dim, n_v) or (n_t, n_v).
 
         Returns:
         --------
-            numpy.ndarray: The reconstructed signal of shape (n_t, n_v) (same shape as the original signal)
+        numpy.ndarray: The reconstructed signal of shape (n_t, n_v) (same shape as the original signal)
         """
         return (shared_response @ individual_tuning).astype(np.float32)
 
@@ -600,27 +613,27 @@ def fit(
         Parameters:
         --------
 
-        - X (array-like):
+        X : array-like
             The training data of shape (n_subjects, n_samples, n_voxels).
-        - searchlights (array-like):
+        searchlights : array-like
             The searchlight indices for each subject, of shape (n_s, n_searchlights).
-        - parcels (array-like):
+        parcels : array-like
             The parcel indices for each subject, of shape (n_s, n_parcels) (if not using searchlights)
-        - dists (array-like):
+        dists : array-like
             The distances of vertices to the center of their searchlight, of shape (n_searchlights, n_vertices_sl)
-        - radius (int, optional):
+        radius : int(optional)
             The radius of the searchlight sphere, in milimeters. Defaults to 20.
-        - tuning (bool, optional):
+        tuning :bool(optional)
             Whether to compute the tuning weights. Defaults to True.
-        - verbose (bool, optional):
+        verbose : bool(optional)
             Whether to print progress information. Defaults to True.
-        - id (str, optional):
+        id : str(optional)
             An identifier for caching purposes. Defaults to None.
 
         Returns:
         --------
 
-        - self (IndividualizedNeuralTuning):
+        self : Instance of IndividualizedNeuralTuning)
             The fitted model.
         """
 
@@ -641,7 +654,10 @@ def fit(
             self.radius = radius
 
         denoiser = PiecewiseAlignment(
-            alignment_method=self.alignment_method, n_jobs=self.n_jobs, verbose=verbose
+            alignment_method=self.alignment_method,
+            template_kind=self.decomp_method,
+            n_jobs=self.n_jobs,
+            verbose=verbose,
         )
         self.denoised_signal = denoiser.fit_transform(
             X_,
@@ -673,12 +689,14 @@ def transform(self, X, verbose=False):
 
         Parameters:
         --------
-        - X (array-like): The test data of shape (n_subjects, n_samples, n_voxels).
-        - verbose (bool, optional): Whether to print progress information. Defaults to False.
+        X : array-like
+            The test data of shape (n_subjects, n_samples, n_voxels).
+        verbose : bool(optional)
+            Whether to print progress information. Defaults to False.
 
         Returns:
         --------
-        - array-like: The transformed data of shape (n_subjects, n_samples, n_voxels).
+        array-like: The transformed data of shape (n_subjects, n_samples, n_voxels).
         """
 
         full_signal = np.concatenate(X, axis=1, dtype=np.float32)

From 9b11b864511df1bb2b2f409972850cef396979b8 Mon Sep 17 00:00:00 2001
From: DF <95576336+denisfouchard@users.noreply.github.com>
Date: Mon, 19 Feb 2024 10:16:22 +0100
Subject: [PATCH 42/69] Update examples/plot_int_alignment.py

Co-authored-by: Elizabeth DuPre <emd222@cornell.edu>
---
 examples/plot_int_alignment.py | 12 ++----------
 1 file changed, 2 insertions(+), 10 deletions(-)

diff --git a/examples/plot_int_alignment.py b/examples/plot_int_alignment.py
index e12a2da..baf7f6a 100644
--- a/examples/plot_int_alignment.py
+++ b/examples/plot_int_alignment.py
@@ -41,16 +41,8 @@
 
 from fmralign.fetch_example_data import fetch_ibc_subjects_contrasts
 
-imgs, df, mask_img = fetch_ibc_subjects_contrasts(
-    [
-        "sub-01",
-        "sub-02",
-        "sub-04",
-        "sub-05",
-        "sub-06",
-        "sub-07",
-    ],
-)
+sub_list = ["sub-01", "sub-02", "sub-04", "sub-05", "sub-06", "sub-07"]
+imgs, df, mask_img = fetch_ibc_subjects_contrasts(sub_list)
 
 ###############################################################################
 # Define a masker

From ace93e44a7315b9e07f475ad2df40eabe43096e2 Mon Sep 17 00:00:00 2001
From: DF <95576336+denisfouchard@users.noreply.github.com>
Date: Mon, 19 Feb 2024 10:17:23 +0100
Subject: [PATCH 43/69] Update fmralign/alignment_methods.py

Co-authored-by: Elizabeth DuPre <emd222@cornell.edu>
---
 fmralign/alignment_methods.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/fmralign/alignment_methods.py b/fmralign/alignment_methods.py
index 9c4ff80..61b5f5b 100644
--- a/fmralign/alignment_methods.py
+++ b/fmralign/alignment_methods.py
@@ -488,9 +488,9 @@ def __init__(
     ):
         """
         Initialize the IndividualizedNeuralTuning object.
+
         Parameters:
         --------
-
         decomp_method : str
              The decomposition method to use. If None, "pca" will be used. Can be ["pca", "procrustes"] Default is None.
         alignment_method : str

From ac39cedd0b64a6bac1194387e00c5e869ac36388 Mon Sep 17 00:00:00 2001
From: FOUCHARD Denis <denis.fouchard@inria.fr>
Date: Mon, 19 Feb 2024 14:47:22 +0100
Subject: [PATCH 44/69] Adressing first part of reviews

---
 examples/plot_int_alignment.py                |   4 +-
 examples/plot_srm_alignment.py                |   2 +-
 fmralign/alignment_methods.py                 |  37 +--
 fmralign/fetch_example_data.py                |   8 +-
 fmralign/generate_data.py                     |  48 ++--
 fmralign/hyperalignment/correlation.py        | 214 +++++++++---------
 .../individualized_neural_tuning.py           |  21 +-
 fmralign/hyperalignment/linalg.py             |  36 +--
 fmralign/hyperalignment/local_template.py     | 120 ++++++----
 .../hyperalignment/piecewise_alignment.py     |   1 -
 fmralign/hyperalignment/regions.py            |  95 ++++----
 .../hyperalignment/test_hyperalignment.py     |  24 +-
 fmralign/hyperalignment/toy_experiment.py     |  13 +-
 fmralign/tests/test_alignment_methods.py      |  10 +-
 14 files changed, 323 insertions(+), 310 deletions(-)

diff --git a/examples/plot_int_alignment.py b/examples/plot_int_alignment.py
index baf7f6a..ca93312 100644
--- a/examples/plot_int_alignment.py
+++ b/examples/plot_int_alignment.py
@@ -135,13 +135,13 @@
 
 parcels = compute_parcels(niimg=template_train[0], mask=masker, n_parcels=100, n_jobs=5)
 model = IndividualizedNeuralTuning(
-    n_jobs=8, alignment_method="parcellation", n_components=None
+    n_jobs=8, alignment_method="parcelation", n_components=20
 )
 model.fit(train_data, parcels=parcels, verbose=False)
 train_stimulus = np.copy(model.shared_response)
 train_tuning = np.linalg.pinv(train_stimulus) @ model.denoised_signal[-1]
 model_bis = IndividualizedNeuralTuning(
-    n_jobs=8, alignment_method="parcellation", n_components=None
+    n_jobs=8, alignment_method="parcelation", n_components=20
 )
 model_bis.fit(test_data, parcels=parcels, verbose=False)
 test_stimulus = np.copy(model_bis.shared_response)
diff --git a/examples/plot_srm_alignment.py b/examples/plot_srm_alignment.py
index 643e8d5..3352638 100644
--- a/examples/plot_srm_alignment.py
+++ b/examples/plot_srm_alignment.py
@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 
 """
-Hyperalignment-base prediction using Feilong Ma's IndividualNeuralTuning Model.
+Hyperalignment-base prediction using Hugo Richard's FastSRM.
 See article : https://doi.org/10.1162/imag_a_00032
 
 ==========================
diff --git a/fmralign/alignment_methods.py b/fmralign/alignment_methods.py
index 61b5f5b..a24b63d 100644
--- a/fmralign/alignment_methods.py
+++ b/fmralign/alignment_methods.py
@@ -481,7 +481,7 @@ class IndividualizedNeuralTuning(Alignment):
 
     def __init__(
         self,
-        decomp_method=None,
+        decomp_method="pca",
         n_components=None,
         alignment_method="searchlight",
         n_jobs=1,
@@ -492,7 +492,7 @@ def __init__(
         Parameters:
         --------
         decomp_method : str
-             The decomposition method to use. If None, "pca" will be used. Can be ["pca", "procrustes"] Default is None.
+             The decomposition method to use. Can be ["pca", "pcav1", "procrustes"] Default is "pca".
         alignment_method : str
              The alignment method to use. Can be either "searchlight" or "parcelation", Default is "searchlight".
         n_components : int
@@ -535,7 +535,8 @@ def _tuning_estimator(shared_response, target):
         Parameters:
         --------
         shared_response : array-like
-             The shared response matrix.
+             The shared response matrix of shape (n_timepoints, k)
+             where k is the dimension of the sources latent space.
         target : array-like
              The target matrix.
         latent_dim : int, optional
@@ -551,31 +552,31 @@ def _tuning_estimator(shared_response, target):
         return np.linalg.inv(shared_response).dot(target)
 
     @staticmethod
-    def _stimulus_estimator(full_signal, n_t, n_s, latent_dim=None, scaling=True):
+    def _stimulus_estimator(full_signal, n_subjects, latent_dim=None, scaling=True):
         """
         Estimates the stimulus matrix for the Individualized Neural Tuning model.
 
         Parameters:
         --------
         full_signal : numpy.ndarray
-             The full signal of shape (n_t, n_s).
-        n_t : int
-             The number of time points.
-        n_s : int
+             Concatenated signal for all subjects,
+             of shape (n_timepoints, n_subjects * n_voxels).
+        n_subjects : int
              The number of subjects.
         latent_dim : int, optional
              The number of latent dimensions to use. Defaults to None.
         scaling : bool, optional
              Whether to scale the stimulus matrix sources. Defaults to True.
         """
+        n_timepoints = full_signal.shape[0]
         if scaling:
             U = svd_pca(full_signal)
         else:
             U, _, _ = safe_svd(full_signal)
-        if latent_dim is not None and latent_dim < n_t:
+        if latent_dim is not None and latent_dim < n_timepoints:
             U = U[:, :latent_dim]
 
-        stimulus = np.sqrt(n_s) * U
+        stimulus = np.sqrt(n_subjects) * U
         stimulus = stimulus.astype(np.float32)
         return stimulus
 
@@ -639,10 +640,10 @@ def fit(
 
         X_ = np.array(X, copy=True, dtype=np.float32)
 
-        self.n_s, self.n_t, self.n_v = X_.shape
+        self.n_subjects, self.n_time_points, self.n_voxels = X_.shape
 
-        self.tuning_data = np.empty(self.n_s, dtype=np.float32)
-        self.denoised_signal = np.empty(self.n_s, dtype=np.float32)
+        self.tuning_data = np.empty(self.n_subjects, dtype=np.float32)
+        self.denoised_signal = np.empty(self.n_subjects, dtype=np.float32)
 
         if searchlights is None:
             self.regions = parcels
@@ -670,7 +671,7 @@ def fit(
 
         full_signal = np.concatenate(self.denoised_signal, axis=1)
         self.shared_response = self._stimulus_estimator(
-            full_signal, self.n_t, self.n_s, self.n_components
+            full_signal, self.n_subjects, self.n_components
         )
         if tuning:
             self.tuning_data = Parallel(n_jobs=self.n_jobs)(
@@ -678,7 +679,7 @@ def fit(
                     self.shared_response,
                     self.denoised_signal[i],
                 )
-                for i in range(self.n_s)
+                for i in range(self.n_subjects)
             )
 
         return self
@@ -690,13 +691,13 @@ def transform(self, X, verbose=False):
         Parameters:
         --------
         X : array-like
-            The test data of shape (n_subjects, n_samples, n_voxels).
+            The test data of shape (n_subjects, n_timepoints, n_voxels).
         verbose : bool(optional)
             Whether to print progress information. Defaults to False.
 
         Returns:
         --------
-        array-like: The transformed data of shape (n_subjects, n_samples, n_voxels).
+        array-like: The transformed data of shape (n_subjects, n_timepoints, n_voxels).
         """
 
         full_signal = np.concatenate(X, axis=1, dtype=np.float32)
@@ -705,7 +706,7 @@ def transform(self, X, verbose=False):
             print("Predict : Computing stimulus matrix...")
 
         stimulus_ = self._stimulus_estimator(
-            full_signal, self.n_t, self.n_s, self.n_components
+            full_signal, self.n_time_points, self.n_subjects, self.n_components
         )
 
         if verbose:
diff --git a/fmralign/fetch_example_data.py b/fmralign/fetch_example_data.py
index bb36e72..67c6dfd 100644
--- a/fmralign/fetch_example_data.py
+++ b/fmralign/fetch_example_data.py
@@ -43,9 +43,7 @@ def fetch_ibc_subjects_contrasts(subjects, data_dir=None, verbose=1):
     """
     # The URLs can be retrieved from the nilearn account on OSF
     if subjects == "all":
-        subjects = [
-            "sub-{i:02d}" for i in [1, 2, 4, 5, 6, 7, 8, 9, 11, 12, 13, 14, 15]
-        ]
+        subjects = ["sub-{i:02d}" for i in [1, 2, 4, 5, 6, 7, 8, 9, 11, 12, 13, 14, 15]]
     dataset_name = "ibc"
     data_dir = get_dataset_dir(dataset_name, data_dir=data_dir, verbose=verbose)
 
@@ -64,9 +62,7 @@ def fetch_ibc_subjects_contrasts(subjects, data_dir=None, verbose=1):
     )
     metadata_df = pd.read_csv(metadata_path[0])
     conditions = metadata_df.condition.unique()
-    metadata_df["path"] = metadata_df["path"].str.replace(
-        "path_to_dir", data_dir
-    )
+    metadata_df["path"] = metadata_df["path"].str.replace("path_to_dir", data_dir)
     # filter the dataframe to return only rows relevant for subjects argument
     metadata_df = metadata_df[metadata_df.subject.isin(subjects)]
 
diff --git a/fmralign/generate_data.py b/fmralign/generate_data.py
index 4a6c9a4..2cb5b89 100644
--- a/fmralign/generate_data.py
+++ b/fmralign/generate_data.py
@@ -5,9 +5,9 @@
 
 
 def generate_dummy_signal(
-    n_s: int,
-    n_t: int,
-    n_v: int,
+    n_subjects: int,
+    n_timepoints: int,
+    n_voxels: int,
     S_std=1,
     latent_dim=None,
     T_mean=0,
@@ -56,30 +56,30 @@ def generate_dummy_signal(
         Tuning matrices.
     """
     if latent_dim is None:
-        latent_dim = n_t
+        latent_dim = n_timepoints
 
     rng = np.random.RandomState(seed=seed)
 
     if generative_method == "custom":
-        sigma = n_s * np.arange(1, latent_dim + 1)
+        sigma = n_subjects * np.arange(1, latent_dim + 1)
         # np.random.shuffle(sigma)
         # Generate common signal matrix
-        S_train = S_std * np.random.randn(n_t, latent_dim)
+        S_train = S_std * np.random.randn(n_timepoints, latent_dim)
         # Normalize each row to have unit norm
         S_train = S_train / np.linalg.norm(S_train, axis=0, keepdims=True)
         S_train = S_train @ np.diag(sigma)
-        S_test = S_std * np.random.randn(n_t, latent_dim)
+        S_test = S_std * np.random.randn(n_timepoints, latent_dim)
         S_test = S_test / np.linalg.norm(S_test, axis=0, keepdims=True)
         S_test = S_test @ np.diag(sigma)
 
     elif generative_method == "fastsrm":
         Sigma = rng.dirichlet(np.ones(latent_dim), 1).flatten()
-        S_train = np.sqrt(Sigma)[:, None] * rng.randn(n_t, latent_dim)
-        S_test = np.sqrt(Sigma)[:, None] * rng.randn(n_t, latent_dim)
+        S_train = np.sqrt(Sigma)[:, None] * rng.randn(n_timepoints, latent_dim)
+        S_test = np.sqrt(Sigma)[:, None] * rng.randn(n_timepoints, latent_dim)
 
     elif generative_method == "multiviewica":
-        S_train = np.random.laplace(size=(n_t, latent_dim))
-        S_test = np.random.laplace(size=(n_t, latent_dim))
+        S_train = np.random.laplace(size=(n_timepoints, latent_dim))
+        S_test = np.random.laplace(size=(n_timepoints, latent_dim))
 
     else:
         raise ValueError("Unknown generative method")
@@ -87,30 +87,30 @@ def generate_dummy_signal(
     # Generate indiivdual spatial components
     data_train, data_test = [], []
     Ts = []
-    for _ in range(n_s):
+    for _ in range(n_subjects):
         if generative_method == "custom" or generative_method == "multiviewica":
-            W = T_mean + T_std * np.random.randn(latent_dim, n_v)
+            W = T_mean + T_std * np.random.randn(latent_dim, n_voxels)
         else:
-            W = projection(rng.randn(latent_dim, n_v))
+            W = projection(rng.randn(latent_dim, n_voxels))
 
         Ts.append(W)
         X_train = S_train @ W
-        N = np.random.randn(n_t, n_v)
-        N = (
-            N
+        noise = np.random.randn(n_timepoints, n_voxels)
+        noise = (
+            noise
             * np.linalg.norm(X_train)
-            / (SNR * np.linalg.norm(N, axis=0, keepdims=True))
+            / (SNR * np.linalg.norm(noise, axis=0, keepdims=True))
         )
-        X_train += N
+        X_train += noise
         data_train.append(X_train)
         X_test = S_test @ W
-        N = np.random.randn(n_t, n_v)
-        N = (
-            N
+        noise = np.random.randn(n_timepoints, n_voxels)
+        noise = (
+            noise
             * np.linalg.norm(X_test)
-            / (SNR * np.linalg.norm(N, axis=0, keepdims=True))
+            / (SNR * np.linalg.norm(noise, axis=0, keepdims=True))
         )
-        X_test += N
+        X_test += noise
         data_test.append(X_test)
 
     data_train = np.array(data_train)
diff --git a/fmralign/hyperalignment/correlation.py b/fmralign/hyperalignment/correlation.py
index cb73465..f74099e 100644
--- a/fmralign/hyperalignment/correlation.py
+++ b/fmralign/hyperalignment/correlation.py
@@ -2,83 +2,28 @@
 meant to be used as a test for the hyperalignment algorithm only."""
 
 import numpy as np
-from scipy.stats import spearmanr
-from sklearn.metrics import pairwise_distances
-from scipy.spatial.distance import cdist
 from sklearn.manifold import MDS
 from scipy.optimize import linear_sum_assignment
-
-
-#############################################################################
-# METRICS
-#############################################################################
-
-
-def compute_correlation(X, Y, metric: str = "correlation"):
-    """Compute correlation between X and Y.
-
-    Parameters
-    ----------
-    X : ndarray of shape (n_samples, n_features)
-        The data.
-    Y : ndarray of shape (n_samples, n_features)
-        The data.
-    metric : {'correlation', 'spearman', 'euclidean'}
-        The metric to use.
-
-    Returns
-    -------
-    corr : ndarray of shape (n_samples, n_samples)
-        The correlation matrix.
-    """
-    if metric == "correlation":
-        corr = 1 - pairwise_distances(X, Y, metric="correlation")
-    elif metric == "spearman":
-        corr = np.zeros((X.shape[0], Y.shape[0]))
-        for i in range(X.shape[0]):
-            for j in range(Y.shape[0]):
-                corr[i, j] = spearmanr(X[i], Y[j])[0]
-    elif metric == "euclidean":
-        corr = -pairwise_distances(X, Y, metric="euclidean")
-    else:
-        raise ValueError("Unknown metric")
-    return corr
-
-
-def compute_similarity(X, Y, metric: str = "euclidean"):
-    """
-    Compute the similarity matrix between two sets of vectors.
-
-    Parameters:
-    X (ndarray): First set of vectors.
-    Y (ndarray): Second set of vectors.
-    metric (str, optional): The distance metric to use. Defaults to "euclidean".
-
-    Returns:
-    ndarray: The similarity matrix between X and Y.
-    """
-    assert X.shape == Y.shape
-    n = X.shape[0]
-
-    def vector_sim(u, v):
-        if metric == "euclidean":
-            return 1 - (np.linalg.norm(u - v) / np.linalg.norm(u + v)) ** 2
-        else:
-            return (1 - cdist(u, v, metric)).mean()
-
-    sim = np.zeros((n, n))
-    for i in range(n):
-        for j in range(i, n):
-            d = vector_sim(X[i], Y[j])
-            sim[i, j] = d
-            sim[j, i] = d
-    return sim
+from itertools import combinations
 
 
 def compute_pearson_corr(X, Y, linear_assignment: bool = False):
     """Compute Pearson correlation between X and Y.
     X and Y are two lists of matrices of the same shape.
     The returned matrix will be of shape 2N x 2N, where N is the number of matrices in X and Y.
+
+    Parameters:
+    ----------
+    X : ndarray
+        First set of matrices.
+    Y : ndarray
+        Second set of matrices.
+    linear_assignment : (bool, optional)
+        Whether to perform linear assignment optimization. Defaults to False.
+
+    Returns:
+    -------
+    ndarray: Pearson correlation matrix.
     """
 
     assert X.shape == Y.shape
@@ -103,16 +48,23 @@ def pearson_corr_coeff(
     linear_assignment: bool = True,
 ):
     """
-    Compute Pearson correlation coefficient between matrices M1 and M2.
+    Compute Pearson correlation coefficient between matrices M1 and M2 by averaging diagonal elements of the correlation matrix.
+    This function can also perform linear assignment optimization to maximize the correlation coefficient.
 
     Parameters:
-        M1 (numpy.ndarray): First matrix.
-        M2 (numpy.ndarray): Second matrix.
-        absolute (bool, optional): Whether to compute absolute correlation coefficients. Defaults to True.
-        linear_assignment (bool, optional): Whether to perform linear assignment optimization. Defaults to True.
+    ----------
+    M1 : ndarray
+        First matrix.
+    M2 : ndarray)
+        Second matrix.
+    absolute : (bool, optional)
+        Whether to compute absolute correlation coefficients. Defaults to True.
+    linear_assignment : (bool, optional)
+        Whether to perform linear assignment optimization. Defaults to True.
 
     Returns:
-        float: Pearson correlation coefficient.
+    -------
+    float: Pearson correlation coefficient.
     """
     assert M1.shape == M2.shape
 
@@ -136,7 +88,20 @@ def pearson_corr_coeff(
 
 
 def tuning_correlation(X, Y):
-    """Compute pairwise Pearson correlation matrix between two sets of matrices."""
+    """Compute pairwise Pearson correlation matrix between two sets of matrices.
+    X and Y are two lists of matrices of the same shape.
+
+    Parameters:
+    ----------
+    X : ndarray
+        First set of matrices, with shape (n_subjects, n_samples, n_features).
+    Y : ndarray
+        Second set of matrices, with shape (n_subjects, n_samples, n_features).
+
+    Returns:
+    -------
+    ndarray: Pearson correlation matrix.
+    """
     assert X.shape == Y.shape
     n = X.shape[0]
     corr_mat = np.zeros((n, n))
@@ -150,7 +115,23 @@ def tuning_correlation(X, Y):
 
 
 def stimulus_correlation(X, Y, linear_assignment=True, absolute=True):
-    """Compute pairwise Pearson correlation matrix between two stimulus matrices."""
+    """Compute pairwise Pearson correlation matrix between two stimulus matrices.
+
+    Parameters:
+    ----------
+    X : ndarray
+        First stimulus matrix, with shape (n_samples, n_features).
+    Y : ndarray
+        Second stimulus matrix, with shape (n_samples, n_features).
+    linear_assignment : (bool, optional)
+        Whether to perform linear assignment optimization. Defaults to True.
+    absolute : (bool, optional)
+        Whether to compute absolute correlation coefficients. Defaults to True.
+
+    Returns:
+    -------
+    ndarray: Pearson correlation matrix.
+    """
     assert X.shape == Y.shape
     n = X.shape[0]
     corr_mat = np.corrcoef(X, Y)[:n, n:]
@@ -171,16 +152,23 @@ def matrix_MDS(X, Y, n_components=2, dissimilarity="euclidean"):
     Perform multidimensional scaling (MDS) on the given data matrices X and Y.
 
     Parameters:
-        X (list): The first data matrix.
-        Y (list): The second data matrix.
-        n_components (int): The number of dimensions in the output space (default is 2).
-        dissimilarity (str or array-like): The dissimilarity measure to use. If it is a string other than "precomputed",
-                                          the dissimilarity is computed using the Euclidean distance between flattened data points.
-                                          If it is "precomputed", the dissimilarity is assumed to be a precomputed dissimilarity matrix.
+    ----------
+
+    X : list of ndarray
+        The first data matrix.
+    Y : list of ndarray
+        The second data matrix.
+    n_components : int
+        The number of dimensions in the output space (default is 2).
+    dissimilarity : str or array-like
+        The dissimilarity measure to use. If it is a string other than "precomputed",
+        the dissimilarity is computed using the Euclidean distance between flattened data points.
+        If it is "precomputed", the dissimilarity is assumed to be a precomputed dissimilarity matrix.
 
     Returns:
-        tuple: A tuple containing two arrays. The first array represents the transformed data points from matrix X,
-               and the second array represents the transformed data points from matrix Y.
+    -------
+    tuple: A tuple containing two arrays. The first array represents the transformed data points from matrix X,
+            and the second array represents the transformed data points from matrix Y.
     """
     assert len(X) == len(Y)
 
@@ -203,20 +191,26 @@ def thread_compute_correlation(X, Y, i, j):
     Compute the correlation between two time series X_i and Y_i.
 
     Parameters:
-    - X (ndarray):
+    ----------
+    X : ndarray
         ndrray of shape (n_samples, n_features) representing the first time series.
-    - Y (ndarray):
+    Y : ndarray
         ndrray of shape (n_samples, n_features) representing the second time series.
-    - i (int):
+    i : int
         Index of the first time series.
-    - j (int):
+    j : int
         Index of the second time series.
 
     Returns:
-    diff_TR_corr (ndarray): Array of shape (n_samples * (n_samples - 1),) containing the correlations between different time points.
-    same_TR_corr (ndarray): Array of shape (n_samples,) containing the correlations between the same time points.
-    empty_diff_TR_corr (ndarray): Empty array.
-    empty_same_TR_corr (ndarray): Empty array.
+    -------
+    ndarray
+        Correlations between different time points if i != j, otherwise an empty array.
+    ndarray
+        Correlations between the same time points if i != j, otherwise an empty array.
+    ndarray
+        Correlations between different time points if i == j, otherwise an empty array.
+    ndarray
+        Correlations between the same time points if i == j, otherwise an empty array.
     """
     X_i, Y_i = X[i], Y[i]
     corr = stimulus_correlation(X_i, Y_i, absolute=False)
@@ -242,19 +236,29 @@ def multithread_compute_correlation(
     """
     Compute correlations between pairs of samples in X and Y using multiple threads.
 
-    Args:
-        X (ndarray): The first set of samples, with shape (n_samples, n_features).
-        Y (ndarray): The second set of samples, with shape (n_samples, n_features).
-        absolute (bool, optional): Whether to compute absolute correlations. Defaults to False.
-        linear_assignment (bool, optional): Whether to use linear assignment for correlation computation. Defaults to True.
-        n_jobs (int, optional): The number of threads to use for parallel computation. Defaults to 40.
+    Parameters:
+    ----------
+
+    X : ndarray)
+        The first set of samples, with shape (n_samples, n_features).
+    Y : ndarray
+        The second set of samples, with shape (n_samples, n_features).
+    absolute : bool(optional)
+        Whether to compute absolute correlations. Defaults to False.
+    linear_assignment : bool(optional)
+        Whether to use linear assignment for correlation computation. Defaults to True.
+    n_jobs : int(optional)
+        The number of threads to use for parallel computation. Defaults to 40.
 
     Returns:
-        tuple: A tuple containing four arrays:
-            - corr_same_sub_diff_TR: Correlations between different time points of the same subject.
-            - corr_same_sub_same_TR: Correlations between the same time points of the same subject.
-            - corr_diff_sub_diff_TR: Correlations between different time points of different subjects.
-            - corr_diff_sub_same_TR: Correlations between the same time points of different subjects.
+    -------
+
+    tuple
+        A tuple containing four arrays
+        - corr_same_sub_diff_TR: Correlations between different time points of the same subject.
+        - corr_same_sub_same_TR: Correlations between the same time points of the same subject.
+        - corr_diff_sub_diff_TR: Correlations between different time points of different subjects.
+        - corr_diff_sub_same_TR: Correlations between the same time points of different subjects.
     """
     from joblib import Parallel, delayed
 
@@ -282,8 +286,6 @@ def thread_compute_correlation(X, Y, i, j):
                 np.array([]),
             )
 
-    from itertools import combinations
-
     assert X.shape == Y.shape
     n_s = X.shape[0]
     coordinates = list(combinations(range(n_s), 2)) + [(i, i) for i in range(n_s)]
diff --git a/fmralign/hyperalignment/individualized_neural_tuning.py b/fmralign/hyperalignment/individualized_neural_tuning.py
index 6215db8..8acad48 100644
--- a/fmralign/hyperalignment/individualized_neural_tuning.py
+++ b/fmralign/hyperalignment/individualized_neural_tuning.py
@@ -30,21 +30,23 @@ def __init__(
 
         Parameters:
         -----------
-        - tmpl_kind (str):
+        tmpl_kind : str
             The type of template used for alignment. Default is "pca".
-        - decomp_method (str):
+        decomp_method : str
             The decomposition method used for template construction. Default is None.
-        - alignment_method (str) in ["searchlight", "parcellation"]:
+        alignment_method : str in ["searchlight", "parcellation"
             The alignment method used. Default is "searchlight".
-        - n_pieces (int): The number of pieces to divide the data into if using parcellation. Default is 150.
-        - radius (int):
+        n_pieces : int
+            The number of pieces to divide the data into if using parcellation. Default is 150.
+        radius : int
             The radius of the searchlight sphere in millimeters.
             Only used if alignment_method is "searchlight".
             Default is 20.
-        - n_components (int):
+        n_components : int
             The number of latent dimensions to use. If None, all the components are used.
             Default is None.
-        - n_jobs (int): The number of parallel jobs to run. Default is 1.
+        n_jobs : int
+            The number of parallel jobs to run. Default is 1.
         """
         super().__init__(
             template=template,
@@ -86,6 +88,11 @@ def fit(
             Not used.
         verbose : int
             The verbosity level.
+
+        Returns
+        -------
+        self : IndividualizedNeuralTuning
+            The fitted model.
         """
 
         if mask_img is None:
diff --git a/fmralign/hyperalignment/linalg.py b/fmralign/hyperalignment/linalg.py
index c0ae5ff..1448349 100644
--- a/fmralign/hyperalignment/linalg.py
+++ b/fmralign/hyperalignment/linalg.py
@@ -89,6 +89,7 @@ def svd_pca(X, remove_mean=True):
 
 def ridge(X, Y, alpha=10):
     """Solve ridge regression problem for matrix target using SVD.
+
     Parameters
     ----------
     X : ndarray
@@ -97,6 +98,7 @@ def ridge(X, Y, alpha=10):
         The target matrix.
     alpha : float
         The regularization parameter.
+
     Returns
     -------
     betas : ndarray of shape (n_features, n_targets)
@@ -109,33 +111,10 @@ def ridge(X, Y, alpha=10):
     return betas
 
 
-def _ridge(X, Y, alpha):
-    from sklearn.linear_model import Ridge
-
-    ridge = Ridge(alpha=alpha, fit_intercept=False)
-    ridge.fit(X, Y)
-    return ridge.coef_.T
-
-
 def procrustes(X, Y, reflection=True, scaling=False):
     r"""
     The orthogonal Procrustes algorithm.
 
-    The orthogonal Procrustes algorithm, also known as the classic
-    hyperalignment algorithm, is the first hyperalignment algorithm, which
-    was introduced in Haxby et al. (2011). It tries to align two
-    configurations in a high-dimensional space through an orthogonal
-    transformation. The transformation is an improper rotation, which is a
-    rotation with an optional reflection. Neither rotation nor reflection
-    changes the geometry of the configuration. Therefore, the geometry,
-    such as a representational dissimilarity matrix (RDM), remains the
-    same in the process.
-
-    Optionally, a global scaling can be added to the algorithm to further
-    improve alignment quality. That is, scaling the data with the same
-    factor for all directions. Different from rotation and reflection,
-    global scaling can change the geometry of the data.
-
     Parameters
     ----------
     X : ndarray
@@ -159,17 +138,6 @@ def procrustes(X, Y, reflection=True, scaling=False):
         Depending on the parameters ``reflection`` and ``scaling``, the
         transformation can be a pure rotation, an improper rotation, or a
         pure/improper rotation with global scaling.
-
-    Notes
-    -----
-    The algorithm tries to minimize the Frobenius norm of the difference
-    between transformed data ``X @ T`` and the target ``Y``:
-
-    .. math:: \underset{T}{\arg\min} \lVert XT - Y \rVert_F
-
-    The solution ``T`` differs depending on whether reflection and global
-    scaling are allowed. When it's a rotation (pure or improper), it's
-    often denoted as ``R``.
     """
 
     A = Y.T.dot(X).T
diff --git a/fmralign/hyperalignment/local_template.py b/fmralign/hyperalignment/local_template.py
index 34f7943..fd67f8a 100644
--- a/fmralign/hyperalignment/local_template.py
+++ b/fmralign/hyperalignment/local_template.py
@@ -1,9 +1,7 @@
 """ Local template computation functions. Those functions are part of the warp hyperalignment
-introducted by Feilong Ma et al. 2023. The functions are adapted from the original code and
-adapted for more general regionations approches.
+introducted by Feilong Ma et al. 2023.
 """
 
-
 import numpy as np
 from scipy.stats import zscore
 from sklearn.decomposition import PCA
@@ -20,7 +18,7 @@ def PCA_decomposition(
 
     Parameters
     ----------
-    X : ndarray of shape (ns, nt, nv)
+    X : ndarray of shape (n_subjects, n_timepoints, n_voxels)
         The input data array.
     n_components : int or None
         The number of components to keep. If None, all components are kept.
@@ -33,8 +31,8 @@ def PCA_decomposition(
 
     Returns
     -------
-    XX : ndarray of shape (nt, npc)
-    cc : ndarray of shape (npc, ns, nv)
+    XX : ndarray of shape (n_timepoints, n_components)
+    cc : ndarray of shape (n_components, n_subjects, n_voxels)
     """
     ns, nt, nv = X.shape
     X = X.transpose(1, 0, 2).reshape(nt, ns * nv).astype(np.float32)
@@ -88,21 +86,21 @@ def compute_PCA_template(X, sl=None, n_components=None, flavor="sklearn", demean
     Parameters:
     -----------
 
-    - X (ndarray):
-        The input data array of shape (n_samples, n_features, n_timepoints).
-    - sl (slice, optional):
-        The slice of timepoints to consider. Defaults to None.
-    - n_components (int, optional):
-        The maximum number of principal components to keep. Defaults to None.
-    - flavor (str, optional):
+    X : ndarray of shape (n_samples, n_features, n_timepoints)
+        The input data array.
+    sl : slice(optional)
+        The region indices for searchlight-based template computation. Defaults to None.
+    n_components : int(optional)
+        The maximum number of principal components to keep. If None, all components will be kept. Defaults to None.
+    flavor : str(optional)
         The flavor of PCA algorithm to use. Defaults to "sklearn".
-    - demean (bool, optional):
+    demean : bool(optional)
         Whether to demean the data before performing PCA. Defaults to False.
 
     Returns:
     --------
 
-    XX (ndarray):
+    XX : ndarray
         The PCA template array of shape (n_samples, n_features, n_components).
     """
     if sl is not None:
@@ -118,16 +116,39 @@ def compute_PCA_template(X, sl=None, n_components=None, flavor="sklearn", demean
 
 
 def compute_PCA_var1_template(
-    dss, sl=None, n_components=None, flavor="sklearn", demean=True
+    X, sl=None, n_components=None, flavor="sklearn", demean=True
 ):
+    """
+    Compute the PCA template from the input data.
+
+    Parameters:
+    -----------
+
+    X : ndarray of shape (n_samples, n_features, n_timepoints)
+        The input data array.
+    sl : slice(optional)
+        The region indices for searchlight-based template computation. Defaults to None.
+    n_components : int(optional)
+        The maximum number of principal components to keep. If None, all components will be kept. Defaults to None.
+    flavor : str(optional)
+        The flavor of PCA algorithm to use. Defaults to "sklearn".
+    demean : bool(optional)
+        Whether to demean the data before performing PCA. Defaults to False.
+
+    Returns:
+    --------
+
+    XX : ndarray
+        The PCA template array of shape (n_samples, n_features, n_components).
+    """
     if sl is not None:
-        dss = dss[:, :, sl]
+        X = X[:, :, sl]
     XX, cc = PCA_decomposition(
-        dss, n_components=n_components, flavor=flavor, adjust_ns=False, demean=demean
+        X, n_components=n_components, flavor=flavor, adjust_ns=False, demean=demean
     )
     w = np.sqrt(np.sum(cc**2, axis=2)).mean(axis=1)
     XX *= w[np.newaxis]
-    return XX
+    return XX.astype(np.float32)
 
 
 def compute_procrustes_template(
@@ -142,19 +163,27 @@ def compute_procrustes_template(
     """
     Compute the Procrustes template for a given set of data.
 
-    Args:
-        X (ndarray): The input data array of shape (n_samples, n_features, n_regions).
-        region (int or None, optional): The index of the region to consider. If None, all regions are considered. Defaults to None.
-        reflection (bool, optional): Whether to allow reflection in the Procrustes alignment. Defaults to True.
-        scaling (bool, optional): Whether to allow scaling in the Procrustes alignment. Defaults to False.
-        zscore_common (bool, optional): Whether to z-score the aligned data to have zero mean and unit variance. Defaults to True.
-        level2_iter (int, optional): The number of iterations for the level 2 alignment. Defaults to 1.
-        X2 (ndarray or None, optional): The second set of input data array of shape (n_samples, n_features, n_regions). Only used for level 2 alignment. Defaults to None.
-        debug (bool, optional): Whether to display progress bars during alignment. Defaults to False.
+    Parameters:
+    -----------
+    X : ndarray of shape (n_subjects, n_timepoints, n_voxels)
+        The input data array.
+    region : arraylike
+        The index of the region to consider. If None, all regions are considered. Defaults to None.
+    reflection : bool (optional)
+        Whether to allow reflection in the Procrustes alignment. Defaults to True.
+    scaling bool (optional)
+        Whether to allow scaling in the Procrustes alignment. Defaults to False.
+    zscore_common : bool(optional)
+        Whether to z-score the aligned data to have zero mean and unit variance. Defaults to True.
+    level2_iter : int(optional)
+        The number of iterations for the level 2 alignment. Defaults to 1.
+    X2 : ndarray(optional)
+        The second set of input data array of shape (n_samples, n_features, n_regions). Only used for level 2 alignment. Defaults to None.
 
     Returns:
-        ndarray: The computed Procrustes template.
-
+    --------
+    common_space: ndarray
+        The computed Procrustes template.
     """
     if region is not None:
         X = X[:, :, region]
@@ -216,21 +245,30 @@ def compute_template(
     """
     Compute a template from a set of datasets.
 
-    ----------
     Parameters:
+    -----------
 
-       - X (ndarray): The input datasets.
-       - region (ndarray or None): The region indices for searchlight or region-based template computation.
-       - sl (ndarray or None): The searchlight indices for searchlight-based template computation.
-       - region (int or None, optional): The index of the region to consider. If None, all regions are considered (or searchlights). Defaults to None.
-       - kind (str): The type of template computation algorithm to use. Can be "pca", "pcav1", "pcav2", or "cls".
-       - n_components (int or None): The maximum number of principal components to use for PCA-based template computation.
-       - common_topography (bool): Whether to enforce common topography across datasets.
-       - demean (bool): Whether to demean the datasets before template computation.
+    X : ndarray of shape (n_subjects, n_timepoints, n_voxels)
+        The input datasets.
+    region : ndarray or None
+        The region indices for searchlight or region-based template computation.
+    sl : ndarray or None
+        The searchlight indices for searchlight-based template computation.
+    region : int
+        The index of the region to consider.
+    kind : str
+        The type of template computation algorithm to use. Can be "pca", "pcav1", "pcav2", or "cls".
+    n_components : int(optional)
+        The maximum number of principal components to use for PCA-based template computation. Defaults to 150.
+    common_topography : bool(optional)
+        Whether to enforce common topography across datasets. Defaults to True.
+    demean : bool(optional)
+        Whether to demean the datasets before template computation. Defaults to True.
 
-    ----------
     Returns:
-        tmpl : The computed template on all parcels (or searchlights).
+    --------
+    tmpl : ndaray of shape (n_timepoints, n_voxels)
+        The computed template on all parcels (or searchlights).
     """
     mapping = {
         "pca": compute_PCA_template,
diff --git a/fmralign/hyperalignment/piecewise_alignment.py b/fmralign/hyperalignment/piecewise_alignment.py
index a51dd7f..462b6ce 100644
--- a/fmralign/hyperalignment/piecewise_alignment.py
+++ b/fmralign/hyperalignment/piecewise_alignment.py
@@ -148,7 +148,6 @@ def fit_transform(
             regions=regions,
             n_jobs=self.n_jobs,
             template_kind=self.template_kind,
-            verbose=self.verbose,
             common_topography=self.common_topography,
             weights=self.weights,
         )
diff --git a/fmralign/hyperalignment/regions.py b/fmralign/hyperalignment/regions.py
index 2bec23c..551100f 100644
--- a/fmralign/hyperalignment/regions.py
+++ b/fmralign/hyperalignment/regions.py
@@ -32,14 +32,16 @@
 
 def create_parcels_from_labels(labels: np.ndarray):
     """
-    Create parcels from labels.
 
-    Args:
-        labels (np.ndarray): Array of labels.
+    Parameters:
+    ----------
+    labels : ndarray
+        Array of labels.
 
     Returns:
-        list: List of parcels, where each parcel is an array of indices.
-
+    -------
+    parcles : list
+        List of parcels, where each parcel is an array of indices.
     """
     n_labels = labels.max()
     parcels = []
@@ -355,7 +357,6 @@ def iter_hyperalignment(
     searchlights,
     sl_func,
     return_betas=False,
-    verbose=False,
 ):
     """
     Tool function to iterate hyperalignment over pieces of data.
@@ -375,8 +376,6 @@ def iter_hyperalignment(
     return_betas : bool, optional
         Whether to return the coefficients of regression instead of the prediciton.
         Defaults to False.
-    verbose : bool, optional
-        Whether to display progress. Defaults to False.
 
     Returns
     -------
@@ -411,21 +410,26 @@ def piece_procrustes(
     scaling=False,
 ):
     """
-    Applies searchlight hyperalignment using Procrustes alignment.
-
-    Args:
-        X (array-like): The source data matrix of shape (n_samples, n_features).
-        Y (array-like): The target data matrix of shape (n_samples, n_features).
-        searchlights (array-like): The indices of the searchlight regions.
-        dists (array-like): distances of vertices to the center of their searchlight, of shape (n_searchlights, n_vertices_sl)
-        radius (float): The radius of the searchlight region.
-        T0 (array-like, optional): The initial transformation matrix. Defaults to None.
-        reflection (bool, optional): Whether to allow reflection in the alignment. Defaults to True.
-        scaling (bool, optional): Whether to allow scaling in the alignment. Defaults to False.
-        weighted (bool, optional): Whether to use weighted Procrustes alignment. Defaults to True.
+    Computes a transformation matrix from a template and a target signal using Procrustes hyperalignment.
+
+    Parameters:
+    ----------
+    X : ndarray
+        The source data matrix of shape (n_samples, n_features).
+    Y : ndarray
+        The target data matrix of shape (n_samples, n_features).
+    regions : list of arrays
+        List of brain regions. Contains the indices of the voxels in each region (either parcels or searchlights).
+    T0 : array-like, optional
+        The initial transformation matrix. Defaults to None.
+    reflection : bool, optional
+        Whether to allow reflection. Defaults to True.
 
     Returns:
-        array-like: The transformation matrix T.
+    -------
+    T : array-like
+        The transformation matrix T.
+
 
     """
     sl_func = functools.partial(procrustes, reflection=reflection, scaling=scaling)
@@ -444,25 +448,29 @@ def piece_ridge(
     Y,
     regions,
     alpha=1e3,
-    weights=None,
     verbose=False,
     return_betas=False,
 ):
     """
     Perform searchlight ridge regression for hyperalignment.
 
-    Args:
-        X (array-like): The source data matrix of shape (n_samples, n_features).
-        Y (array-like): The target data matrix of shape (n_samples, n_features).
-        searchlights (array-like): The indices of the searchlight regions.
-        dists (array-like): distances of vertices to the center of their searchlight, of shape (n_searchlights, n_vertices_sl)
-        radius (float): The radius of the searchlight region.
-        T0 (array-like, optional): The initial transformation matrix. Defaults to None.
-        alpha (float, optional): The regularization parameter for ridge regression. Defaults to 1e3.
-        weighted (bool, optional): Whether to use weighted ridge regression. Defaults to True.
+    Parameters:
+    ----------
+    X : ndarray
+        The source data matrix of shape (n_samples, n_features).
+    Y : ndarray
+        The target data matrix of shape (n_samples, n_features).
+    regions : list of arrays
+        List of brain regions. Contains the indices of the voxels in each region (either parcels or searchlights).
+    alpha : float(optional)
+        The regularization parameter for Ridge regression. Defaults to 1e3.
+    return_betas : bool(optional)
+        Whether to return the coefficients of regression instead of the prediciton. Defaults to False.
 
     Returns:
-        array-like: The transformation matrix T.
+    -------
+    T : array-like
+        The transformation matrix T.
 
     """
     sl_func = functools.partial(ridge, alpha=alpha)
@@ -483,22 +491,29 @@ def template(
     regions,
     n_jobs=1,
     template_kind="pca",
-    verbose=False,
     common_topography=True,
     weights=None,
 ):
     """
     Compute a template by aggregating local templates within searchlights.
 
-    Args:
-        X (numpy.ndarray): The input data matrix of shape (n_subjects, n_samples, n_features).
-        regions (list): The indices of the searchlight/parcels regions.
-        n_jobs (int, optional): The number of parallel jobs to run. Defaults to -1.
-        template_kind (str, optional): The kind of template to compute. Defaults to "pca".
-        verbose (bool, optional): Whether to display progress. Defaults to False.
+    Parameters:
+    ----------
+
+    X : ndarray
+        The input data matrix of shape (n_subjects, n_samples, n_features).
+    regions : list of ndarrays
+        List of regions composed of indices of voxels.
+    n_jobs : int(optional)
+        The number of parallel jobs to run. Defaults to 1.
+    template_kind : str(optional)
+        The kind of template to compute. Defaults to "pca".
+
 
     Returns:
-        numpy.ndarray: The computed template of shape (n_features,).
+    -------
+    template : ndarray of shape (n_timepoints, n_voxels)
+        The computed template.
 
     """
     with Parallel(n_jobs=n_jobs, batch_size=1, verbose=1) as parallel:
diff --git a/fmralign/hyperalignment/test_hyperalignment.py b/fmralign/hyperalignment/test_hyperalignment.py
index aee883b..7e9f9bf 100644
--- a/fmralign/hyperalignment/test_hyperalignment.py
+++ b/fmralign/hyperalignment/test_hyperalignment.py
@@ -7,9 +7,9 @@ def test_int_fit_predict():
     """Test if the outputs and arguments of the INT are the correct format"""
     # Create random data
     X_train, X_test, S_true_first_part, S_true_second_part, Ts = generate_dummy_signal(
-        n_s=5,
-        n_t=50,
-        n_v=300,
+        n_subjects=5,
+        n_timepoints=50,
+        n_voxels=300,
         S_std=1,
         T_std=1,
         latent_dim=6,
@@ -61,17 +61,17 @@ def test_int_fit_predict():
     assert 3 * np.mean(corr2_out) < np.mean(np.diag(corr2))
     assert 3 * np.mean(corr3_out) < np.mean(np.diag(corr3))
     assert 3 * np.mean(corr4_out) < np.mean(np.diag(corr4))
-    assert int1.tuning_data[0].shape == (6, int1.n_v)
-    assert int2.tuning_data[0].shape == (6, int2.n_v)
-    assert int1.shared_response.shape == (int1.n_t, 6)
+    assert int1.tuning_data[0].shape == (6, int1.n_voxels)
+    assert int2.tuning_data[0].shape == (6, int2.n_voxels)
+    assert int1.shared_response.shape == (int1.n_time_points, 6)
     assert X_pred.shape == X_test.shape
 
 
 def test_int_with_searchlight():
     X_train, X_test, stimulus_train, stimulus_test, _ = generate_dummy_signal(
-        n_s=5,
-        n_t=50,
-        n_v=300,
+        n_subjects=5,
+        n_timepoints=50,
+        n_voxels=300,
         S_std=1,
         T_std=1,
         latent_dim=6,
@@ -118,7 +118,7 @@ def test_int_with_searchlight():
     assert 3 * np.mean(corr2_out) < np.mean(np.diag(corr2))
     assert 3 * np.mean(corr3_out) < np.mean(np.diag(corr3))
     assert 3 * np.mean(corr4_out) < np.mean(np.diag(corr4))
-    assert model1.tuning_data[0].shape == (6, model1.n_v)
-    assert model2.tuning_data[0].shape == (6, model2.n_v)
-    assert model1.shared_response.shape == (model1.n_t, 6)
+    assert model1.tuning_data[0].shape == (6, model1.n_voxels)
+    assert model2.tuning_data[0].shape == (6, model2.n_voxels)
+    assert model1.shared_response.shape == (model1.n_time_points, 6)
     assert X_pred.shape == X_test.shape
diff --git a/fmralign/hyperalignment/toy_experiment.py b/fmralign/hyperalignment/toy_experiment.py
index 81899d6..d4ad5f6 100644
--- a/fmralign/hyperalignment/toy_experiment.py
+++ b/fmralign/hyperalignment/toy_experiment.py
@@ -1,6 +1,3 @@
-import os
-import sys
-
 import numpy as np
 import matplotlib.pyplot as plt
 from fmralign.alignment_methods import IndividualizedNeuralTuning as INT
@@ -18,10 +15,6 @@
 #############################################################################
 
 
-file_dir = os.path.dirname(__file__)
-sys.path.append(file_dir)
-
-
 n_s = 10
 n_t = 200
 n_v = 500
@@ -42,9 +35,9 @@
     stimulus_run_2,
     data_tuning,
 ) = generate_dummy_signal(
-    n_s=n_s,
-    n_t=n_t,
-    n_v=n_v,
+    n_subjects=n_s,
+    n_timepoints=n_t,
+    n_voxels=n_v,
     S_std=S_std,
     T_std=T_std,
     latent_dim=latent_dim,
diff --git a/fmralign/tests/test_alignment_methods.py b/fmralign/tests/test_alignment_methods.py
index 1e936b9..d946c9a 100644
--- a/fmralign/tests/test_alignment_methods.py
+++ b/fmralign/tests/test_alignment_methods.py
@@ -5,12 +5,6 @@
 from scipy.sparse import csc_matrix
 from scipy.linalg import orthogonal_procrustes
 
-# add directory to path
-import sys
-import os
-
-sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-
 from fmralign.alignment_methods import (
     DiagonalAlignment,
     Hungarian,
@@ -232,7 +226,7 @@ def test_searchlight_alignment_with_ridge():
     )
 
     X_train, X_test, _, _, _ = generate_dummy_signal(
-        n_t=n_time_points, n_v=n_voxels, n_s=n_subjects
+        n_timepoints=n_time_points, n_voxels=n_voxels, n_subjects=n_subjects
     )
 
     model = INT(n_jobs=5)
@@ -251,7 +245,7 @@ def test_parcel_alignment():
     parcels = create_parcels_from_labels(labels)
 
     X_train, X_test, _, _, _ = generate_dummy_signal(
-        n_t=n_time_points, n_v=n_voxels, n_s=n_subjects
+        n_timepoints=n_time_points, n_voxels=n_voxels, n_subjects=n_subjects
     )
 
     model = INT(n_jobs=5, alignment_method="parcel")

From 585365672c953f0f40958daed987effc59d5c1ff Mon Sep 17 00:00:00 2001
From: FOUCHARD Denis <denis.fouchard@inria.fr>
Date: Mon, 19 Feb 2024 15:13:24 +0100
Subject: [PATCH 45/69] fix more doc

---
 examples/plot_alignment_methods_benchmark.py |  8 +---
 fmralign/hyperalignment/regions.py           | 39 ++++++++++++--------
 2 files changed, 25 insertions(+), 22 deletions(-)

diff --git a/examples/plot_alignment_methods_benchmark.py b/examples/plot_alignment_methods_benchmark.py
index 3f4eb77..e343639 100644
--- a/examples/plot_alignment_methods_benchmark.py
+++ b/examples/plot_alignment_methods_benchmark.py
@@ -129,13 +129,7 @@
 from fmralign.metrics import score_voxelwise
 from fmralign.pairwise_alignment import PairwiseAlignment
 
-methods = [
-    "identity",
-    "scaled_orthogonal",
-    "ridge_cv",
-    "optimal_transport",
-    "individualized_neural_tuning",
-]
+methods = ["identity", "scaled_orthogonal", "ridge_cv", "optimal_transport"]
 
 for method in methods:
     alignment_estimator = PairwiseAlignment(
diff --git a/fmralign/hyperalignment/regions.py b/fmralign/hyperalignment/regions.py
index 551100f..f99d284 100644
--- a/fmralign/hyperalignment/regions.py
+++ b/fmralign/hyperalignment/regions.py
@@ -236,7 +236,8 @@ def compute_searchlights(
     return_dist_mat=False,
     n_jobs=1,
 ):
-    """Implement search_light analysis using an arbitrary type of classifier.
+    """
+    Compute searchlights for a given 4D image and mask.
 
     Parameters
     ----------
@@ -260,11 +261,6 @@ def compute_searchlights(
         Whether to return the distance matrix between voxels in the mask.
         Defaults to False.
 
-    groups : array-like of shape (n_samples,), optional
-        Labels of samples for each subject. If provided, the searchlights
-        will be computed within each group separately.
-        Defaults to None.
-
     verbose : int, optional
         Verbosity level (0 means no message).
         Defaults to 0.
@@ -279,10 +275,6 @@ def compute_searchlights(
         Contains the boolean indices for each sphere.
         shape: (number of seeds, number of voxels)
 
-    dist_matrix : 2D numpy.ndarray
-        Distance matrix between voxels in the mask.
-        shape: (number of voxels, number of voxels)
-
     dists : list of lists
         Contains the distance between each voxel and the seed.
         shape: (number of seeds, number of voxels)
@@ -332,6 +324,24 @@ def compute_searchlights(
 
 
 def searchlight_weights(searchlights, dists, radius):
+    """
+    Calculate the weights for each searchlight based on the distances from the center.
+
+    Parameters:
+    ----------
+    searchlights :list of arrays
+        List of searchlights, where each searchlight is represented as an array of voxel indices.
+    dists : array
+        Array of distances from the center for each searchlight.
+    radius : float
+        Radius of the searchlight.
+
+    Returns:
+    --------
+    weights : list
+        List of weights for each searchlight.
+
+    """
     nv = np.concatenate(searchlights).max() + 1
     weights_sum = np.zeros((nv,))
     for sl, d in zip(searchlights, dists):
@@ -354,7 +364,7 @@ def searchlight_weights(searchlights, dists, radius):
 def iter_hyperalignment(
     X,
     Y,
-    searchlights,
+    regions,
     sl_func,
     return_betas=False,
 ):
@@ -367,8 +377,8 @@ def iter_hyperalignment(
         The source data matrix.
     Y : array-like of shape (n_samples, n_features)
         The target data matrix.
-    searchlights : array-like
-        The indices of the searchlight regions.
+    regions : array-like
+        The indices of the regions.
     sl_func : function
         The function to use for hyperalignment.
     weights : array-like, optional
@@ -388,7 +398,7 @@ def iter_hyperalignment(
     else:
         Yhat = np.zeros_like(X, dtype=np.float32)
 
-    searchlights_iter = searchlights
+    searchlights_iter = regions
     for sl in searchlights_iter:
         x, y = X[:, sl], Y[:, sl]
         t = sl_func(x, y)
@@ -480,7 +490,6 @@ def piece_ridge(
         Y,
         regions,
         sl_func=sl_func,
-        verbose=verbose,
         return_betas=return_betas,
     )
     return T

From c1fc2002469ae19109f78d3a20d100178dc315fd Mon Sep 17 00:00:00 2001
From: FOUCHARD Denis <denis.fouchard@inria.fr>
Date: Mon, 19 Feb 2024 16:10:30 +0100
Subject: [PATCH 46/69] fix stimulus + fix doc

---
 examples/plot_int_alignment.py                | 14 ++++--
 fmralign/alignment_methods.py                 | 11 ++--
 fmralign/generate_data.py                     |  6 +--
 .../hyperalignment/test_hyperalignment.py     |  6 +--
 fmralign/tests/test_alignment_methods.py      | 50 -------------------
 5 files changed, 21 insertions(+), 66 deletions(-)

diff --git a/examples/plot_int_alignment.py b/examples/plot_int_alignment.py
index ca93312..d8169e5 100644
--- a/examples/plot_int_alignment.py
+++ b/examples/plot_int_alignment.py
@@ -73,13 +73,13 @@
 template_train = []
 for i in range(6):
     template_train.append(concat_imgs(imgs[i]))
-target_train = df[df.subject == "sub-07"][df.acquisition == "ap"].path.values
+
 
 # For subject sub-07, we split it in two folds:
 #   - target train: sub-07 AP contrasts, used to learn alignment to template
 #   - target test: sub-07 PA contrasts, used as a ground truth to score predictions
 # We make a single 4D Niimg from our list of 3D filenames
-
+target_train = df[df.subject == "sub-07"][df.acquisition == "ap"].path.values
 target_train = concat_imgs(target_train)
 target_train_data = masker.transform(target_train)
 target_test = df[df.subject == "sub-07"][df.acquisition == "pa"].path.values
@@ -138,8 +138,12 @@
     n_jobs=8, alignment_method="parcelation", n_components=20
 )
 model.fit(train_data, parcels=parcels, verbose=False)
+
 train_stimulus = np.copy(model.shared_response)
-train_tuning = np.linalg.pinv(train_stimulus) @ model.denoised_signal[-1]
+
+train_tuning = model._tuning_estimator(
+    shared_response=train_stimulus, target=model.denoised_signal[-1]
+)
 model_bis = IndividualizedNeuralTuning(
     n_jobs=8, alignment_method="parcelation", n_components=20
 )
@@ -150,7 +154,9 @@
 # than one subject for which we'd want to predict : [train_1, train_2 ...]
 
 stimulus_ = np.copy(model_bis.shared_response)
-prediction_from_template = stimulus_[test_index] @ train_tuning
+prediction_from_template = model._reconstruct_signal(
+    shared_response=stimulus_[test_index], individual_tuning=train_tuning
+)
 prediction_from_template = masker.inverse_transform(prediction_from_template)
 
 
diff --git a/fmralign/alignment_methods.py b/fmralign/alignment_methods.py
index a24b63d..0bf3299 100644
--- a/fmralign/alignment_methods.py
+++ b/fmralign/alignment_methods.py
@@ -588,13 +588,13 @@ def _reconstruct_signal(shared_response, individual_tuning):
         Parameters:
         --------
         shared_response : numpy.ndarray
-             The shared response of shape (n_t, n_t) or (n_t, latent_dim).
+             The shared response of shape (n_timeframes, n_timeframes) or (n_timeframes, latent_dim).
         individual_tuning : numpy.ndarray
-             The individual tuning of shape (latent_dim, n_v) or (n_t, n_v).
+             The individual tuning of shape (latent_dim, n_voxels) or (n_timeframes, n_voxels).
 
         Returns:
         --------
-        numpy.ndarray: The reconstructed signal of shape (n_t, n_v) (same shape as the original signal)
+        ndarray: The reconstructed signal of shape (n_timeframes, n_voxels) (same shape as the original signal)
         """
         return (shared_response @ individual_tuning).astype(np.float32)
 
@@ -628,8 +628,6 @@ def fit(
             Whether to compute the tuning weights. Defaults to True.
         verbose : bool(optional)
             Whether to print progress information. Defaults to True.
-        id : str(optional)
-            An identifier for caching purposes. Defaults to None.
 
         Returns:
         --------
@@ -706,8 +704,9 @@ def transform(self, X, verbose=False):
             print("Predict : Computing stimulus matrix...")
 
         stimulus_ = self._stimulus_estimator(
-            full_signal, self.n_time_points, self.n_subjects, self.n_components
+            full_signal, self.n_subjects, self.n_components
         )
+        print("Predict : stimulus matrix shape: ", stimulus_.shape)
 
         if verbose:
             print("Predict : stimulus matrix shape: ", stimulus_.shape)
diff --git a/fmralign/generate_data.py b/fmralign/generate_data.py
index 2cb5b89..0c0cb46 100644
--- a/fmralign/generate_data.py
+++ b/fmralign/generate_data.py
@@ -120,7 +120,7 @@ def generate_dummy_signal(
 
 def generate_dummy_searchlights(
     n_searchlights: int,
-    n_v: int,
+    n_voxels: int,
     radius: int,
     sl_size: int = 5,
     seed: int = 0,
@@ -131,7 +131,7 @@ def generate_dummy_searchlights(
     ----------
     n_searchlights : int
         Number of searchlights.
-    n_v : int
+    n_voxels : int
         Number of voxels.
     radius : int
         Radius of searchlights.
@@ -148,6 +148,6 @@ def generate_dummy_searchlights(
         Distances.
     """
     rng = np.random.RandomState(seed=seed)
-    searchlights = rng.randint(n_v, size=(n_searchlights, sl_size))
+    searchlights = rng.randint(n_voxels, size=(n_searchlights, sl_size))
     dists = rng.randint(radius, size=searchlights.shape)
     return searchlights, dists
diff --git a/fmralign/hyperalignment/test_hyperalignment.py b/fmralign/hyperalignment/test_hyperalignment.py
index 7e9f9bf..1236d13 100644
--- a/fmralign/hyperalignment/test_hyperalignment.py
+++ b/fmralign/hyperalignment/test_hyperalignment.py
@@ -6,8 +6,8 @@
 def test_int_fit_predict():
     """Test if the outputs and arguments of the INT are the correct format"""
     # Create random data
-    X_train, X_test, S_true_first_part, S_true_second_part, Ts = generate_dummy_signal(
-        n_subjects=5,
+    X_train, X_test, S_true_first_part, S_true_second_part, _ = generate_dummy_signal(
+        n_subjects=7,
         n_timepoints=50,
         n_voxels=300,
         S_std=1,
@@ -80,7 +80,7 @@ def test_int_with_searchlight():
         seed=0,
     )
     searchlights, dists = generate_dummy_searchlights(
-        n_searchlights=10, n_v=300, radius=5, seed=0
+        n_searchlights=10, n_voxels=30, radius=5, seed=0
     )
     from fmralign.hyperalignment.correlation import (
         tuning_correlation,
diff --git a/fmralign/tests/test_alignment_methods.py b/fmralign/tests/test_alignment_methods.py
index d946c9a..7dbddaa 100644
--- a/fmralign/tests/test_alignment_methods.py
+++ b/fmralign/tests/test_alignment_methods.py
@@ -13,18 +13,11 @@
     POTAlignment,
     RidgeAlignment,
     ScaledOrthogonalAlignment,
-    IndividualizedNeuralTuning as INT,
     _voxelwise_signal_projection,
     optimal_permutation,
     scaled_procrustes,
 )
 from fmralign.tests.utils import zero_mean_coefficient_determination
-from fmralign.generate_data import (
-    generate_dummy_signal,
-    generate_dummy_searchlights,
-)
-
-from fmralign.hyperalignment.regions import create_parcels_from_labels
 
 
 def test_scaled_procrustes_algorithmic():
@@ -209,46 +202,3 @@ def test_ott_backend():
     algo.fit(X, Y)
     old_implem.fit(X, Y)
     assert_array_almost_equal(algo.R, old_implem.R, decimal=3)
-
-
-def test_searchlight_alignment_with_ridge():
-    n_voxels = 99
-    n_time_points = 149
-    n_searchlights = 92
-    n_subjects = 5
-
-    radius = 20
-    searchlights, dists = generate_dummy_searchlights(
-        n_searchlights=n_searchlights,
-        n_v=n_voxels,
-        radius=radius,
-        seed=0,
-    )
-
-    X_train, X_test, _, _, _ = generate_dummy_signal(
-        n_timepoints=n_time_points, n_voxels=n_voxels, n_subjects=n_subjects
-    )
-
-    model = INT(n_jobs=5)
-    model.fit(X_train, searchlights=searchlights, dists=dists, radius=radius)
-    X_pred = model.transform(X_test)
-    assert X_pred.shape == X_test.shape
-
-
-def test_parcel_alignment():
-    n_voxels = 99
-    n_time_points = 149
-    n_subjects = 5
-
-    n_parcels = 10
-    labels = np.arange(n_voxels) % n_parcels + 1
-    parcels = create_parcels_from_labels(labels)
-
-    X_train, X_test, _, _, _ = generate_dummy_signal(
-        n_timepoints=n_time_points, n_voxels=n_voxels, n_subjects=n_subjects
-    )
-
-    model = INT(n_jobs=5, alignment_method="parcel")
-    model.fit(X_train, parcels=parcels)
-    X_pred = model.transform(X_test)
-    assert X_pred.shape == X_test.shape

From ffd13acf50465f86b6db44ec3e9bc8a53d6a72e8 Mon Sep 17 00:00:00 2001
From: FOUCHARD Denis <denis.fouchard@inria.fr>
Date: Mon, 19 Feb 2024 16:57:59 +0100
Subject: [PATCH 47/69] adressing plot code issues

---
 examples/plot_int_alignment.py            | 29 ++++++++++-------------
 fmralign/alignment_methods.py             |  2 +-
 fmralign/hyperalignment/toy_experiment.py |  4 +---
 3 files changed, 14 insertions(+), 21 deletions(-)

diff --git a/examples/plot_int_alignment.py b/examples/plot_int_alignment.py
index d8169e5..5b51dbc 100644
--- a/examples/plot_int_alignment.py
+++ b/examples/plot_int_alignment.py
@@ -112,6 +112,7 @@
 
 from nilearn.image import index_img
 from fmralign.alignment_methods import IndividualizedNeuralTuning
+from fmralign.hyperalignment.piecewise_alignment import PiecewiseAlignment
 from fmralign.hyperalignment.regions import compute_parcels
 
 ###############################################################################
@@ -126,36 +127,30 @@
 
 train_index = range(53)
 test_index = range(53, 106)
-full_index = range(106)
 
-train_data = np.array(masked_imgs)[:, train_index, :]
-test_data = np.array(masked_imgs)[:, full_index, :]
+denoising_data = np.array(masked_imgs)[:, train_index, :]
+training_data = np.array(masked_imgs)[:-1]
 target_test_masked = np.array(masked_imgs)[:, test_index, :]
 
 
 parcels = compute_parcels(niimg=template_train[0], mask=masker, n_parcels=100, n_jobs=5)
+denoiser = PiecewiseAlignment(alignment_method="parcelation", n_jobs=5)
+denoised_signal = denoiser.fit_transform(X=denoising_data, regions=parcels)
 model = IndividualizedNeuralTuning(
-    n_jobs=8, alignment_method="parcelation", n_components=20
+    n_jobs=8, alignment_method="parcelation", n_components=None
 )
-model.fit(train_data, parcels=parcels, verbose=False)
-
-train_stimulus = np.copy(model.shared_response)
-
-train_tuning = model._tuning_estimator(
-    shared_response=train_stimulus, target=model.denoised_signal[-1]
+model.fit(training_data, parcels=parcels, verbose=False)
+stimulus_ = np.copy(model.shared_response)
+target_tuning = model._tuning_estimator(
+    shared_response=stimulus_[train_index], target=denoised_signal[-1]
 )
-model_bis = IndividualizedNeuralTuning(
-    n_jobs=8, alignment_method="parcelation", n_components=20
-)
-model_bis.fit(test_data, parcels=parcels, verbose=False)
-test_stimulus = np.copy(model_bis.shared_response)
 # %%
 # We input the mapping image target_train in a list, we could have input more
 # than one subject for which we'd want to predict : [train_1, train_2 ...]
 
-stimulus_ = np.copy(model_bis.shared_response)
+
 prediction_from_template = model._reconstruct_signal(
-    shared_response=stimulus_[test_index], individual_tuning=train_tuning
+    shared_response=stimulus_[test_index], individual_tuning=target_tuning
 )
 prediction_from_template = masker.inverse_transform(prediction_from_template)
 
diff --git a/fmralign/alignment_methods.py b/fmralign/alignment_methods.py
index 0bf3299..b813d61 100644
--- a/fmralign/alignment_methods.py
+++ b/fmralign/alignment_methods.py
@@ -547,7 +547,7 @@ def _tuning_estimator(shared_response, target):
         array-like: The estimated tuning matrix for the given target.
 
         """
-        if shared_response.shape[1] < shared_response.shape[0]:
+        if shared_response.shape[1] != shared_response.shape[0]:
             return (np.linalg.pinv(shared_response)).dot(target)
         return np.linalg.inv(shared_response).dot(target)
 
diff --git a/fmralign/hyperalignment/toy_experiment.py b/fmralign/hyperalignment/toy_experiment.py
index d4ad5f6..fe3630c 100644
--- a/fmralign/hyperalignment/toy_experiment.py
+++ b/fmralign/hyperalignment/toy_experiment.py
@@ -22,7 +22,7 @@
 T_std = 1
 SNR = 100
 latent_dim = 15  # if None, latent_dim = n_t
-decomposition_method = None  # if None, SVD is used
+decomposition_method = "pca"  # if None, SVD is used
 
 
 #############################################################################
@@ -74,12 +74,10 @@
     int1 = INT(
         n_components=latent_dim,
         decomp_method=decomposition_method,
-        alignment_method="parcelation",
     )
     int2 = INT(
         n_components=latent_dim,
         decomp_method=decomposition_method,
-        alignment_method="parcelation",
     )
     int_first_part = int1.fit(
         data_run_1, parcels=parcels, verbose=False

From eb4a8dfc77dfe67280487b56fd371a195332b9ed Mon Sep 17 00:00:00 2001
From: FOUCHARD Denis <denis.fouchard@inria.fr>
Date: Mon, 19 Feb 2024 17:17:38 +0100
Subject: [PATCH 48/69] better variable names

---
 fmralign/generate_data.py                 | 18 ++++++++--------
 fmralign/hyperalignment/toy_experiment.py | 26 ++++++++++++++---------
 2 files changed, 25 insertions(+), 19 deletions(-)

diff --git a/fmralign/generate_data.py b/fmralign/generate_data.py
index 0c0cb46..e13aae3 100644
--- a/fmralign/generate_data.py
+++ b/fmralign/generate_data.py
@@ -20,16 +20,16 @@ def generate_dummy_signal(
 
     Parameters
     ----------
-    n_s : int
+    n_subjects : int
         Number of subjects.
-    n_t : int
+    n_timepoints : int
         Number of timepoints.
-    n_v : int
+    n_voxels : int
         Number of voxels.
     S_std : float, default=1
         Standard deviation of latent variables.
     latent_dim: int, defult=None
-        Number of latent dimensions. Defualts to n_t
+        Number of latent dimensions. Defualts to n_timepoints
     T_mean : float
         Mean of weights.
     T_std : float
@@ -44,15 +44,15 @@ def generate_dummy_signal(
 
     Returns
     -------
-    imgs_train : ndarray of shape (n_s, n_t, n_v)
+    imgs_train : ndarray of shape (n_subjects, n_timepoints, n_voxels)
         Training data.
-    imgs_test : ndarray of shape (n_s, n_t, n_v)
+    imgs_test : ndarray of shape (n_subjects, n_timepoints, n_voxels)
         Testing data.
-    S_train : ndarray of shape (n_t, latent_dim)
+    S_train : ndarray of shape (n_timepoints, latent_dim)
         Training latent variables.
-    S_test : ndarray of shape (n_t, latent_dim)
+    S_test : ndarray of shape (n_timepoints, latent_dim)
         Testing latent variables.
-    Ts : ndarray of shape (n_s, latent_dim , n_v)
+    Ts : ndarray of shape (n_subjects, latent_dim , n_voxels)
         Tuning matrices.
     """
     if latent_dim is None:
diff --git a/fmralign/hyperalignment/toy_experiment.py b/fmralign/hyperalignment/toy_experiment.py
index fe3630c..499530c 100644
--- a/fmralign/hyperalignment/toy_experiment.py
+++ b/fmralign/hyperalignment/toy_experiment.py
@@ -1,3 +1,9 @@
+"""Toy experiment to test INT on the two parts of the data (ie different runs
+of the experiment) to acess the validity of tuning computation. This code as
+no vocation to be an explanatory example, it is only used to test the INT
+method. It is not intended to be used as a tutorial.
+"""
+
 import numpy as np
 import matplotlib.pyplot as plt
 from fmralign.alignment_methods import IndividualizedNeuralTuning as INT
@@ -15,9 +21,9 @@
 #############################################################################
 
 
-n_s = 10
-n_t = 200
-n_v = 500
+n_subjects = 10
+n_timepoints = 200
+n_voxels = 500
 S_std = 5
 T_std = 1
 SNR = 100
@@ -35,9 +41,9 @@
     stimulus_run_2,
     data_tuning,
 ) = generate_dummy_signal(
-    n_subjects=n_s,
-    n_timepoints=n_t,
-    n_voxels=n_v,
+    n_subjects=n_subjects,
+    n_timepoints=n_timepoints,
+    n_voxels=n_voxels,
     S_std=S_std,
     T_std=T_std,
     latent_dim=latent_dim,
@@ -49,7 +55,7 @@
 
 if SEARCHLIGHT:
     searchlights, dists = generate_dummy_searchlights(
-        n_searchlights=n_v, n_v=n_v, radius=5
+        n_searchlights=n_voxels, n_v=n_voxels, radius=5
     )
     int1 = INT(
         n_components=latent_dim,
@@ -70,7 +76,7 @@
 
 
 else:
-    parcels = [range(n_v)]
+    parcels = [range(n_voxels)]
     int1 = INT(
         n_components=latent_dim,
         decomp_method=decomposition_method,
@@ -116,7 +122,7 @@
 ax[0, 0].set_ylabel("Subjects, Run 2")
 fig.colorbar(ax[0, 0].imshow(correlation_tuning), ax=ax[0, 0])
 
-random_colors = np.random.rand(n_s, 3)
+random_colors = np.random.rand(n_subjects, 3)
 # MDS of predicted images
 corr_tunning = compute_pearson_corr(data_pred, data_run_2)
 data_pred_reduced, data_test_reduced = matrix_MDS(
@@ -190,7 +196,7 @@
 plt.rc("font", size=10)
 # Define small font for titles
 fig.suptitle(
-    f"Correlation Run 1/2\n ns={n_s}, nt={n_t}, nv={n_v}, S_std={S_std}, T_std={T_std}, SNR={SNR}, latent space dim={latent_dim}"
+    f"Correlation Run 1/2\n ns={n_subjects}, nt={n_timepoints}, nv={n_voxels}, S_std={S_std}, T_std={T_std}, SNR={SNR}, latent space dim={latent_dim}"
 )
 plt.tight_layout()
 

From 5bc14828ca89f362aa62fa09e3bec53c8a7e171a Mon Sep 17 00:00:00 2001
From: FOUCHARD Denis <denis.fouchard@inria.fr>
Date: Fri, 23 Feb 2024 13:33:33 +0100
Subject: [PATCH 49/69] delete useless srm experiment (should already be in
 templateAlignment anyway)

---
 examples/plot_int_alignment.py |   2 +-
 examples/plot_srm_alignment.py | 205 ---------------------------------
 2 files changed, 1 insertion(+), 206 deletions(-)
 delete mode 100644 examples/plot_srm_alignment.py

diff --git a/examples/plot_int_alignment.py b/examples/plot_int_alignment.py
index 5b51dbc..0687102 100644
--- a/examples/plot_int_alignment.py
+++ b/examples/plot_int_alignment.py
@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 
 """
-Hyperalignment-base prediction using the IndividualNeuralTuning Model.
+Co-smoothing Prediction using the IndividualNeuralTuning Model.
 See article : https://doi.org/10.1162/imag_a_00032
 
 ==========================
diff --git a/examples/plot_srm_alignment.py b/examples/plot_srm_alignment.py
deleted file mode 100644
index 3352638..0000000
--- a/examples/plot_srm_alignment.py
+++ /dev/null
@@ -1,205 +0,0 @@
-# -*- coding: utf-8 -*-
-
-"""
-Hyperalignment-base prediction using Hugo Richard's FastSRM.
-See article : https://doi.org/10.1162/imag_a_00032
-
-==========================
-
-In this tutorial, we show how to better predict new contrasts for a target
-subject using many source subjects corresponding contrasts. For this purpose,
-we create a template to which we align the target subject, using shared information.
-We then predict new images for the target and compare them to a baseline.
-
-We mostly rely on Python common packages and on nilearn to handle
-functional data in a clean fashion.
-
-
-To run this example, you must launch IPython via ``ipython
---matplotlib`` in a terminal, or use ``jupyter-notebook``.
-
-.. contents:: **Contents**
-    :local:
-    :depth: 1
-
-"""
-
-###############################################################################
-# Retrieve the data
-# -----------------
-# In this example we use the IBC dataset, which includes a large number of
-# different contrasts maps for 12 subjects.
-# We download the images for subjects sub-01, sub-02, sub-04, sub-05, sub-06
-# and sub-07 (or retrieve them if they were already downloaded).
-# imgs is the list of paths to available statistical images for each subjects.
-# df is a dataframe with metadata about each of them.
-# mask is a binary image used to extract grey matter regions.
-#
-
-from fmralign.fetch_example_data import fetch_ibc_subjects_contrasts
-
-imgs, df, mask_img = fetch_ibc_subjects_contrasts(
-    ["sub-01", "sub-02", "sub-04", "sub-05", "sub-06", "sub-07"]
-)
-
-###############################################################################
-# Definine a masker
-# -----------------
-# We define a nilearn masker that will be used to handle relevant data.
-#   For more information, visit :
-#   'http://nilearn.github.io/manipulating_images/masker_objects.html'
-#
-
-from nilearn.maskers import NiftiMasker
-
-masker = NiftiMasker(mask_img=mask_img).fit()
-
-###############################################################################
-# Prepare the data
-# ----------------
-# For each subject, we will use two series of contrasts acquired during
-# two independent sessions with a different phase encoding:
-# Antero-posterior(AP) or Postero-anterior(PA).
-#
-
-
-# To infer a template for subjects sub-01 to sub-06 for both AP and PA data,
-# we make a list of 4D niimgs from our list of list of files containing 3D images
-
-from nilearn.image import concat_imgs
-
-template_train = []
-for i in range(6):
-    template_train.append(concat_imgs(imgs[i]))
-target_train = df[df.subject == "sub-07"][df.acquisition == "ap"].path.values
-
-# For subject sub-07, we split it in two folds:
-#   - target train: sub-07 AP contrasts, used to learn alignment to template
-#   - target test: sub-07 PA contrasts, used as a ground truth to score predictions
-# We make a single 4D Niimg from our list of 3D filenames
-
-target_train = concat_imgs(target_train)
-target_train_data = masker.transform(target_train)
-target_test = df[df.subject == "sub-07"][df.acquisition == "pa"].path.values
-
-###############################################################################
-# Compute a baseline (average of subjects)
-# ----------------------------------------
-# We create an image with as many contrasts as any subject representing for
-# each contrast the average of all train subjects maps.
-#
-
-import numpy as np
-
-masked_imgs = [masker.transform(img) for img in template_train]
-average_img = np.mean(masked_imgs[:-1], axis=0)
-average_subject = masker.inverse_transform(average_img)
-
-###############################################################################
-# Create a template from the training subjects.
-# ---------------------------------------------
-# We define an estimator using the class TemplateAlignment:
-#   * We align the whole brain through 'multiple' local alignments.
-#   * These alignments are calculated on a parcellation of the brain in 150 pieces,
-#     this parcellation creates group of functionnally similar voxels.
-#   * The template is created iteratively, aligning all subjects data into a
-#     common space, from which the template is inferred and aligning again to this
-#     new template space.
-#
-
-from nilearn.image import index_img
-from fastsrm.identifiable_srm import IdentifiableFastSRM as SRM
-
-
-train_index = range(53)
-model = SRM(n_components=20, n_iter=500)
-
-train_imgs = np.array(masked_imgs)[:, train_index, :]
-train_imgs = [x.T for x in train_imgs]
-model.fit(train_imgs)
-
-
-###############################################################################
-# Predict new data for left-out subject
-# -------------------------------------
-# We use target_train data to fit the transform, indicating it corresponds to
-# the contrasts indexed by train_index and predict from this learnt alignment
-# contrasts corresponding to template test_index numbers.
-# For each train subject and for the template, the AP contrasts are sorted from
-# 0, to 53, and then the PA contrasts from 53 to 106.
-#
-test_index = range(53, 106)
-
-test_imgs = np.array(masked_imgs)[:, test_index, :][:-1]
-test_imgs = [x.T for x in test_imgs]
-
-model2 = SRM(n_components=20, n_iter=500)
-shared_response = model2.fit_transform(test_imgs)
-
-# We input the mapping image target_train in a list, we could have input more
-# than one subject for which we'd want to predict : [train_1, train_2 ...]
-
-prediction_from_template = model.inverse_transform(
-    shared_response, subjects_indexes=[5]
-)
-
-prediction_from_template = prediction_from_template[0].T
-prediction_from_template = [masker.inverse_transform(prediction_from_template)]
-
-# As a baseline prediction, let's just take the average of activations across subjects.
-
-prediction_from_average = index_img(average_subject, test_index)
-
-###############################################################################
-# Score the baseline and the prediction
-# -------------------------------------
-# We use a utility scoring function to measure the voxelwise correlation
-# between the prediction and the ground truth. That is, for each voxel, we
-# measure the correlation between its profile of activation without and with
-# alignment, to see if alignment was able to predict a signal more alike the ground truth.
-#
-
-from fmralign.metrics import score_voxelwise
-
-# Now we use this scoring function to compare the correlation of predictions
-# made from group average and from template with the real PA contrasts of sub-07
-
-average_score = masker.inverse_transform(
-    C_avg := score_voxelwise(target_test, prediction_from_average, masker, loss="corr")
-)
-
-# I choose abs value in reference to the work we did with the INT
-template_score = masker.inverse_transform(
-    C_temp := score_voxelwise(
-        target_test, prediction_from_template[0], masker, loss="corr"
-    )
-)
-
-print("============Mean correlation============")
-print(f"Baseline : {np.mean(C_avg)}")
-print(f"Template : {np.mean(C_temp)}")
-
-###############################################################################
-# Plotting the measures
-# ---------------------
-# Finally we plot both scores
-#
-
-from nilearn import plotting
-
-baseline_display = plotting.plot_stat_map(
-    average_score, display_mode="z", vmax=1, cut_coords=[-15, -5]
-)
-baseline_display.title("Group average correlation wt ground truth")
-display = plotting.plot_stat_map(
-    template_score, display_mode="z", cut_coords=[-15, -5], vmax=1
-)
-display.title("SRM-based prediction correlation wt ground truth")
-
-###############################################################################
-# We observe that creating a template and aligning a new subject to it yields
-# a prediction that is better correlated with the ground truth than just using
-# the average activations of subjects.
-#
-
-plotting.show()

From 8e691e4dd704f759cea9d81380a00f952d1fd5e5 Mon Sep 17 00:00:00 2001
From: FOUCHARD Denis <denis.fouchard@inria.fr>
Date: Mon, 26 Feb 2024 10:19:38 +0100
Subject: [PATCH 50/69] fix linting issues

---
 fmralign/alignment_methods.py | 43 +++++++++++++++++++++++------------
 1 file changed, 29 insertions(+), 14 deletions(-)

diff --git a/fmralign/alignment_methods.py b/fmralign/alignment_methods.py
index b813d61..e2ce51b 100644
--- a/fmralign/alignment_methods.py
+++ b/fmralign/alignment_methods.py
@@ -6,6 +6,7 @@
 import scipy
 from joblib import Parallel, delayed
 from scipy import linalg
+import ot
 from scipy.optimize import linear_sum_assignment
 from scipy.sparse import diags
 from scipy.spatial.distance import cdist
@@ -372,7 +373,6 @@ def fit(self, X, Y):
         Y: (n_samples, n_features) nd array
             target data
         """
-        import ot
 
         n = len(X.T)
         if n > 5000:
@@ -474,8 +474,10 @@ def transform(self, X):
 class IndividualizedNeuralTuning(Alignment):
     """
     Method of alignment based on the Individualized Neural Tuning model.
-    It works on 4D fMRI data, and is based on the assumption that the neural response to a stimulus is shared across subjects.
-    It uses searchlight/parcelation alignment to denoise the data, and then computes the stimulus response matrix.
+    It works on 4D fMRI data, and is based on the assumption that the neural
+    response to a stimulus is shared across subjects. It uses searchlight/
+    parcelation alignment to denoise the data, and then computes the stimulus
+    response matrix.
     See article : https://doi.org/10.1162/imag_a_00032
     """
 
@@ -492,11 +494,17 @@ def __init__(
         Parameters:
         --------
         decomp_method : str
-             The decomposition method to use. Can be ["pca", "pcav1", "procrustes"] Default is "pca".
+             The decomposition method to use.
+             Can be ["pca", "pcav1", "procrustes"]
+             Default is "pca".
         alignment_method : str
-             The alignment method to use. Can be either "searchlight" or "parcelation", Default is "searchlight".
+             The alignment method to use.
+             Can be either "searchlight" or "parcelation",
+             Default is "searchlight".
         n_components : int
-             The number of latent dimensions to use in the shared stimulus information matrix. Default is None.
+             The number of latent dimensions to use in the shared stimulus
+             information
+             matrix. Default is None.
         n_jobs : int
              The number of parallel jobs to run. Default is -1.
 
@@ -524,7 +532,7 @@ def __init__(
         self.n_components = n_components
         self.n_jobs = n_jobs
 
-    #######################################################################################
+    ################################################################
     # Computing decomposition
 
     @staticmethod
@@ -588,13 +596,16 @@ def _reconstruct_signal(shared_response, individual_tuning):
         Parameters:
         --------
         shared_response : numpy.ndarray
-             The shared response of shape (n_timeframes, n_timeframes) or (n_timeframes, latent_dim).
+             The shared response of shape (n_timeframes, n_timeframes) or
+             (n_timeframes, latent_dim).
         individual_tuning : numpy.ndarray
-             The individual tuning of shape (latent_dim, n_voxels) or (n_timeframes, n_voxels).
+             The individual tuning of shape (latent_dim, n_voxels) or
+             (n_timeframes, n_voxels).
 
         Returns:
         --------
-        ndarray: The reconstructed signal of shape (n_timeframes, n_voxels) (same shape as the original signal)
+        ndarray: The reconstructed signal of shape (n_timeframes, n_voxels)
+        (same shape as the original signal)
         """
         return (shared_response @ individual_tuning).astype(np.float32)
 
@@ -617,13 +628,17 @@ def fit(
         X : array-like
             The training data of shape (n_subjects, n_samples, n_voxels).
         searchlights : array-like
-            The searchlight indices for each subject, of shape (n_s, n_searchlights).
+            The searchlight indices for each subject,
+            of shape (n_s, n_searchlights).
         parcels : array-like
-            The parcel indices for each subject, of shape (n_s, n_parcels) (if not using searchlights)
+            The parcel indices for each subject,
+            of shape (n_s, n_parcels) (if not using searchlights)
         dists : array-like
-            The distances of vertices to the center of their searchlight, of shape (n_searchlights, n_vertices_sl)
+            The distances of vertices to the center of their searchlight,
+            of shape (n_searchlights, n_vertices_sl)
         radius : int(optional)
-            The radius of the searchlight sphere, in milimeters. Defaults to 20.
+            The radius of the searchlight sphere, in milimeters.
+            Defaults to 20.
         tuning :bool(optional)
             Whether to compute the tuning weights. Defaults to True.
         verbose : bool(optional)

From a9a329d9a5f8cd77383c22f13a3abc3c611d2f52 Mon Sep 17 00:00:00 2001
From: FOUCHARD Denis <denis.fouchard@inria.fr>
Date: Mon, 26 Feb 2024 15:21:47 +0100
Subject: [PATCH 51/69] adressing other PR comments

---
 .../individualized_neural_tuning.py           |  2 +-
 fmralign/hyperalignment/local_template.py     | 96 +------------------
 2 files changed, 2 insertions(+), 96 deletions(-)

diff --git a/fmralign/hyperalignment/individualized_neural_tuning.py b/fmralign/hyperalignment/individualized_neural_tuning.py
index 8acad48..159d816 100644
--- a/fmralign/hyperalignment/individualized_neural_tuning.py
+++ b/fmralign/hyperalignment/individualized_neural_tuning.py
@@ -5,7 +5,7 @@
 import numpy as np
 
 
-class IndividualizedNeuralTuning(BaseINT):
+class HA_INT(BaseINT):
     """
     Wrapper for the IndividualTuningModel class to be used in fmralign with Niimg objects.
     Preprocessing and searchlight/parcellation alignment are done without any user input.
diff --git a/fmralign/hyperalignment/local_template.py b/fmralign/hyperalignment/local_template.py
index fd67f8a..ec90426 100644
--- a/fmralign/hyperalignment/local_template.py
+++ b/fmralign/hyperalignment/local_template.py
@@ -3,7 +3,6 @@
 """
 
 import numpy as np
-from scipy.stats import zscore
 from sklearn.decomposition import PCA
 from sklearn.utils.extmath import randomized_svd
 
@@ -151,89 +150,6 @@ def compute_PCA_var1_template(
     return XX.astype(np.float32)
 
 
-def compute_procrustes_template(
-    X,
-    region=None,
-    reflection=True,
-    scaling=False,
-    zscore_common=True,
-    level2_iter=1,
-    X2=None,
-):
-    """
-    Compute the Procrustes template for a given set of data.
-
-    Parameters:
-    -----------
-    X : ndarray of shape (n_subjects, n_timepoints, n_voxels)
-        The input data array.
-    region : arraylike
-        The index of the region to consider. If None, all regions are considered. Defaults to None.
-    reflection : bool (optional)
-        Whether to allow reflection in the Procrustes alignment. Defaults to True.
-    scaling bool (optional)
-        Whether to allow scaling in the Procrustes alignment. Defaults to False.
-    zscore_common : bool(optional)
-        Whether to z-score the aligned data to have zero mean and unit variance. Defaults to True.
-    level2_iter : int(optional)
-        The number of iterations for the level 2 alignment. Defaults to 1.
-    X2 : ndarray(optional)
-        The second set of input data array of shape (n_samples, n_features, n_regions). Only used for level 2 alignment. Defaults to None.
-
-    Returns:
-    --------
-    common_space: ndarray
-        The computed Procrustes template.
-    """
-    if region is not None:
-        X = X[:, :, region]
-    common_space = np.copy(X[0])
-    aligned_X = [X[0]]
-    iter_X = X[1:]
-    for x in iter_X:
-        T = procrustes(x, common_space, reflection=reflection, scaling=scaling)
-        aligned_x = x.dot(T)
-        if zscore_common:
-            aligned_x = np.nan_to_num(zscore(aligned_x, axis=0))
-        aligned_X.append(aligned_x)
-        common_space = (common_space + aligned_x) * 0.5
-        if zscore_common:
-            common_space = np.nan_to_num(zscore(common_space, axis=0))
-
-    aligned_X2 = []
-    iter2 = range(level2_iter)
-
-    for level2 in iter2:
-        common_space = np.zeros_like(X[0])
-        for x in aligned_X:
-            common_space += x
-        for i, x in enumerate(X):
-            reference = (common_space - aligned_X[i]) / float(len(X) - 1)
-            if zscore_common:
-                reference = np.nan_to_num(zscore(reference, axis=0))
-            T = procrustes(x, reference, reflection=reflection, scaling=scaling)
-            if level2 == level2_iter - 1 and X2 is not None:
-                aligned_X2.append(X2[i].dot(T))
-            aligned_X[i] = x.dot(T)
-
-    common_space = np.sum(aligned_X, axis=0)
-    if zscore_common:
-        common_space = np.nan_to_num(zscore(common_space, axis=0))
-    else:
-        common_space /= float(len(X))
-    if X2 is not None:
-        common_space2 = np.zeros_like(X2[0])
-        for x in aligned_X2:
-            common_space2 += x
-        if zscore_common:
-            common_space2 = np.nan_to_num(zscore(common_space2, axis=0))
-        else:
-            common_space2 /= float(len(X))
-        return common_space, common_space2
-
-    return common_space
-
-
 def compute_template(
     X,
     region,
@@ -273,18 +189,8 @@ def compute_template(
     mapping = {
         "pca": compute_PCA_template,
         "pcav1": compute_PCA_var1_template,
-        "procrustes": compute_procrustes_template,
     }
-
-    if kind == "procrustes":
-        tmpl = compute_procrustes_template(
-            X=X,
-            region=region,
-            reflection=True,
-            scaling=False,
-            zscore_common=True,
-        )
-    elif kind in mapping:
+    if kind in mapping:
         tmpl = mapping[kind](X, sl=region, n_components=n_components, demean=demean)
     else:
         raise ValueError("Unknown template kind")

From 4074bf9aeb93c3f180f3cbda1ea306948373e687 Mon Sep 17 00:00:00 2001
From: FOUCHARD Denis <denis.fouchard@inria.fr>
Date: Mon, 26 Feb 2024 15:55:16 +0100
Subject: [PATCH 52/69] Fix typo in flavor parameter description

---
 fmralign/hyperalignment/local_template.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/fmralign/hyperalignment/local_template.py b/fmralign/hyperalignment/local_template.py
index ec90426..169acd4 100644
--- a/fmralign/hyperalignment/local_template.py
+++ b/fmralign/hyperalignment/local_template.py
@@ -22,7 +22,7 @@ def PCA_decomposition(
     n_components : int or None
         The number of components to keep. If None, all components are kept.
     flavor : {'sklearn', 'svd'}
-        Wethter to use sklearn or the custom SVD implementation.
+        Whether to use sklearn or the custom SVD implementation.
     adjust_ns : bool
         Whether to adjust the variance of the output so that it doesn't increase with the number of subjects.
     demean : bool

From 008422e296453129d025dd61af6e5b51a788b060 Mon Sep 17 00:00:00 2001
From: FOUCHARD Denis <denis.fouchard@inria.fr>
Date: Mon, 26 Feb 2024 15:57:17 +0100
Subject: [PATCH 53/69] update int tests

---
 .../hyperalignment/test_hyperalignment.py     | 21 ++++++++-----------
 1 file changed, 9 insertions(+), 12 deletions(-)

diff --git a/fmralign/hyperalignment/test_hyperalignment.py b/fmralign/hyperalignment/test_hyperalignment.py
index 1236d13..b4a677e 100644
--- a/fmralign/hyperalignment/test_hyperalignment.py
+++ b/fmralign/hyperalignment/test_hyperalignment.py
@@ -2,6 +2,11 @@
 from fmralign.generate_data import generate_dummy_signal, generate_dummy_searchlights
 import numpy as np
 
+from fmralign.hyperalignment.correlation import (
+    tuning_correlation,
+    stimulus_correlation,
+)
+
 
 def test_int_fit_predict():
     """Test if the outputs and arguments of the INT are the correct format"""
@@ -17,10 +22,6 @@ def test_int_fit_predict():
         generative_method="custom",
         seed=0,
     )
-    from fmralign.hyperalignment.correlation import (
-        tuning_correlation,
-        stimulus_correlation,
-    )
 
     # Testing without searchlights
     searchlights = [np.arange(300)]
@@ -50,8 +51,6 @@ def test_int_fit_predict():
     corr3 = stimulus_correlation(S_estimated_second_part.T, S_true_second_part.T)
     corr4 = tuning_correlation(X_pred, X_test)
 
-    # Check that predicted components have the same shape as original data
-
     # Check that the correlation between the two parts of the data is high
     corr1_out = corr1 - np.diag(corr1)
     corr2_out = corr2 - np.diag(corr2)
@@ -61,6 +60,8 @@ def test_int_fit_predict():
     assert 3 * np.mean(corr2_out) < np.mean(np.diag(corr2))
     assert 3 * np.mean(corr3_out) < np.mean(np.diag(corr3))
     assert 3 * np.mean(corr4_out) < np.mean(np.diag(corr4))
+
+    # Check that predicted components have the same shape as original data
     assert int1.tuning_data[0].shape == (6, int1.n_voxels)
     assert int2.tuning_data[0].shape == (6, int2.n_voxels)
     assert int1.shared_response.shape == (int1.n_time_points, 6)
@@ -82,10 +83,6 @@ def test_int_with_searchlight():
     searchlights, dists = generate_dummy_searchlights(
         n_searchlights=10, n_voxels=30, radius=5, seed=0
     )
-    from fmralign.hyperalignment.correlation import (
-        tuning_correlation,
-        stimulus_correlation,
-    )
 
     # Test INT on the two parts of the data (ie different runs of the experiment)
     model1 = INT(n_components=6)
@@ -107,8 +104,6 @@ def test_int_with_searchlight():
     corr3 = stimulus_correlation(stimulus_run_2.T, stimulus_test.T)
     corr4 = tuning_correlation(X_pred, X_test)
 
-    # Check that predicted components have the same shape as original data
-
     # Check that the correlation between the two parts of the data is high
     corr1_out = corr1 - np.diag(corr1)
     corr2_out = corr2 - np.diag(corr2)
@@ -118,6 +113,8 @@ def test_int_with_searchlight():
     assert 3 * np.mean(corr2_out) < np.mean(np.diag(corr2))
     assert 3 * np.mean(corr3_out) < np.mean(np.diag(corr3))
     assert 3 * np.mean(corr4_out) < np.mean(np.diag(corr4))
+
+    # Check that predicted components have the same shape as original data
     assert model1.tuning_data[0].shape == (6, model1.n_voxels)
     assert model2.tuning_data[0].shape == (6, model2.n_voxels)
     assert model1.shared_response.shape == (model1.n_time_points, 6)

From 8fd68639b950e33bcaa69e3457b31499d22c7c3d Mon Sep 17 00:00:00 2001
From: FOUCHARD Denis <denis.fouchard@inria.fr>
Date: Mon, 26 Feb 2024 15:59:06 +0100
Subject: [PATCH 54/69] Better rst  in alignment_methods.py

---
 fmralign/alignment_methods.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/fmralign/alignment_methods.py b/fmralign/alignment_methods.py
index e2ce51b..14c7c6a 100644
--- a/fmralign/alignment_methods.py
+++ b/fmralign/alignment_methods.py
@@ -367,7 +367,7 @@ def fit(self, X, Y):
         """
 
         Parameters
-        --------------
+        ----------
         X: (n_samples, n_features) nd array
             source data
         Y: (n_samples, n_features) nd array
@@ -440,7 +440,7 @@ def fit(self, X, Y):
         """
 
         Parameters
-        --------------
+        ----------
         X: (n_samples, n_features) nd array
             source data
         Y: (n_samples, n_features) nd array
@@ -492,7 +492,7 @@ def __init__(
         Initialize the IndividualizedNeuralTuning object.
 
         Parameters:
-        --------
+        ----------
         decomp_method : str
              The decomposition method to use.
              Can be ["pca", "pcav1", "procrustes"]
@@ -541,7 +541,7 @@ def _tuning_estimator(shared_response, target):
         Estimate the tuning matrix for individualized neural tuning.
 
         Parameters:
-        --------
+        ----------
         shared_response : array-like
              The shared response matrix of shape (n_timepoints, k)
              where k is the dimension of the sources latent space.
@@ -565,7 +565,7 @@ def _stimulus_estimator(full_signal, n_subjects, latent_dim=None, scaling=True):
         Estimates the stimulus matrix for the Individualized Neural Tuning model.
 
         Parameters:
-        --------
+        -----------
         full_signal : numpy.ndarray
              Concatenated signal for all subjects,
              of shape (n_timepoints, n_subjects * n_voxels).
@@ -623,7 +623,7 @@ def fit(
         Fits the IndividualizedNeuralTuning model to the training data.
 
         Parameters:
-        --------
+        -----------
 
         X : array-like
             The training data of shape (n_subjects, n_samples, n_voxels).
@@ -702,7 +702,7 @@ def transform(self, X, verbose=False):
         Transforms the input test data using the hyperalignment model.
 
         Parameters:
-        --------
+        ----------
         X : array-like
             The test data of shape (n_subjects, n_timepoints, n_voxels).
         verbose : bool(optional)

From dad2a61f4245d902a7281e9224efe45acd754069 Mon Sep 17 00:00:00 2001
From: DF <95576336+denisfouchard@users.noreply.github.com>
Date: Mon, 26 Feb 2024 17:18:02 +0100
Subject: [PATCH 55/69] Update fmralign/hyperalignment/correlation.py

Co-authored-by: bthirion <bertrand.thirion@inria.fr>
---
 fmralign/hyperalignment/correlation.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/fmralign/hyperalignment/correlation.py b/fmralign/hyperalignment/correlation.py
index f74099e..41a0e9c 100644
--- a/fmralign/hyperalignment/correlation.py
+++ b/fmralign/hyperalignment/correlation.py
@@ -149,7 +149,7 @@ def stimulus_correlation(X, Y, linear_assignment=True, absolute=True):
 
 def matrix_MDS(X, Y, n_components=2, dissimilarity="euclidean"):
     """
-    Perform multidimensional scaling (MDS) on the given data matrices X and Y.
+    Perform multidimensional scaling (MDS) on the rows of X and Y.
 
     Parameters:
     ----------

From 923f1003e81799d93dee72f103c7f5b72d7ccdda Mon Sep 17 00:00:00 2001
From: FOUCHARD Denis <denis.fouchard@inria.fr>
Date: Mon, 26 Feb 2024 17:18:37 +0100
Subject: [PATCH 56/69] PR comments adress

---
 examples/plot_int_alignment.py            | 4 ++--
 fmralign/generate_data.py                 | 4 ++--
 fmralign/hyperalignment/local_template.py | 2 ++
 3 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/examples/plot_int_alignment.py b/examples/plot_int_alignment.py
index 0687102..4d7159d 100644
--- a/examples/plot_int_alignment.py
+++ b/examples/plot_int_alignment.py
@@ -149,10 +149,10 @@
 # than one subject for which we'd want to predict : [train_1, train_2 ...]
 
 
-prediction_from_template = model._reconstruct_signal(
+pred = model._reconstruct_signal(
     shared_response=stimulus_[test_index], individual_tuning=target_tuning
 )
-prediction_from_template = masker.inverse_transform(prediction_from_template)
+prediction_from_template = masker.inverse_transform(pred)
 
 
 # As a baseline prediction, let's just take the average of activations across subjects.
diff --git a/fmralign/generate_data.py b/fmralign/generate_data.py
index e13aae3..2669558 100644
--- a/fmralign/generate_data.py
+++ b/fmralign/generate_data.py
@@ -121,7 +121,7 @@ def generate_dummy_signal(
 def generate_dummy_searchlights(
     n_searchlights: int,
     n_voxels: int,
-    radius: int,
+    radius: float,
     sl_size: int = 5,
     seed: int = 0,
 ):
@@ -133,7 +133,7 @@ def generate_dummy_searchlights(
         Number of searchlights.
     n_voxels : int
         Number of voxels.
-    radius : int
+    radius : float,
         Radius of searchlights.
     sl_size : int, default=5
         Size of each searchlight (easier for dummy signal generation).
diff --git a/fmralign/hyperalignment/local_template.py b/fmralign/hyperalignment/local_template.py
index 169acd4..60197d8 100644
--- a/fmralign/hyperalignment/local_template.py
+++ b/fmralign/hyperalignment/local_template.py
@@ -31,7 +31,9 @@ def PCA_decomposition(
     Returns
     -------
     XX : ndarray of shape (n_timepoints, n_components)
+        The decomposed data array with reduced dimensionality.
     cc : ndarray of shape (n_components, n_subjects, n_voxels)
+        Column-wise principal components (from Vt)
     """
     ns, nt, nv = X.shape
     X = X.transpose(1, 0, 2).reshape(nt, ns * nv).astype(np.float32)

From a70b4a8b59316ce442a631e35e60eb3990111b77 Mon Sep 17 00:00:00 2001
From: FOUCHARD Denis <denis.fouchard@inria.fr>
Date: Mon, 26 Feb 2024 17:21:43 +0100
Subject: [PATCH 57/69] Refactor dissimilarity measure in matrix_MDS function

---
 fmralign/hyperalignment/correlation.py | 69 ++++++--------------------
 1 file changed, 16 insertions(+), 53 deletions(-)

diff --git a/fmralign/hyperalignment/correlation.py b/fmralign/hyperalignment/correlation.py
index 41a0e9c..f698088 100644
--- a/fmralign/hyperalignment/correlation.py
+++ b/fmralign/hyperalignment/correlation.py
@@ -161,14 +161,17 @@ def matrix_MDS(X, Y, n_components=2, dissimilarity="euclidean"):
     n_components : int
         The number of dimensions in the output space (default is 2).
     dissimilarity : str or array-like
-        The dissimilarity measure to use. If it is a string other than "precomputed",
-        the dissimilarity is computed using the Euclidean distance between flattened data points.
-        If it is "precomputed", the dissimilarity is assumed to be a precomputed dissimilarity matrix.
+        The dissimilarity measure to use.
+        If it is a string other than "precomputed", the dissimilarity is
+        computed using the Euclidean distance between flattened data points.
+        If it is "precomputed", the dissimilarity is assumed to be a
+        precomputed dissimilarity matrix.
 
     Returns:
     -------
-    tuple: A tuple containing two arrays. The first array represents the transformed data points from matrix X,
-            and the second array represents the transformed data points from matrix Y.
+    tuple: A tuple containing two arrays. The first array represents
+        the transformed data points from matrix X, and the second array
+        represents the transformed data points from matrix Y.
     """
     assert len(X) == len(Y)
 
@@ -186,50 +189,6 @@ def matrix_MDS(X, Y, n_components=2, dissimilarity="euclidean"):
     return np.array(transformed[: len(X)]), np.array(transformed[len(X) :])
 
 
-def thread_compute_correlation(X, Y, i, j):
-    """
-    Compute the correlation between two time series X_i and Y_i.
-
-    Parameters:
-    ----------
-    X : ndarray
-        ndrray of shape (n_samples, n_features) representing the first time series.
-    Y : ndarray
-        ndrray of shape (n_samples, n_features) representing the second time series.
-    i : int
-        Index of the first time series.
-    j : int
-        Index of the second time series.
-
-    Returns:
-    -------
-    ndarray
-        Correlations between different time points if i != j, otherwise an empty array.
-    ndarray
-        Correlations between the same time points if i != j, otherwise an empty array.
-    ndarray
-        Correlations between different time points if i == j, otherwise an empty array.
-    ndarray
-        Correlations between the same time points if i == j, otherwise an empty array.
-    """
-    X_i, Y_i = X[i], Y[i]
-    corr = stimulus_correlation(X_i, Y_i, absolute=False)
-    same_TR_corr = np.diag(corr)
-    # Get all the values except the diagonal in a list
-    diff_TR_corr = corr[np.where(~np.eye(corr.shape[0], dtype=bool))]
-    diff_TR_corr = diff_TR_corr.flatten()
-    if i == j:
-        return (
-            np.array([]),
-            np.array([]),
-            [x for x in diff_TR_corr],
-            [x for x in same_TR_corr],
-        )
-
-    else:
-        return diff_TR_corr, same_TR_corr, np.array([]), np.array([])
-
-
 def multithread_compute_correlation(
     X, Y, absolute=False, linear_assignment=True, n_jobs=1
 ):
@@ -255,10 +214,14 @@ def multithread_compute_correlation(
 
     tuple
         A tuple containing four arrays
-        - corr_same_sub_diff_TR: Correlations between different time points of the same subject.
-        - corr_same_sub_same_TR: Correlations between the same time points of the same subject.
-        - corr_diff_sub_diff_TR: Correlations between different time points of different subjects.
-        - corr_diff_sub_same_TR: Correlations between the same time points of different subjects.
+        - corr_same_sub_diff_TR: Correlations between different time points
+        of the same subject.
+        - corr_same_sub_same_TR: Correlations between the same time points
+        of the same subject.
+        - corr_diff_sub_diff_TR: Correlations between different time points
+        of different subjects.
+        - corr_diff_sub_same_TR: Correlations between the same time points
+        of different subjects.
     """
     from joblib import Parallel, delayed
 

From 1711823c3bcb1a02a814ad44766a2ac0a6980bb5 Mon Sep 17 00:00:00 2001
From: FOUCHARD Denis <denis.fouchard@inria.fr>
Date: Mon, 26 Feb 2024 17:24:06 +0100
Subject: [PATCH 58/69] Refactor test_hyperalignment.py to include
 decomposition and searchlight input

---
 fmralign/hyperalignment/test_hyperalignment.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/fmralign/hyperalignment/test_hyperalignment.py b/fmralign/hyperalignment/test_hyperalignment.py
index b4a677e..9ad7540 100644
--- a/fmralign/hyperalignment/test_hyperalignment.py
+++ b/fmralign/hyperalignment/test_hyperalignment.py
@@ -9,7 +9,8 @@
 
 
 def test_int_fit_predict():
-    """Test if the outputs and arguments of the INT are the correct format"""
+    """Test if the outputs and arguments of the INT are the correct format,
+    and if decomposition is working"""
     # Create random data
     X_train, X_test, S_true_first_part, S_true_second_part, _ = generate_dummy_signal(
         n_subjects=7,
@@ -69,6 +70,8 @@ def test_int_fit_predict():
 
 
 def test_int_with_searchlight():
+    """Test if the outputs and arguments of the INT are the correct format and
+    if the decomposition is working, with searchlight input"""
     X_train, X_test, stimulus_train, stimulus_test, _ = generate_dummy_signal(
         n_subjects=5,
         n_timepoints=50,

From 995f82a202f4df5648fc89bdb3b5821f52408f5f Mon Sep 17 00:00:00 2001
From: FOUCHARD Denis <denis.fouchard@inria.fr>
Date: Tue, 27 Feb 2024 10:06:11 +0100
Subject: [PATCH 59/69] Add toy experiment to examples

---
 .../plot_toy_int_experiment.py                | 149 ++++++++++--------
 1 file changed, 80 insertions(+), 69 deletions(-)
 rename fmralign/hyperalignment/toy_experiment.py => examples/plot_toy_int_experiment.py (53%)

diff --git a/fmralign/hyperalignment/toy_experiment.py b/examples/plot_toy_int_experiment.py
similarity index 53%
rename from fmralign/hyperalignment/toy_experiment.py
rename to examples/plot_toy_int_experiment.py
index 499530c..409a3bb 100644
--- a/fmralign/hyperalignment/toy_experiment.py
+++ b/examples/plot_toy_int_experiment.py
@@ -1,13 +1,30 @@
-"""Toy experiment to test INT on the two parts of the data (ie different runs
-of the experiment) to acess the validity of tuning computation. This code as
-no vocation to be an explanatory example, it is only used to test the INT
-method. It is not intended to be used as a tutorial.
+# -*- coding: utf-8 -*-
+
 """
+Co-smoothing Prediction using the IndividualNeuralTuning Model.
+See article : https://doi.org/10.1162/imag_a_00032
+
+==========================
+
+This is a toy experiment to test Individual Tuning Model (INT) on two parts of the 
+data (or different runs) to assess the validity of tuning computation. This code has
+no vocation to be an explanatory example, but rather a test to check the validity of
+the INT model.
+
+
+To run this example, you must launch IPython via ``ipython
+--matplotlib`` in a terminal, or use ``jupyter-notebook``.
+
+.. contents:: **Contents**
+    :local:
+    :depth: 1
 
+"""
+# %%
 import numpy as np
 import matplotlib.pyplot as plt
 from fmralign.alignment_methods import IndividualizedNeuralTuning as INT
-from fmralign.generate_data import generate_dummy_signal, generate_dummy_searchlights
+from fmralign.generate_data import generate_dummy_signal
 from fmralign.hyperalignment.correlation import (
     tuning_correlation,
     stimulus_correlation,
@@ -15,25 +32,29 @@
     matrix_MDS,
 )
 
-
-#############################################################################
-# INT
-#############################################################################
-
+# %%
+###############################################################################
+# Generate the data
+# -----------------
+# In this example we use toy data to test the INT model. We generate two runs of
+# the experiment, and we use the INT model to align the two runs. We then compare
+# the tuning matrices and the shared response to assess the validity of the INT model.
+# We also compare the reconstructed images to the ground truth to assess the validity
+# of the INT model.
+# The toy generation function allows us to get the ground truth stimulus and tuning
+# matrices that were used to generate the data, and we can also control the level of
+# noise in the data.
 
 n_subjects = 10
 n_timepoints = 200
 n_voxels = 500
-S_std = 5
+S_std = 5  # Standard deviation of the source components
 T_std = 1
-SNR = 100
+SNR = 100  # Signal to noise ratio
 latent_dim = 15  # if None, latent_dim = n_t
 decomposition_method = "pca"  # if None, SVD is used
 
 
-#############################################################################
-# GENERATE DUMMY SIGNAL
-
 (
     data_run_1,
     data_run_2,
@@ -51,48 +72,27 @@
     seed=42,
 )
 
-SEARCHLIGHT = False
-
-if SEARCHLIGHT:
-    searchlights, dists = generate_dummy_searchlights(
-        n_searchlights=n_voxels, n_v=n_voxels, radius=5
-    )
-    int1 = INT(
-        n_components=latent_dim,
-        decomp_method=decomposition_method,
-        alignment_method="searchlight",
-    )
-    int2 = INT(
-        n_components=latent_dim,
-        decomp_method=decomposition_method,
-        alignment_method="searchlight",
-    )
-    int_first_part = int1.fit(
-        data_run_1, searchlights=searchlights, dists=dists, radius=5, verbose=False
-    )  # S is provided if we cheat and know the ground truth
-    int_second_part = int2.fit(
-        data_run_2, searchlights=searchlights, dists=dists, radius=5, verbose=False
-    )
-
-
-else:
-    parcels = [range(n_voxels)]
-    int1 = INT(
-        n_components=latent_dim,
-        decomp_method=decomposition_method,
-    )
-    int2 = INT(
-        n_components=latent_dim,
-        decomp_method=decomposition_method,
-    )
-    int_first_part = int1.fit(
-        data_run_1, parcels=parcels, verbose=False
-    )  # S is provided if we cheat and know the ground truth
-    int_second_part = int2.fit(data_run_2, parcels=parcels, verbose=False)
+parcels = [range(n_voxels)]
 
+# %%
 #############################################################################
-# Test INT on the two parts of the data (ie different runs of the experiment)
-
+# Create two independant instances of the model
+# ---------------------------------------------
+# We create two instances of the INT model to align the two runs of the experiment.
+# We then extract the tuning matrices and the shared from the two runs to compare them.
+#
+#
+
+int1 = INT(
+    n_components=latent_dim,
+    decomp_method=decomposition_method,
+)
+int2 = INT(
+    n_components=latent_dim,
+    decomp_method=decomposition_method,
+)
+int1.fit(data_run_1, parcels=parcels, verbose=False)
+int2.fit(data_run_2, parcels=parcels, verbose=False)
 
 # save individual components
 tuning_pred_run_1 = int1.tuning_data
@@ -105,19 +105,25 @@
 
 data_pred = int1.transform(data_run_2)
 
+# %%
+###############################################################################
+# Plotting validation metrics
+# ---------------------------
+# We compare the tuning matrices and the shared response to assess the validity
+# of the INT model. To achieve that, we use Pearson correlation between true and
+# estimated stimulus, as well as between true and estimated tuning matrices.
+# For tuning matrices, this is dones by first computing the correlation between
+# every pair of tuning matrices from the two runs of the experiment, and then
+# averaging the correlation across the diagonal (ie the correlation between
+# the same timepoint of the two runs).
 
-#############################################################################
-# Plot
-#############################################################################
-
-plt.rc("font", size=6)
-fig, ax = plt.subplots(2, 3, figsize=(10, 5))
+fig, ax = plt.subplots(2, 3, figsize=(15, 8))
 
 
 # Tunning matrices
 correlation_tuning = tuning_correlation(tuning_pred_run_1, tuning_pred_run_2)
 ax[0, 0].imshow(correlation_tuning)
-ax[0, 0].set_title("Correlation Tuning Run 1 vs Run 2")
+ax[0, 0].set_title("Pearson correlation of tuning matrices (Run 1 vs Run 2)")
 ax[0, 0].set_xlabel("Subjects, Run 1")
 ax[0, 0].set_ylabel("Subjects, Run 2")
 fig.colorbar(ax[0, 0].imshow(correlation_tuning), ax=ax[0, 0])
@@ -149,7 +155,6 @@
     tuning_pred_run_1, tuning_pred_run_2, n_components=2, dissimilarity=1 - corr_tunning
 )
 
-
 ax[0, 2].scatter(
     T_first_part_transformed[:, 0],
     T_first_part_transformed[:, 1],
@@ -163,13 +168,16 @@
     c=random_colors,
 )
 ax[0, 2].set_title("MDS of tunning matrices, dim=2")
+# Set square aspect
+ax[0, 1].set_aspect("equal", "box")
+ax[0, 2].set_aspect("equal", "box")
 
 # Stimulus matrix correlation
 correlation_stimulus_true_est_first_part = stimulus_correlation(
     stimulus_pred_run_1.T, stimulus_run_1.T
 )
 ax[1, 0].imshow(correlation_stimulus_true_est_first_part)
-ax[1, 0].set_title("Stimumus Estimated vs ground truth (Run 1)")
+ax[1, 0].set_title("Correlation of estimated stimulus vs ground truth (Run 1)")
 ax[1, 0].set_xlabel("Latent components, Run 1")
 ax[1, 0].set_ylabel("Latent components, ground truth")
 fig.colorbar(ax[1, 0].imshow(correlation_stimulus_true_est_first_part), ax=ax[1, 0])
@@ -178,7 +186,7 @@
     stimulus_pred_run_2.T, stimulus_run_2.T
 )
 ax[1, 1].imshow(correlation_stimulus_true_est_second_part)
-ax[1, 1].set_title("Stimulus Estimated vs ground truth (Run 2)")
+ax[1, 1].set_title("Correlation of estimated stimulus vs ground truth (Run 2))")
 ax[1, 1].set_xlabel("Latent components, Run 2")
 ax[1, 1].set_ylabel("Latent components, ground truth")
 fig.colorbar(ax[1, 1].imshow(correlation_stimulus_true_est_second_part), ax=ax[1, 1])
@@ -187,17 +195,20 @@
 # Reconstruction
 corr_reconstruction = tuning_correlation(data_pred, data_run_2)
 ax[1, 2].imshow(corr_reconstruction)
-ax[1, 2].set_title("Reconstruction correlation")
+ax[1, 2].set_title("Correlation of brain response (Run 2 vs Ground truth)")
 ax[1, 2].set_xlabel("Subjects, Run 2")
-ax[1, 2].set_ylabel("Subjects, Run 1")
+ax[1, 2].set_ylabel("Subjects, Ground truth")
 fig.colorbar(ax[1, 2].imshow(corr_reconstruction), ax=ax[1, 2])
 
 
 plt.rc("font", size=10)
 # Define small font for titles
 fig.suptitle(
-    f"Correlation Run 1/2\n ns={n_subjects}, nt={n_timepoints}, nv={n_voxels}, S_std={S_std}, T_std={T_std}, SNR={SNR}, latent space dim={latent_dim}"
+    "Correlation metrics for the Individual Tuning Model\n"
+    + f"{n_subjects} subjects, {n_timepoints} timepoints, {n_voxels} voxels, {latent_dim} latent components\n"
+    + f"SNR={SNR}"
 )
-plt.tight_layout()
 
+plt.tight_layout()
+# %%
 plt.show()

From e03309798e9f40aa6df77c20a89c26dc25e47270 Mon Sep 17 00:00:00 2001
From: FOUCHARD Denis <denis.fouchard@inria.fr>
Date: Tue, 27 Feb 2024 10:07:38 +0100
Subject: [PATCH 60/69] RM wrapper

---
 .../individualized_neural_tuning.py           | 157 ------------------
 1 file changed, 157 deletions(-)
 delete mode 100644 fmralign/hyperalignment/individualized_neural_tuning.py

diff --git a/fmralign/hyperalignment/individualized_neural_tuning.py b/fmralign/hyperalignment/individualized_neural_tuning.py
deleted file mode 100644
index 159d816..0000000
--- a/fmralign/hyperalignment/individualized_neural_tuning.py
+++ /dev/null
@@ -1,157 +0,0 @@
-from fmralign.alignment_methods import IndividualizedNeuralTuning as BaseINT
-from .regions import compute_searchlights, compute_parcels
-from nilearn.maskers import NiftiMasker
-from nibabel import Nifti1Image
-import numpy as np
-
-
-class HA_INT(BaseINT):
-    """
-    Wrapper for the IndividualTuningModel class to be used in fmralign with Niimg objects.
-    Preprocessing and searchlight/parcellation alignment are done without any user input.
-
-    Method of alignment based on the Individualized Neural Tuning model.
-    It uses searchlight/parcelation alignment to denoise the data, and then computes the stimulus response matrix.
-    See article : https://doi.org/10.1162/imag_a_00032
-    """
-
-    def __init__(
-        self,
-        template="pca",
-        decomp_method=None,
-        alignment_method="searchlight",
-        n_pieces=150,
-        searchlight_radius=20,
-        n_components=None,
-        n_jobs=1,
-    ):
-        """
-        Initialize the IndividualizedNeuralTuning object.
-
-        Parameters:
-        -----------
-        tmpl_kind : str
-            The type of template used for alignment. Default is "pca".
-        decomp_method : str
-            The decomposition method used for template construction. Default is None.
-        alignment_method : str in ["searchlight", "parcellation"
-            The alignment method used. Default is "searchlight".
-        n_pieces : int
-            The number of pieces to divide the data into if using parcellation. Default is 150.
-        radius : int
-            The radius of the searchlight sphere in millimeters.
-            Only used if alignment_method is "searchlight".
-            Default is 20.
-        n_components : int
-            The number of latent dimensions to use. If None, all the components are used.
-            Default is None.
-        n_jobs : int
-            The number of parallel jobs to run. Default is 1.
-        """
-        super().__init__(
-            template=template,
-            decomp_method=decomp_method,
-            n_components=n_components,
-            alignment_method=alignment_method,
-            n_jobs=n_jobs,
-        )
-        self.n_pieces = n_pieces
-        self.radius = searchlight_radius
-        self.mask_img = None
-        self.masker = None
-
-    def fit(
-        self,
-        imgs,
-        masker: NiftiMasker = None,
-        mask_img: Nifti1Image = None,
-        tuning: bool = True,
-        y=None,
-        verbose=0,
-    ):
-        """
-        Fit the model to the data.
-        Can either take as entry a masking image or a masker object.
-        This information will be kept for transforming the data.
-
-        Parameters
-        ----------
-        imgs : list of Nifti1Image
-            The images to be fit.
-        masker : NiftiMasker
-            The masker to be used to transform the images into the common space.
-        mask_img : Nifti1Image
-            The mask to be used to transform the images into the common space.
-        tuning : bool
-            Whether to perform tuning or not.
-        y : None
-            Not used.
-        verbose : int
-            The verbosity level.
-
-        Returns
-        -------
-        self : IndividualizedNeuralTuning
-            The fitted model.
-        """
-
-        if mask_img is None:
-            mask_img = masker.mask_img
-        self.mask_img = mask_img
-        self.masker = masker
-
-        X = np.array([masker.transform(img) for img in imgs])
-
-        if self.alignment_method == "searchlight":
-            _, searchlights, dists = compute_searchlights(
-                niimg=imgs[0],
-                mask_img=mask_img,
-            )
-            super().fit(
-                X,
-                searchlights,
-                dists,
-                radius=self.radius,
-                tuning=tuning,
-                verbose=verbose,
-            )
-
-        elif self.alignment_method == "parcellation":
-            parcels = compute_parcels(
-                niimg=imgs[0],
-                mask=masker.mask_img,
-                n_parcels=self.n_pieces,
-                n_jobs=self.n_jobs,
-            )
-            super().fit(X, parcels=parcels, tuning=tuning, verbose=verbose)
-
-        return self
-
-    def transform(self, imgs, y=None, verbose=0):
-        """
-        Transform the data into the common space and return the Nifti1Image objects.
-
-        Parameters
-        ----------
-        imgs : list of Nifti1Image
-            The images to be transformed into the common space.
-        y : None
-            Not used.
-        verbose : int
-            The verbosity level.
-
-        Returns
-        -------
-        preds : list of Nifti1Image
-            The images in the common space.
-        """
-        if self.masker is None:
-            self.masker = NiftiMasker(mask_img=self.mask_img)
-
-        X = self.masker.fit_transform(imgs)
-        Y = super().transform(X, verbose=verbose)
-        preds = []
-        for y in Y:
-            preds.append(self.masker.inverse_transform(y))
-
-        return preds

From 70c0d00317d990bf1d093ae4a8c656a66ca65fe8 Mon Sep 17 00:00:00 2001
From: FOUCHARD Denis <denis.fouchard@inria.fr>
Date: Tue, 27 Feb 2024 10:10:53 +0100
Subject: [PATCH 61/69] adress PR comments

---
 fmralign/hyperalignment/test_hyperalignment.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/fmralign/hyperalignment/test_hyperalignment.py b/fmralign/hyperalignment/test_hyperalignment.py
index 9ad7540..2961f8b 100644
--- a/fmralign/hyperalignment/test_hyperalignment.py
+++ b/fmralign/hyperalignment/test_hyperalignment.py
@@ -10,7 +10,8 @@
 
 def test_int_fit_predict():
     """Test if the outputs and arguments of the INT are the correct format,
-    and if decomposition is working"""
+    and if decomposition is working. Without proper searchlight input
+    (ie all voxels are used)"""
     # Create random data
     X_train, X_test, S_true_first_part, S_true_second_part, _ = generate_dummy_signal(
         n_subjects=7,

From 850a2a15f0d8dcde3df56dd69a824b8b1ce0c62e Mon Sep 17 00:00:00 2001
From: FOUCHARD Denis <denis.fouchard@inria.fr>
Date: Tue, 27 Feb 2024 10:13:29 +0100
Subject: [PATCH 62/69] flake8

---
 examples/plot_toy_int_experiment.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/examples/plot_toy_int_experiment.py b/examples/plot_toy_int_experiment.py
index 409a3bb..af8c456 100644
--- a/examples/plot_toy_int_experiment.py
+++ b/examples/plot_toy_int_experiment.py
@@ -6,7 +6,7 @@
 
 ==========================
 
-This is a toy experiment to test Individual Tuning Model (INT) on two parts of the 
+This is a toy experiment to test Individual Tuning Model (INT) on two parts of the
 data (or different runs) to assess the validity of tuning computation. This code has
 no vocation to be an explanatory example, but rather a test to check the validity of
 the INT model.
@@ -80,8 +80,6 @@
 # ---------------------------------------------
 # We create two instances of the INT model to align the two runs of the experiment.
 # We then extract the tuning matrices and the shared from the two runs to compare them.
-#
-#
 
 int1 = INT(
     n_components=latent_dim,

From 657233a62648c6e779419bc9161baa387dadb561 Mon Sep 17 00:00:00 2001
From: FOUCHARD Denis <denis.fouchard@inria.fr>
Date: Wed, 28 Feb 2024 11:20:39 +0100
Subject: [PATCH 63/69] rm region arguments from fit method + refined doc for
 example

---
 examples/plot_int_alignment.py                |  12 +-
 fmralign/alignment_methods.py                 | 103 +++++++++---------
 .../hyperalignment/piecewise_alignment.py     |   2 -
 .../hyperalignment/test_hyperalignment.py     |  18 ++-
 4 files changed, 67 insertions(+), 68 deletions(-)

diff --git a/examples/plot_int_alignment.py b/examples/plot_int_alignment.py
index 4d7159d..96095c8 100644
--- a/examples/plot_int_alignment.py
+++ b/examples/plot_int_alignment.py
@@ -134,15 +134,19 @@
 
 
 parcels = compute_parcels(niimg=template_train[0], mask=masker, n_parcels=100, n_jobs=5)
-denoiser = PiecewiseAlignment(alignment_method="parcelation", n_jobs=5)
+denoiser = PiecewiseAlignment(n_jobs=5)
 denoised_signal = denoiser.fit_transform(X=denoising_data, regions=parcels)
+target_denoised_data = denoised_signal[-1]
 model = IndividualizedNeuralTuning(
-    n_jobs=8, alignment_method="parcelation", n_components=None
+    parcels=parcels,
 )
-model.fit(training_data, parcels=parcels, verbose=False)
+model.fit(training_data, verbose=False)
 stimulus_ = np.copy(model.shared_response)
+
+# From the denoised data and the stimulus, we can now extract the tuning
+# matrix from sub-07 AP contrasts, and use it to predict the PA contrasts.
 target_tuning = model._tuning_estimator(
-    shared_response=stimulus_[train_index], target=denoised_signal[-1]
+    shared_response=stimulus_[train_index], target=target_denoised_data
 )
 # %%
 # We input the mapping image target_train in a list, we could have input more
diff --git a/fmralign/alignment_methods.py b/fmralign/alignment_methods.py
index 14c7c6a..e3b8b5b 100644
--- a/fmralign/alignment_methods.py
+++ b/fmralign/alignment_methods.py
@@ -485,7 +485,11 @@ def __init__(
         self,
         decomp_method="pca",
         n_components=None,
-        alignment_method="searchlight",
+        searchlights=None,
+        parcels=None,
+        dists=None,
+        radius=20,
+        tuning=True,
         n_jobs=1,
     ):
         """
@@ -497,10 +501,20 @@ def __init__(
              The decomposition method to use.
              Can be ["pca", "pcav1", "procrustes"]
              Default is "pca".
-        alignment_method : str
-             The alignment method to use.
-             Can be either "searchlight" or "parcelation",
-             Default is "searchlight".
+        searchlights : array-like
+            The searchlight indices for each subject,
+            of shape (n_s, n_searchlights).
+        parcels : array-like
+            The parcel indices for each subject,
+            of shape (n_s, n_parcels) (if not using searchlights)
+        dists : array-like
+            The distances of vertices to the center of their searchlight,
+            of shape (n_searchlights, n_vertices_sl)
+        radius : int(optional)
+            The radius of the searchlight sphere, in milimeters.
+            Defaults to 20.
+        tuning : bool(optional)
+            Whether to compute the tuning weights. Defaults to True.
         n_components : int
              The number of latent dimensions to use in the shared stimulus
              information
@@ -517,14 +531,23 @@ def __init__(
         self.n_time_points = None
         self.labels = None
         self.alphas = None
-        self.alignment_method = alignment_method
-        if alignment_method == "parcelation":
-            self.parcels = None
 
-        elif alignment_method == "searchlight":
-            self.searchlights = None
-            self.distances = None
-            self.radius = None
+        if searchlights is None and parcels is None:
+            raise ValueError("searchlights or parcels must be provided")
+
+        if searchlights is not None and parcels is not None:
+            raise ValueError(
+                "searchlights and parcels cannot be provided at the same time"
+            )
+
+        if searchlights is not None:
+            self.regions = searchlights
+        else:
+            self.regions = parcels
+
+        self.dists = dists
+        self.radius = radius
+        self.tuning = tuning
 
         self.tuning_data = []
         self.denoised_signal = []
@@ -566,7 +589,7 @@ def _stimulus_estimator(full_signal, n_subjects, latent_dim=None, scaling=True):
 
         Parameters:
         -----------
-        full_signal : numpy.ndarray
+        full_signal : ndarray
              Concatenated signal for all subjects,
              of shape (n_timepoints, n_subjects * n_voxels).
         n_subjects : int
@@ -575,6 +598,11 @@ def _stimulus_estimator(full_signal, n_subjects, latent_dim=None, scaling=True):
              The number of latent dimensions to use. Defaults to None.
         scaling : bool, optional
              Whether to scale the stimulus matrix sources. Defaults to True.
+
+        Returns:
+        --------
+        stimulus : ndarray
+            The stimulus matrix of shape (n_timepoints, n_subjects * n_voxels)
         """
         n_timepoints = full_signal.shape[0]
         if scaling:
@@ -591,32 +619,28 @@ def _stimulus_estimator(full_signal, n_subjects, latent_dim=None, scaling=True):
     @staticmethod
     def _reconstruct_signal(shared_response, individual_tuning):
         """
-        Reconstructs the signal using the shared response and individual tuning.
+        Reconstructs the signal using the stimulus as shared
+        response and individual tuning.
 
         Parameters:
         --------
-        shared_response : numpy.ndarray
+        shared_response : ndarray
              The shared response of shape (n_timeframes, n_timeframes) or
              (n_timeframes, latent_dim).
-        individual_tuning : numpy.ndarray
+        individual_tuning : ndarray
              The individual tuning of shape (latent_dim, n_voxels) or
              (n_timeframes, n_voxels).
 
         Returns:
         --------
-        ndarray: The reconstructed signal of shape (n_timeframes, n_voxels)
-        (same shape as the original signal)
+        ndarray:
+            The reconstructed signal of shape (n_timeframes, n_voxels).
         """
         return (shared_response @ individual_tuning).astype(np.float32)
 
     def fit(
         self,
         X,
-        searchlights=None,
-        parcels=None,
-        dists=None,
-        radius=20,
-        tuning=True,
         verbose=True,
     ):
         """
@@ -624,23 +648,8 @@ def fit(
 
         Parameters:
         -----------
-
         X : array-like
             The training data of shape (n_subjects, n_samples, n_voxels).
-        searchlights : array-like
-            The searchlight indices for each subject,
-            of shape (n_s, n_searchlights).
-        parcels : array-like
-            The parcel indices for each subject,
-            of shape (n_s, n_parcels) (if not using searchlights)
-        dists : array-like
-            The distances of vertices to the center of their searchlight,
-            of shape (n_searchlights, n_vertices_sl)
-        radius : int(optional)
-            The radius of the searchlight sphere, in milimeters.
-            Defaults to 20.
-        tuning :bool(optional)
-            Whether to compute the tuning weights. Defaults to True.
         verbose : bool(optional)
             Whether to print progress information. Defaults to True.
 
@@ -658,17 +667,7 @@ def fit(
         self.tuning_data = np.empty(self.n_subjects, dtype=np.float32)
         self.denoised_signal = np.empty(self.n_subjects, dtype=np.float32)
 
-        if searchlights is None:
-            self.regions = parcels
-            self.distances = None
-            self.radius = None
-        else:
-            self.regions = searchlights
-            self.distances = dists
-            self.radius = radius
-
         denoiser = PiecewiseAlignment(
-            alignment_method=self.alignment_method,
             template_kind=self.decomp_method,
             n_jobs=self.n_jobs,
             verbose=verbose,
@@ -676,17 +675,16 @@ def fit(
         self.denoised_signal = denoiser.fit_transform(
             X_,
             regions=self.regions,
-            dists=self.distances,
+            dists=self.dists,
             radius=self.radius,
         )
 
         # Stimulus matrix computation
-
         full_signal = np.concatenate(self.denoised_signal, axis=1)
         self.shared_response = self._stimulus_estimator(
             full_signal, self.n_subjects, self.n_components
         )
-        if tuning:
+        if self.tuning:
             self.tuning_data = Parallel(n_jobs=self.n_jobs)(
                 delayed(self._tuning_estimator)(
                     self.shared_response,
@@ -710,7 +708,8 @@ def transform(self, X, verbose=False):
 
         Returns:
         --------
-        array-like: The transformed data of shape (n_subjects, n_timepoints, n_voxels).
+        ndarray :
+            The transformed data of shape (n_subjects, n_timepoints, n_voxels).
         """
 
         full_signal = np.concatenate(X, axis=1, dtype=np.float32)
diff --git a/fmralign/hyperalignment/piecewise_alignment.py b/fmralign/hyperalignment/piecewise_alignment.py
index 462b6ce..78b4cbb 100644
--- a/fmralign/hyperalignment/piecewise_alignment.py
+++ b/fmralign/hyperalignment/piecewise_alignment.py
@@ -26,7 +26,6 @@ class PiecewiseAlignment(BaseEstimator, TransformerMixin):
 
     def __init__(
         self,
-        alignment_method="searchlight_ridge",
         template_kind="pca",
         common_topography=True,
         verbose=True,
@@ -50,7 +49,6 @@ def __init__(
         self.n_s = None
         self.n_t = None
         self.n_v = None
-        self.warp_alignment_method = alignment_method
         self.template_kind = template_kind
         self.verbose = verbose
         self.common_topography = common_topography
diff --git a/fmralign/hyperalignment/test_hyperalignment.py b/fmralign/hyperalignment/test_hyperalignment.py
index 2961f8b..e0fe937 100644
--- a/fmralign/hyperalignment/test_hyperalignment.py
+++ b/fmralign/hyperalignment/test_hyperalignment.py
@@ -30,12 +30,10 @@ def test_int_fit_predict():
     dists = [np.ones((300,))]
 
     # Test INT on the two parts of the data (ie different runs of the experiment)
-    int1 = INT(n_components=6)
-    int2 = INT(n_components=6)
-    int1.fit(
-        X_train, searchlights=searchlights, dists=dists
-    )  # S is provided if we cheat and know the ground truth
-    int2.fit(X_test, searchlights=searchlights, dists=dists)
+    int1 = INT(n_components=6, searchlights=searchlights, dists=dists)
+    int2 = INT(n_components=6, searchlights=searchlights, dists=dists)
+    int1.fit(X_train)
+    int2.fit(X_test)
 
     X_pred = int1.transform(X_test)
     # save individual components
@@ -89,10 +87,10 @@ def test_int_with_searchlight():
     )
 
     # Test INT on the two parts of the data (ie different runs of the experiment)
-    model1 = INT(n_components=6)
-    model2 = INT(n_components=6)
-    model1.fit(X_train, searchlights=searchlights, dists=dists, radius=5)
-    model2.fit(X_test, searchlights=searchlights, dists=dists, radius=5)
+    model1 = INT(n_components=6, searchlights=searchlights, dists=dists, radius=5)
+    model2 = INT(n_components=6, searchlights=searchlights, dists=dists, radius=5)
+    model1.fit(X_train)
+    model2.fit(X_test)
     X_pred = model1.transform(X_test)
 
     tuning_data_run_1 = model1.tuning_data

From c5131dbbec1843c3656b789cfd9cae212de5584b Mon Sep 17 00:00:00 2001
From: FOUCHARD Denis <denis.fouchard@inria.fr>
Date: Wed, 28 Feb 2024 11:33:07 +0100
Subject: [PATCH 64/69] fix toy example with new API

---
 examples/plot_toy_int_experiment.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/examples/plot_toy_int_experiment.py b/examples/plot_toy_int_experiment.py
index af8c456..0fe9a05 100644
--- a/examples/plot_toy_int_experiment.py
+++ b/examples/plot_toy_int_experiment.py
@@ -78,19 +78,22 @@
 #############################################################################
 # Create two independant instances of the model
 # ---------------------------------------------
-# We create two instances of the INT model to align the two runs of the experiment.
-# We then extract the tuning matrices and the shared from the two runs to compare them.
+# We create two instances of the INT model to align the two runs of
+# the experiment, then extract the tuning matrices and the shared from the two
+# runs to compare them.
 
 int1 = INT(
     n_components=latent_dim,
+    parcels=parcels,
     decomp_method=decomposition_method,
 )
 int2 = INT(
     n_components=latent_dim,
+    parcels=parcels,
     decomp_method=decomposition_method,
 )
-int1.fit(data_run_1, parcels=parcels, verbose=False)
-int2.fit(data_run_2, parcels=parcels, verbose=False)
+int1.fit(data_run_1, verbose=False)
+int2.fit(data_run_2, verbose=False)
 
 # save individual components
 tuning_pred_run_1 = int1.tuning_data

From c2943f18a8b1c5ef767e16b92ed1d531b97ebaa0 Mon Sep 17 00:00:00 2001
From: DF <95576336+denisfouchard@users.noreply.github.com>
Date: Wed, 6 Mar 2024 16:58:42 +0100
Subject: [PATCH 65/69] Update examples/plot_int_alignment.py

Co-authored-by: Elizabeth DuPre <emd222@cornell.edu>
---
 examples/plot_int_alignment.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/plot_int_alignment.py b/examples/plot_int_alignment.py
index 96095c8..0fbc6e2 100644
--- a/examples/plot_int_alignment.py
+++ b/examples/plot_int_alignment.py
@@ -49,7 +49,7 @@
 # -----------------
 # We define a nilearn masker that will be used to handle relevant data.
 #   For more information, visit :
-#   'http://nilearn.github.io/manipulating_images/masker_objects.html'
+#   'https://nilearn.github.io/stable/manipulating_images/masker_objects.html'
 #
 
 from nilearn.maskers import NiftiMasker

From a7525b0f86ec8fb9dd36c8d1a1eca9ad35c8b90e Mon Sep 17 00:00:00 2001
From: DF <95576336+denisfouchard@users.noreply.github.com>
Date: Wed, 6 Mar 2024 16:59:50 +0100
Subject: [PATCH 66/69] Update examples/plot_int_alignment.py

Co-authored-by: Elizabeth DuPre <emd222@cornell.edu>
---
 examples/plot_int_alignment.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/plot_int_alignment.py b/examples/plot_int_alignment.py
index 0fbc6e2..2c92de3 100644
--- a/examples/plot_int_alignment.py
+++ b/examples/plot_int_alignment.py
@@ -199,7 +199,7 @@
 baseline_display = plotting.plot_stat_map(
     average_score, display_mode="z", vmax=1, cut_coords=[-15, -5]
 )
-baseline_display.title("Group average correlation wt ground truth")
+baseline_display.title("Group average correlation wrt ground truth")
 display = plotting.plot_stat_map(
     template_score, display_mode="z", cut_coords=[-15, -5], vmax=1
 )

From a19aefb1ef3e384090fcc8c29c7959e8a062fca7 Mon Sep 17 00:00:00 2001
From: DF <95576336+denisfouchard@users.noreply.github.com>
Date: Wed, 6 Mar 2024 17:00:02 +0100
Subject: [PATCH 67/69] Update examples/plot_toy_int_experiment.py

Co-authored-by: Elizabeth DuPre <emd222@cornell.edu>
---
 examples/plot_toy_int_experiment.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/plot_toy_int_experiment.py b/examples/plot_toy_int_experiment.py
index 0fe9a05..d9985ed 100644
--- a/examples/plot_toy_int_experiment.py
+++ b/examples/plot_toy_int_experiment.py
@@ -8,7 +8,7 @@
 
 This is a toy experiment to test Individual Tuning Model (INT) on two parts of the
 data (or different runs) to assess the validity of tuning computation. This code has
-no vocation to be an explanatory example, but rather a test to check the validity of
+no intention to be an explanatory example, but rather a test to check the validity of
 the INT model.
 
 

From 3a1ccc992c4f9ccb45c5158f5fcaedd04d15475a Mon Sep 17 00:00:00 2001
From: DF <95576336+denisfouchard@users.noreply.github.com>
Date: Wed, 6 Mar 2024 17:01:05 +0100
Subject: [PATCH 68/69] Update fmralign/hyperalignment/piecewise_alignment.py

Co-authored-by: Elizabeth DuPre <emd222@cornell.edu>
---
 fmralign/hyperalignment/piecewise_alignment.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/fmralign/hyperalignment/piecewise_alignment.py b/fmralign/hyperalignment/piecewise_alignment.py
index 78b4cbb..dd56002 100644
--- a/fmralign/hyperalignment/piecewise_alignment.py
+++ b/fmralign/hyperalignment/piecewise_alignment.py
@@ -1,8 +1,8 @@
 """Piecewise alignment model. This model decomposes the data into regions (pieces).
 Those can either be searchlights or parcels (computed with standard parcellation algorithms).
 See the ```nilearn``` documentation for more details:
-- https://nilearn.github.io/modules/generated/nilearn.regions.Parcellations.html
-- https://nilearn.github.io/dev/modules/generated/nilearn.decoding.SearchLight.html
+- https://nilearn.github.io/stable/modules/generated/nilearn.regions.Parcellations.html
+- https://nilearn.github.io/stable/modules/generated/nilearn.decoding.SearchLight.html
 """
 
 import numpy as np

From 5e4b904aabff4f4599f6741a02548c9e55577170 Mon Sep 17 00:00:00 2001
From: Denis Fouchard <dffouchard@gmail.com>
Date: Wed, 6 Mar 2024 17:11:43 +0100
Subject: [PATCH 69/69] adressing last PR comments

---
 examples/plot_toy_int_experiment.py           |   2 +-
 fmralign/fetch_example_data.py                | 151 +++++++++++++++++
 fmralign/generate_data.py                     | 153 ------------------
 .../hyperalignment/test_hyperalignment.py     |   5 +-
 4 files changed, 156 insertions(+), 155 deletions(-)
 delete mode 100644 fmralign/generate_data.py

diff --git a/examples/plot_toy_int_experiment.py b/examples/plot_toy_int_experiment.py
index d9985ed..ad06d57 100644
--- a/examples/plot_toy_int_experiment.py
+++ b/examples/plot_toy_int_experiment.py
@@ -24,7 +24,7 @@
 import numpy as np
 import matplotlib.pyplot as plt
 from fmralign.alignment_methods import IndividualizedNeuralTuning as INT
-from fmralign.generate_data import generate_dummy_signal
+from fmralign.fetch_example_data import generate_dummy_signal
 from fmralign.hyperalignment.correlation import (
     tuning_correlation,
     stimulus_correlation,
diff --git a/fmralign/fetch_example_data.py b/fmralign/fetch_example_data.py
index 67c6dfd..f43770e 100644
--- a/fmralign/fetch_example_data.py
+++ b/fmralign/fetch_example_data.py
@@ -3,6 +3,8 @@
 
 import pandas as pd
 from nilearn.datasets._utils import fetch_files, get_dataset_dir
+from fastsrm.srm import projection
+import numpy as np
 
 
 def fetch_ibc_subjects_contrasts(subjects, data_dir=None, verbose=1):
@@ -114,3 +116,152 @@ def fetch_ibc_subjects_contrasts(subjects, data_dir=None, verbose=1):
         )
         files.append(fetch_files(data_dir, filenames, verbose=verbose))
     return files, metadata_df, mask
+
+
+def generate_dummy_signal(
+    n_subjects: int,
+    n_timepoints: int,
+    n_voxels: int,
+    S_std=1,
+    latent_dim=None,
+    T_mean=0,
+    T_std=1,
+    SNR=1,
+    generative_method="custom",
+    seed=0,
+):
+    """Generate dummy signal for testing INT model
+
+    Parameters
+    ----------
+    n_subjects : int
+        Number of subjects.
+    n_timepoints : int
+        Number of timepoints.
+    n_voxels : int
+        Number of voxels.
+    S_std : float, default=1
+        Standard deviation of latent variables.
+    latent_dim: int, defult=None
+        Number of latent dimensions. Defualts to n_timepoints
+    T_mean : float
+        Mean of weights.
+    T_std : float
+        Standard deviation of weights.
+    SNR : float
+        Signal-to-noise ratio.
+    generative_method : str, default="custom"
+        Method for generating data. Options are "custom", "fastsrm".
+    seed : int
+        Random seed.
+
+
+    Returns
+    -------
+    imgs_train : ndarray of shape (n_subjects, n_timepoints, n_voxels)
+        Training data.
+    imgs_test : ndarray of shape (n_subjects, n_timepoints, n_voxels)
+        Testing data.
+    S_train : ndarray of shape (n_timepoints, latent_dim)
+        Training latent variables.
+    S_test : ndarray of shape (n_timepoints, latent_dim)
+        Testing latent variables.
+    Ts : ndarray of shape (n_subjects, latent_dim , n_voxels)
+        Tuning matrices.
+    """
+    if latent_dim is None:
+        latent_dim = n_timepoints
+
+    rng = np.random.RandomState(seed=seed)
+
+    if generative_method == "custom":
+        sigma = n_subjects * np.arange(1, latent_dim + 1)
+        # np.random.shuffle(sigma)
+        # Generate common signal matrix
+        S_train = S_std * np.random.randn(n_timepoints, latent_dim)
+        # Normalize each row to have unit norm
+        S_train = S_train / np.linalg.norm(S_train, axis=0, keepdims=True)
+        S_train = S_train @ np.diag(sigma)
+        S_test = S_std * np.random.randn(n_timepoints, latent_dim)
+        S_test = S_test / np.linalg.norm(S_test, axis=0, keepdims=True)
+        S_test = S_test @ np.diag(sigma)
+
+    elif generative_method == "fastsrm":
+        Sigma = rng.dirichlet(np.ones(latent_dim), 1).flatten()
+        S_train = np.sqrt(Sigma)[:, None] * rng.randn(n_timepoints, latent_dim)
+        S_test = np.sqrt(Sigma)[:, None] * rng.randn(n_timepoints, latent_dim)
+
+    elif generative_method == "multiviewica":
+        S_train = np.random.laplace(size=(n_timepoints, latent_dim))
+        S_test = np.random.laplace(size=(n_timepoints, latent_dim))
+
+    else:
+        raise ValueError("Unknown generative method")
+
+    # Generate indiivdual spatial components
+    data_train, data_test = [], []
+    Ts = []
+    for _ in range(n_subjects):
+        if generative_method == "custom" or generative_method == "multiviewica":
+            W = T_mean + T_std * np.random.randn(latent_dim, n_voxels)
+        else:
+            W = projection(rng.randn(latent_dim, n_voxels))
+
+        Ts.append(W)
+        X_train = S_train @ W
+        noise = np.random.randn(n_timepoints, n_voxels)
+        noise = (
+            noise
+            * np.linalg.norm(X_train)
+            / (SNR * np.linalg.norm(noise, axis=0, keepdims=True))
+        )
+        X_train += noise
+        data_train.append(X_train)
+        X_test = S_test @ W
+        noise = np.random.randn(n_timepoints, n_voxels)
+        noise = (
+            noise
+            * np.linalg.norm(X_test)
+            / (SNR * np.linalg.norm(noise, axis=0, keepdims=True))
+        )
+        X_test += noise
+        data_test.append(X_test)
+
+    data_train = np.array(data_train)
+    data_test = np.array(data_test)
+    return data_train, data_test, S_train, S_test, Ts
+
+
+def generate_dummy_searchlights(
+    n_searchlights: int,
+    n_voxels: int,
+    radius: float,
+    sl_size: int = 5,
+    seed: int = 0,
+):
+    """Generate dummy searchlights for testing INT model
+
+    Parameters
+    ----------
+    n_searchlights : int
+        Number of searchlights.
+    n_voxels : int
+        Number of voxels.
+    radius : float,
+        Radius of searchlights.
+    sl_size : int, default=5
+        Size of each searchlight (easier for dummy signal generation).
+    seed : int
+        Random seed.
+
+    Returns
+    -------
+    searchlights : ndarray of shape (n_searchlights, sl_size)
+        Searchlights.
+    dists : ndarray of shape (n_searchlights, sl_size)
+        Distances.
+    """
+    rng = np.random.RandomState(seed=seed)
+    searchlights = rng.randint(n_voxels, size=(n_searchlights, sl_size))
+    dists = rng.randint(radius, size=searchlights.shape)
+    return searchlights, dists
diff --git a/fmralign/generate_data.py b/fmralign/generate_data.py
deleted file mode 100644
index 2669558..0000000
--- a/fmralign/generate_data.py
+++ /dev/null
@@ -1,153 +0,0 @@
-"""Some functions to generate toy fMRI data using shared stimulus information."""
-
-import numpy as np
-from fastsrm.srm import projection
-
-
-def generate_dummy_signal(
-    n_subjects: int,
-    n_timepoints: int,
-    n_voxels: int,
-    S_std=1,
-    latent_dim=None,
-    T_mean=0,
-    T_std=1,
-    SNR=1,
-    generative_method="custom",
-    seed=0,
-):
-    """Generate dummy signal for testing INT model
-
-    Parameters
-    ----------
-    n_subjects : int
-        Number of subjects.
-    n_timepoints : int
-        Number of timepoints.
-    n_voxels : int
-        Number of voxels.
-    S_std : float, default=1
-        Standard deviation of latent variables.
-    latent_dim: int, defult=None
-        Number of latent dimensions. Defualts to n_timepoints
-    T_mean : float
-        Mean of weights.
-    T_std : float
-        Standard deviation of weights.
-    SNR : float
-        Signal-to-noise ratio.
-    generative_method : str, default="custom"
-        Method for generating data. Options are "custom", "fastsrm".
-    seed : int
-        Random seed.
-
-
-    Returns
-    -------
-    imgs_train : ndarray of shape (n_subjects, n_timepoints, n_voxels)
-        Training data.
-    imgs_test : ndarray of shape (n_subjects, n_timepoints, n_voxels)
-        Testing data.
-    S_train : ndarray of shape (n_timepoints, latent_dim)
-        Training latent variables.
-    S_test : ndarray of shape (n_timepoints, latent_dim)
-        Testing latent variables.
-    Ts : ndarray of shape (n_subjects, latent_dim , n_voxels)
-        Tuning matrices.
-    """
-    if latent_dim is None:
-        latent_dim = n_timepoints
-
-    rng = np.random.RandomState(seed=seed)
-
-    if generative_method == "custom":
-        sigma = n_subjects * np.arange(1, latent_dim + 1)
-        # np.random.shuffle(sigma)
-        # Generate common signal matrix
-        S_train = S_std * np.random.randn(n_timepoints, latent_dim)
-        # Normalize each row to have unit norm
-        S_train = S_train / np.linalg.norm(S_train, axis=0, keepdims=True)
-        S_train = S_train @ np.diag(sigma)
-        S_test = S_std * np.random.randn(n_timepoints, latent_dim)
-        S_test = S_test / np.linalg.norm(S_test, axis=0, keepdims=True)
-        S_test = S_test @ np.diag(sigma)
-
-    elif generative_method == "fastsrm":
-        Sigma = rng.dirichlet(np.ones(latent_dim), 1).flatten()
-        S_train = np.sqrt(Sigma)[:, None] * rng.randn(n_timepoints, latent_dim)
-        S_test = np.sqrt(Sigma)[:, None] * rng.randn(n_timepoints, latent_dim)
-
-    elif generative_method == "multiviewica":
-        S_train = np.random.laplace(size=(n_timepoints, latent_dim))
-        S_test = np.random.laplace(size=(n_timepoints, latent_dim))
-
-    else:
-        raise ValueError("Unknown generative method")
-
-    # Generate indiivdual spatial components
-    data_train, data_test = [], []
-    Ts = []
-    for _ in range(n_subjects):
-        if generative_method == "custom" or generative_method == "multiviewica":
-            W = T_mean + T_std * np.random.randn(latent_dim, n_voxels)
-        else:
-            W = projection(rng.randn(latent_dim, n_voxels))
-
-        Ts.append(W)
-        X_train = S_train @ W
-        noise = np.random.randn(n_timepoints, n_voxels)
-        noise = (
-            noise
-            * np.linalg.norm(X_train)
-            / (SNR * np.linalg.norm(noise, axis=0, keepdims=True))
-        )
-        X_train += noise
-        data_train.append(X_train)
-        X_test = S_test @ W
-        noise = np.random.randn(n_timepoints, n_voxels)
-        noise = (
-            noise
-            * np.linalg.norm(X_test)
-            / (SNR * np.linalg.norm(noise, axis=0, keepdims=True))
-        )
-        X_test += noise
-        data_test.append(X_test)
-
-    data_train = np.array(data_train)
-    data_test = np.array(data_test)
-    return data_train, data_test, S_train, S_test, Ts
-
-
-def generate_dummy_searchlights(
-    n_searchlights: int,
-    n_voxels: int,
-    radius: float,
-    sl_size: int = 5,
-    seed: int = 0,
-):
-    """Generate dummy searchlights for testing INT model
-
-    Parameters
-    ----------
-    n_searchlights : int
-        Number of searchlights.
-    n_voxels : int
-        Number of voxels.
-    radius : float,
-        Radius of searchlights.
-    sl_size : int, default=5
-        Size of each searchlight (easier for dummy signal generation).
-    seed : int
-        Random seed.
-
-    Returns
-    -------
-    searchlights : ndarray of shape (n_searchlights, sl_size)
-        Searchlights.
-    dists : ndarray of shape (n_searchlights, sl_size)
-        Distances.
-    """
-    rng = np.random.RandomState(seed=seed)
-    searchlights = rng.randint(n_voxels, size=(n_searchlights, sl_size))
-    dists = rng.randint(radius, size=searchlights.shape)
-    return searchlights, dists
diff --git a/fmralign/hyperalignment/test_hyperalignment.py b/fmralign/hyperalignment/test_hyperalignment.py
index e0fe937..9100381 100644
--- a/fmralign/hyperalignment/test_hyperalignment.py
+++ b/fmralign/hyperalignment/test_hyperalignment.py
@@ -1,5 +1,8 @@
 from fmralign.alignment_methods import IndividualizedNeuralTuning as INT
-from fmralign.generate_data import generate_dummy_signal, generate_dummy_searchlights
+from fmralign.fetch_example_data import (
+    generate_dummy_signal,
+    generate_dummy_searchlights,
+)
 import numpy as np
 
 from fmralign.hyperalignment.correlation import (