class FastSRM(BaseEstimator, TransformerMixin):
    """SRM decomposition using a very low amount of memory and \
computational power thanks to the use of an atlas \
as described in [Richard2019]_.

    Given multi-subject data, factorize it as a shared response S \
among all subjects and an orthogonal transform (basis) W per subject:

    .. math:: X_i \\approx W_i S, \\forall i=1 \\dots N

    Parameters
    ----------

    atlas :  array, shape=[n_supervoxels, n_voxels] or array,\
shape=[n_voxels] or str or None, default=None
        Probabilistic or deterministic atlas on which to project the data. \
Deterministic atlas is an array of shape [n_voxels,] \
where values range from 1 \
to n_supervoxels. Voxels labelled 0 will be ignored. If atlas is a str the \
corresponding array is loaded with numpy.load and expected shape \
is (n_voxels,) for a deterministic atlas and \
(n_supervoxels, n_voxels) for a probabilistic atlas.

    n_components : int
        Number of timecourses of the shared coordinates

    n_iter : int
        Number of iterations to perform

    temp_dir : str or None
        Path to dir where temporary results are stored. If None \
temporary results will be stored in memory. This can results in memory \
errors when the number of subjects and/or sessions is large

    low_ram : bool
        If True and temp_dir is not None, reduced_data will be saved on \
disk. This increases the number of IO but reduces memory complexity when \
the number of subject and/or sessions is large

    seed : int
        Seed used for random sampling.

    n_jobs : int, optional, default=1
        The number of CPUs to use to do the computation. \
-1 means all CPUs, -2 all CPUs but one, and so on.

    verbose : bool or "warn"
        If True, logs are enabled. If False, logs are disabled. \
If "warn" only warnings are printed.

    aggregate: str or None, default="mean"
        If "mean", shared_response is the mean shared response \
from all subjects. If None, shared_response contains all \
subject-specific responses in shared space

    Attributes
    ----------

    `basis_list`: list of array, element i has \
shape=[n_components, n_voxels] or list of str
        - if basis is a list of array, element i is the basis of subject i
        - if basis is a list of str, element i is the path to the basis \
of subject i that is loaded with np.load yielding an array of \
shape [n_components, n_voxels].

        Note that any call to the clean method erases this attribute

    Note
    -----

        **References:**
        H. Richard, L. Martin, A. Pinho, J. Pillow, B. Thirion, 2019: \
Fast shared response model for fMRI data (https://arxiv.org/pdf/1909.12537.pdf)

    """
    def __init__(self,
                 atlas=None,
                 n_components=20,
                 n_iter=100,
                 temp_dir=None,
                 low_ram=False,
                 seed=None,
                 n_jobs=1,
                 verbose="warn",
                 aggregate="mean"):

        self.seed = seed
        self.n_jobs = n_jobs
        self.verbose = verbose
        self.n_components = n_components
        self.n_iter = n_iter
        self.atlas = atlas

        if aggregate is not None and aggregate != "mean":
            raise ValueError("aggregate can have only value mean or None")

        self.aggregate = aggregate

        self.basis_list = None

        if temp_dir is None:
            if self.verbose == "warn" or self.verbose is True:
                logger.warning("temp_dir has value None. "
                               "All basis (spatial maps) and reconstructed "
                               "data will therefore be kept in memory."
                               "This can lead to memory errors when the "
                               "number of subjects "
                               "and/or sessions is large.")
            self.temp_dir = None
            self.low_ram = False

        if temp_dir is not None:
            self.temp_dir = os.path.join(temp_dir,
                                         "fastsrm" + str(uuid.uuid4()))
            self.low_ram = low_ram

    def clean(self):
        """This erases temporary files and basis_list attribute to \
free memory. This method should be called when fitted model \
is not needed anymore.
        """
        if self.temp_dir is not None:
            if os.path.exists(self.temp_dir):
                for root, dirs, files in os.walk(self.temp_dir, topdown=False):
                    for name in files:
                        os.remove(os.path.join(root, name))
                os.rmdir(self.temp_dir)

        if self.basis_list is not None:
            self.basis_list is None

    def fit(self, imgs):
        """Computes basis across subjects from input imgs

        Parameters
        ----------
        imgs : array of str, shape=[n_subjects, n_sessions] or \
list of list of arrays or list of arrays
            Element i, j of the array is a path to the data of subject i \
collected during session j. Data are loaded with numpy.load and expected \
shape is [n_voxels, n_timeframes] n_timeframes and n_voxels are assumed \
to be the same across subjects n_timeframes can vary across sessions. \
Each voxel's timecourse is assumed to have mean 0 and variance 1

            imgs can also be a list of list of arrays where element i, j \
of the array is a numpy array of shape [n_voxels, n_timeframes] \
that contains the data of subject i collected during session j.

            imgs can also be a list of arrays where element i \
of the array is a numpy array of shape [n_voxels, n_timeframes] \
that contains the data of subject i (number of sessions is implicitly 1)

        Returns
        -------
        self : object
           Returns the instance itself. Contains attributes listed \
at the object level.
        """
        atlas_shape = check_atlas(self.atlas, self.n_components)
        reshaped_input, imgs, shapes = check_imgs(
            imgs, n_components=self.n_components, atlas_shape=atlas_shape)
        self.clean()
        create_temp_dir(self.temp_dir)

        if self.verbose is True:
            logger.info("[FastSRM.fit] Reducing data")

        reduced_data = reduce_data(imgs,
                                   atlas=self.atlas,
                                   n_jobs=self.n_jobs,
                                   low_ram=self.low_ram,
                                   temp_dir=self.temp_dir)

        if self.verbose is True:
            logger.info("[FastSRM.fit] Finds shared "
                        "response using reduced data")
        shared_response_list = fast_srm(reduced_data,
                                        n_iter=self.n_iter,
                                        n_components=self.n_components,
                                        low_ram=self.low_ram,
                                        seed=self.seed)
        if self.verbose is True:
            logger.info("[FastSRM.fit] Finds basis using "
                        "full data and shared response")

        if self.n_jobs == 1:
            basis = []
            for i, sessions in enumerate(imgs):
                basis_i = _compute_basis_subject_online(
                    sessions, shared_response_list)
                if self.temp_dir is None:
                    basis.append(basis_i)
                else:
                    path = os.path.join(self.temp_dir, "basis_%i" % i)
                    np.save(path, basis_i)
                    basis.append(path + ".npy")
                del basis_i
        else:
            if self.temp_dir is None:
                basis = Parallel(n_jobs=self.n_jobs)(
                    delayed(_compute_basis_subject_online)(
                        sessions, shared_response_list) for sessions in imgs)
            else:
                Parallel(n_jobs=self.n_jobs)(
                    delayed(_compute_and_save_corr_mat)(
                        imgs[i][j], shared_response_list[j], self.temp_dir)
                    for j in range(len(imgs[0])) for i in range(len(imgs)))

                basis = Parallel(n_jobs=self.n_jobs)(
                    delayed(_compute_and_save_subject_basis)(i, sessions,
                                                             self.temp_dir)
                    for i, sessions in enumerate(imgs))

        self.basis_list = basis
        return self

    def fit_transform(self, imgs, subjects_indexes=None):
        """Computes basis across subjects and shared response from input imgs
        return shared response.

        Parameters
        ----------
        imgs : array of str, shape=[n_subjects, n_sessions] or \
list of list of arrays or list of arrays
            Element i, j of the array is a path to the data of subject i \
collected during session j. Data are loaded with numpy.load and expected \
shape is [n_voxels, n_timeframes] n_timeframes and n_voxels are assumed \
to be the same across subjects n_timeframes can vary across sessions. \
Each voxel's timecourse is assumed to have mean 0 and variance 1

            imgs can also be a list of list of arrays where element i, j \
of the array is a numpy array of shape [n_voxels, n_timeframes] \
that contains the data of subject i collected during session j.

            imgs can also be a list of arrays where element i \
of the array is a numpy array of shape [n_voxels, n_timeframes] \
that contains the data of subject i (number of sessions is implicitly 1)

        subjects_indexes : list or None:
            if None imgs[i] will be transformed using basis_list[i]. \
Otherwise imgs[i] will be transformed using basis_list[subjects_index[i]]

        Returns
        --------
        shared_response : list of arrays, list of list of arrays or array
            - if imgs is a list of array and self.aggregate="mean": shared \
response is an array of shape (n_components, n_timeframes)
            - if imgs is a list of array and self.aggregate=None: shared \
response is a list of array, element i is the projection of data of \
subject i in shared space.
            - if imgs is an array or a list of list of array and \
self.aggregate="mean": shared response is a list of array, \
element j is the shared response during session j
            - if imgs is an array or a list of list of array and \
self.aggregate=None: shared response is a list of list of array, \
element i, j is the projection of data of subject i collected \
during session j in shared space.
        """
        self.fit(imgs)
        return self.transform(imgs, subjects_indexes=subjects_indexes)

    def transform(self, imgs, subjects_indexes=None):
        """From data in imgs and basis from training data,
        computes shared response.

        Parameters
        ----------
        imgs : array of str, shape=[n_subjects, n_sessions] or \
list of list of arrays or list of arrays
            Element i, j of the array is a path to the data of subject i \
collected during session j. Data are loaded with numpy.load and expected \
shape is [n_voxels, n_timeframes] n_timeframes and n_voxels are assumed \
to be the same across subjects n_timeframes can vary across sessions. \
Each voxel's timecourse is assumed to have mean 0 and variance 1

            imgs can also be a list of list of arrays where element i, j \
of the array is a numpy array of shape [n_voxels, n_timeframes] \
that contains the data of subject i collected during session j.

            imgs can also be a list of arrays where element i \
of the array is a numpy array of shape [n_voxels, n_timeframes] \
that contains the data of subject i (number of sessions is implicitly 1)

        subjects_indexes : list or None:
            if None imgs[i] will be transformed using basis_list[i]. \
Otherwise imgs[i] will be transformed using basis[subjects_index[i]]

        Returns
        --------
        shared_response : list of arrays, list of list of arrays or array
            - if imgs is a list of array and self.aggregate="mean": shared \
response is an array of shape (n_components, n_timeframes)
            - if imgs is a list of array and self.aggregate=None: shared \
response is a list of array, element i is the projection of data of \
subject i in shared space.
            - if imgs is an array or a list of list of array and \
self.aggregate="mean": shared response is a list of array, \
element j is the shared response during session j
            - if imgs is an array or a list of list of array and \
self.aggregate=None: shared response is a list of list of array, \
element i, j is the projection of data of subject i collected \
during session j in shared space.
         """
        aggregate = self.aggregate
        if self.basis_list is None:
            raise NotFittedError("The model fit has not been run yet.")

        atlas_shape = check_atlas(self.atlas, self.n_components)
        reshaped_input, imgs, shapes = check_imgs(
            imgs,
            n_components=self.n_components,
            atlas_shape=atlas_shape,
            ignore_nsubjects=True)
        check_indexes(subjects_indexes, "subjects_indexes")
        if subjects_indexes is None:
            subjects_indexes = np.arange(len(imgs))
        else:
            subjects_indexes = np.array(subjects_indexes)

        # Transform specific checks
        if len(subjects_indexes) < len(imgs):
            raise ValueError("Input data imgs has len %i whereas "
                             "subject_indexes has len %i. "
                             "The number of basis used to compute "
                             "the shared response should be equal "
                             "to the number of subjects in imgs" %
                             (len(imgs), len(subjects_indexes)))

        assert_valid_index(subjects_indexes, len(self.basis_list),
                           "subjects_indexes")

        shared_response = _compute_shared_response_online(
            imgs, self.basis_list, self.temp_dir, self.n_jobs,
            subjects_indexes, aggregate)

        # If shared response has only 1 session we need to reshape it
        if reshaped_input:
            if aggregate == "mean":
                shared_response = shared_response[0]
            if aggregate is None:
                shared_response = [
                    shared_response[i][0] for i in range(len(subjects_indexes))
                ]

        return shared_response

    def inverse_transform(
            self,
            shared_response,
            subjects_indexes=None,
            sessions_indexes=None,
    ):
        """From shared response and basis from training data
        reconstruct subject's data

        Parameters
        ----------

        shared_response : list of arrays, list of list of arrays or array
            - if imgs is a list of array and self.aggregate="mean": shared \
response is an array of shape (n_components, n_timeframes)
            - if imgs is a list of array and self.aggregate=None: shared \
response is a list of array, element i is the projection of data of \
subject i in shared space.
            - if imgs is an array or a list of list of array and \
self.aggregate="mean": shared response is a list of array, \
element j is the shared response during session j
            - if imgs is an array or a list of list of array and \
self.aggregate=None: shared response is a list of list of array, \
element i, j is the projection of data of subject i collected \
during session j in shared space.

        subjects_indexes : list or None
            if None reconstructs data of all subjects used during train. \
Otherwise reconstructs data of subjects specified by subjects_indexes.

        sessions_indexes : list or None
            if None reconstructs data of all sessions. \
Otherwise uses reconstructs data of sessions specified by sessions_indexes.

        Returns
        -------
        reconstructed_data: list of list of arrays or list of arrays
            - if reconstructed_data is a list of list : element i, j is \
the reconstructed data for subject subjects_indexes[i] and \
session sessions_indexes[j] as an np array of shape n_voxels, \
n_timeframes
            - if reconstructed_data is a list : element i is the \
reconstructed data for subject \
subject_indexes[i] as an np array of shape n_voxels, n_timeframes

        """
        added_session, shared = check_shared_response(
            shared_response, self.aggregate, n_components=self.n_components)
        n_subjects = len(self.basis_list)
        n_sessions = len(shared)

        for j in range(n_sessions):
            assert_array_2axis(shared[j], "shared_response[%i]" % j)

        check_indexes(subjects_indexes, "subjects_indexes")
        check_indexes(sessions_indexes, "sessions_indexes")

        if subjects_indexes is None:
            subjects_indexes = np.arange(n_subjects)
        else:
            subjects_indexes = np.array(subjects_indexes)

        assert_valid_index(subjects_indexes, n_subjects, "subjects_indexes")

        if sessions_indexes is None:
            sessions_indexes = np.arange(len(shared))
        else:
            sessions_indexes = np.array(sessions_indexes)

        assert_valid_index(sessions_indexes, n_sessions, "sessions_indexes")

        data = []
        for i in subjects_indexes:
            data_ = []
            basis_i = safe_load(self.basis_list[i])
            if added_session:
                data.append(basis_i.T.dot(shared[0]))
            else:
                for j in sessions_indexes:
                    data_.append(basis_i.T.dot(shared[j]))
                data.append(data_)
        return data

    def add_subjects(self, imgs, shared_response):
        """ Add subjects to the current fit. Each new basis will be \
appended at the end of the list of basis (which can \
be accessed using self.basis)

        Parameters
        ----------

        imgs : array of str, shape=[n_subjects, n_sessions] or \
list of list of arrays or list of arrays
            Element i, j of the array is a path to the data of subject i \
collected during session j. Data are loaded with numpy.load and expected \
shape is [n_voxels, n_timeframes] n_timeframes and n_voxels are assumed \
to be the same across subjects n_timeframes can vary across sessions. \
Each voxel's timecourse is assumed to have mean 0 and variance 1

            imgs can also be a list of list of arrays where element i, j \
of the array is a numpy array of shape [n_voxels, n_timeframes] \
that contains the data of subject i collected during session j.

            imgs can also be a list of arrays where element i \
of the array is a numpy array of shape [n_voxels, n_timeframes] \
that contains the data of subject i (number of sessions is implicitly 1)

        shared_response : list of arrays, list of list of arrays or array
            - if imgs is a list of array and self.aggregate="mean": shared \
response is an array of shape (n_components, n_timeframes)
            - if imgs is a list of array and self.aggregate=None: shared \
response is a list of array, element i is the projection of data of \
subject i in shared space.
            - if imgs is an array or a list of list of array and \
self.aggregate="mean": shared response is a list of array, \
element j is the shared response during session j
            - if imgs is an array or a list of list of array and \
self.aggregate=None: shared response is a list of list of array, \
element i, j is the projection of data of subject i collected \
during session j in shared space.
        """
        atlas_shape = check_atlas(self.atlas, self.n_components)
        reshaped_input, imgs, shapes = check_imgs(
            imgs,
            n_components=self.n_components,
            atlas_shape=atlas_shape,
            ignore_nsubjects=True)

        _, shared_response_list = check_shared_response(
            shared_response,
            n_components=self.n_components,
            aggregate=self.aggregate,
            input_shapes=shapes)

        # we need to transpose shared_response_list to be consistent with
        # other functions
        shared_response_list = [
            shared_response_list[j].T for j in range(len(shared_response_list))
        ]

        if self.n_jobs == 1:
            basis = []
            for i, sessions in enumerate(imgs):
                basis_i = _compute_basis_subject_online(
                    sessions, shared_response_list)
                if self.temp_dir is None:
                    basis.append(basis_i)
                else:
                    path = os.path.join(
                        self.temp_dir, "basis_%i" % (len(self.basis_list) + i))
                    np.save(path, basis_i)
                    basis.append(path + ".npy")
                del basis_i
        else:
            if self.temp_dir is None:
                basis = Parallel(n_jobs=self.n_jobs)(
                    delayed(_compute_basis_subject_online)(
                        sessions, shared_response_list) for sessions in imgs)
            else:
                Parallel(n_jobs=self.n_jobs)(
                    delayed(_compute_and_save_corr_mat)(
                        imgs[i][j], shared_response_list[j], self.temp_dir)
                    for j in range(len(imgs[0])) for i in range(len(imgs)))

                basis = Parallel(n_jobs=self.n_jobs)(
                    delayed(_compute_and_save_subject_basis)(
                        len(self.basis_list) + i, sessions, self.temp_dir)
                    for i, sessions in enumerate(imgs))

        self.basis_list += basis def _compute_and_save_corr_mat(img, shared_response, temp_dir):
    """computes correlation matrix and stores it

    Parameters
    ----------
    img : str
        path to data.
        Data are loaded with numpy.load and expected shape is
        [n_timeframes, n_voxels]
        n_timeframes and n_voxels are assumed to be the same across subjects
        n_timeframes can vary across sessions
        Each voxel's timecourse is assumed to have mean 0 and variance 1

    shared_response : array, shape=[n_timeframes, n_components]
        shared response
    """
    data = safe_load(img).T
    name = safe_encode(img)
    path = os.path.join(temp_dir, "corr_mat_" + name)
    np.save(path, shared_response.T.dot(data)) def _compute_and_save_subject_basis(subject_number, sessions, temp_dir):
    """computes correlation matrix for all sessions

    Parameters
    ----------

    subject_number: int
        Number that identifies the subject. Basis will be stored in
         [temp_dir]/basis_[subject_number].npy

    sessions : array of str
        Element i of the array is a path to the data collected during
        session i.
        Data are loaded with numpy.load and expected shape is
        [n_timeframes, n_voxels]
        n_timeframes and n_voxels are assumed to be the same across subjects
        n_timeframes can vary across sessions
        Each voxel's timecourse is assumed to have mean 0 and variance

    temp_dir : str or None
        path to dir where temporary results are stored
        if None temporary results will be stored in memory. This
        can results in memory errors when the number of subjects
        and / or sessions is large

    Returns
    -------

    basis: array, shape=[n_component, n_voxels] or str
        basis of subject [subject_number] or path to this basis
    """
    corr_mat = None
    for session in sessions:
        name = safe_encode(session)
        path = os.path.join(temp_dir, "corr_mat_" + name + ".npy")
        if corr_mat is None:
            corr_mat = np.load(path)
        else:
            corr_mat += np.load(path)
        os.remove(path)
    basis_i = _compute_subject_basis(corr_mat)
    path = os.path.join(temp_dir, "basis_%i" % subject_number)
    np.save(path, basis_i)
    return path + ".npy" def _compute_basis_subject_online(sessions, shared_response_list):
    """Computes subject's basis with shared response fixed

    Parameters
    ----------

    sessions : array of str
        Element i of the array is a path to the data
        collected during session i.
        Data are loaded with numpy.load and expected shape is
         [n_timeframes, n_voxels]
        n_timeframes and n_voxels are assumed to be the same across subjects
        n_timeframes can vary across sessions
        Each voxel's timecourse is assumed to have mean 0 and variance 1

    shared_response_list : list of array, element i has
    shape=[n_timeframes, n_components]
        shared response, element i is the shared response during session i

    Returns
    -------

    basis: array, shape=[n_components, n_voxels]
        basis
    """

    basis_i = None
    i = 0
    for session in sessions:
        data = safe_load(session).T
        if basis_i is None:
            basis_i = shared_response_list[i].T.dot(data)
        else:
            basis_i += shared_response_list[i].T.dot(data)
        i += 1
        del data
    return _compute_subject_basis(basis_i) def _reduced_space_compute_shared_response(reduced_data_list,
                                           reduced_basis_list,
                                           n_components=50):
    """Compute shared response with basis fixed in reduced space

    Parameters
    ----------

    reduced_data_list : array of str, shape=[n_subjects, n_sessions]
    or array, shape=[n_subjects, n_sessions, n_timeframes, n_supervoxels]
        Element i, j of the array is a path to the data of subject i
        collected during session j.
        Data are loaded with numpy.load and expected shape is
        [n_timeframes, n_supervoxels]
        or Element i, j of the array is the data in array of
        shape=[n_timeframes, n_supervoxels]
        n_timeframes and n_supervoxels are
        assumed to be the same across subjects
        n_timeframes can vary across sessions
        Each voxel's timecourse is assumed to have mean 0 and variance 1

    reduced_basis_list : None or list of array, element i has
    shape=[n_components, n_supervoxels]
        each subject's reduced basis
        if None the basis will be generated on the fly

    n_components : int or None
        number of components

    Returns
    -------

    shared_response_list : list of array, element i has
    shape=[n_timeframes, n_components]
        shared response, element i is the shared response during session i

    """
    n_subjects, n_sessions = reduced_data_list.shape[:2]

    s = [None] * n_sessions

    # This is just to check that all subjects have same number of
    # timeframes in a given session
    for n in range(n_subjects):
        for m in range(n_sessions):
            data_nm = safe_load(reduced_data_list[n][m])
            n_timeframes, n_supervoxels = data_nm.shape

            if reduced_basis_list is None:
                reduced_basis_list = []
                for subject in range(n_subjects):
                    q = np.eye(n_components, n_supervoxels)
                    reduced_basis_list.append(q)

            basis_n = reduced_basis_list[n]
            if s[m] is None:
                s[m] = data_nm.dot(basis_n.T)
            else:
                s[m] = s[m] + data_nm.dot(basis_n.T)

    for m in range(n_sessions):
        s[m] = s[m] / float(n_subjects)

    return s def check_atlas(atlas, n_components=None):
    """ Check input atlas

    Parameters
    ----------
    atlas :  array, shape=[n_supervoxels, n_voxels] or array, shape=[n_voxels]
     or str or None
        Probabilistic or deterministic atlas on which to project the data
        Deterministic atlas is an array of shape [n_voxels,] where values
        range from 1 to n_supervoxels. Voxels labelled 0 will be ignored.
        If atlas is a str the corresponding array is loaded with numpy.load
        and expected shape is (n_voxels,) for a deterministic atlas and
        (n_supervoxels, n_voxels) for a probabilistic atlas.

    n_components : int
        Number of timecourses of the shared coordinates

    Returns
    -------
    shape : array or None
        atlas shape
    """
    if atlas is None:
        return None

    if not (isinstance(atlas, np.ndarray) or isinstance(atlas, str)
            or isinstance(atlas, np.str_) or isinstance(atlas, np.str)):
        raise ValueError("Atlas is stored using "
                         "type %s which is neither np.ndarray or str" %
                         type(atlas))

    if isinstance(atlas, np.ndarray):
        shape = atlas.shape
    else:
        shape = get_shape(atlas)

    if len(shape) == 1:
        # We have a deterministic atlas
        atlas_array = safe_load(atlas)
        n_voxels = atlas_array.shape[0]
        n_supervoxels = len(np.unique(atlas_array)) - 1
        shape = (n_supervoxels, n_voxels)
    elif len(shape) != 2:
        raise ValueError(
            "Atlas has %i axes. It should have either 1 or 2 axes." %
            len(shape))

    n_supervoxels, n_voxels = shape

    if n_supervoxels > n_voxels:
        raise ValueError("Number of regions in the atlas is bigger than "
                         "the number of voxels (%i > %i)" %
                         (n_supervoxels, n_voxels))

    if n_components is not None:
        if n_supervoxels < n_components:
            raise ValueError("Number of regions in the atlas is "
                             "lower than the number of components "
                             "(%i < %i)" % (n_supervoxels, n_components))
    return shape def check_imgs(imgs,
               n_components=None,
               atlas_shape=None,
               ignore_nsubjects=False):
    """
    Check input images

    Parameters
    ----------

    imgs : array of str, shape=[n_subjects, n_sessions]
            Element i, j of the array is a path to the data of subject i
            collected during session j.
            Data are loaded with numpy.load and expected
            shape is [n_voxels, n_timeframes]
            n_timeframes and n_voxels are assumed to be the same across
            subjects
            n_timeframes can vary across sessions
            Each voxel's timecourse is assumed to have mean 0 and variance 1

            imgs can also be a list of list of arrays where element i, j of
            the array is a numpy array of shape [n_voxels, n_timeframes] that
            contains the data of subject i collected during session j.

            imgs can also be a list of arrays where element i of the array is
            a numpy array of shape [n_voxels, n_timeframes] that contains the
            data of subject i (number of sessions is implicitly 1)

    Returns
    -------
    reshaped_input: bool
        True if input had to be reshaped to match the
        n_subjects, n_sessions input
    new_imgs: list of list of array or np array
        input imgs reshaped if it is a list of arrays so that it becomes a
        list of list of arrays
    shapes: array
        Shape of input images
    """
    reshaped_input = False
    new_imgs = imgs
    if isinstance(imgs, list):
        _check_imgs_list(imgs)
        if isinstance(imgs[0], list):
            shapes = _check_imgs_list_list(imgs)
        elif isinstance(imgs[0], np.ndarray):
            new_imgs, shapes = _check_imgs_list_array(imgs)
            reshaped_input = True
        else:
            raise ValueError(
                "Since imgs is a list, it should be a list of list "
                "of arrays or a list of arrays but imgs[0] has type %s" %
                type(imgs[0]))
    elif isinstance(imgs, np.ndarray):
        shapes = _check_imgs_array(imgs)
    else:
        raise ValueError(
            "Input imgs should either be a list or an array but has type %s" %
            type(imgs))

    _check_shapes(shapes, n_components, atlas_shape, ignore_nsubjects)

    return reshaped_input, new_imgs, shapes def check_shared_response(shared_response,
                          aggregate="mean",
                          n_components=None,
                          input_shapes=None):
    """
    Check that shared response has valid input and turn it into
    a session-wise shared response

    Returns
    -------
    added_session: bool
        True if an artificial sessions was added to match the list of
        session input type for shared_response
    reshaped_shared_response: list of arrays
        shared response (reshaped to match the list of session input)
    """
    # Depending on aggregate and shape of input we infer what to do
    if isinstance(shared_response, list):
        assert_non_empty_list(shared_response, "shared_response")
        if isinstance(shared_response[0], list):
            if aggregate == "mean":
                raise ValueError("self.aggregate has value 'mean' but "
                                 "shared response is a list of list. This is "
                                 "incompatible")
            return False, _check_shared_response_list_of_list(
                shared_response, n_components, input_shapes)
        elif isinstance(shared_response[0], np.ndarray):
            if aggregate == "mean":
                return False, _check_shared_response_list_sessions(
                    shared_response, n_components, input_shapes)
            else:
                return True, _check_shared_response_list_subjects(
                    shared_response, n_components, input_shapes)
        else:
            raise ValueError("shared_response is a list but "
                             "shared_response[0] is neither a list "
                             "or an array. This is invalid.")
    elif isinstance(shared_response, np.ndarray):
        return True, _check_shared_response_array(shared_response,
                                                  n_components, input_shapes)
    else:
        raise ValueError("shared_response should be either "
                         "a list or an array but is of type %s" %
                         type(shared_response)) def create_temp_dir(temp_dir):
    """
    This check whether temp_dir exists and creates dir otherwise
    """
    if temp_dir is None:
        return None

    if not os.path.exists(temp_dir):
        os.makedirs(temp_dir)
    else:
        raise ValueError("Path %s already exists. "
                         "When a model is used, filesystem should be cleaned "
                         "by using the .clean() method" % temp_dir) def fast_srm(reduced_data_list,
             n_iter=10,
             n_components=None,
             low_ram=False,
             seed=0):
    """Computes shared response and basis in reduced space

    Parameters
    ----------

    reduced_data_list : array, shape=[n_subjects, n_sessions]
    or array, shape=[n_subjects, n_sessions, n_timeframes, n_supervoxels]
        Element i, j of the array is a path to the data of subject i
        collected during session j.
        Data are loaded with numpy.load and expected
        shape is [n_timeframes, n_supervoxels]
        or Element i, j of the array is the data in array of
        shape=[n_timeframes, n_supervoxels]
        n_timeframes and n_supervoxels are
         assumed to be the same across subjects
        n_timeframes can vary across sessions
        Each voxel's timecourse is assumed to have mean 0 and variance 1

    n_iter : int
        Number of iterations performed

    n_components : int or None
        number of components

    Returns
    -------

    shared_response_list : list of array, element i has
     shape=[n_timeframes, n_components]
        shared response, element i is the shared response during session i
    """
    if low_ram:
        return lowram_srm(reduced_data_list, n_iter, n_components)
    else:
        # We need to switch data to DetSRM format
        # Indeed in DetSRM all sessions are concatenated.
        # Whereas in FastSRM multiple sessions are supported.

        n_subjects, n_sessions = reduced_data_list.shape[:2]
        # We store the correspondence between timeframes and session
        timeframes_slices = []
        current_j = 0
        for j in range(n_sessions):
            timeframes_slices.append(
                slice(current_j, current_j + len(reduced_data_list[0, j])))
            current_j += len(reduced_data_list[0][j])
        # Now we can concatenate everything
        X = [
            np.concatenate(reduced_data_list[i], axis=0).T
            for i in range(n_subjects)
        ]

        srm = DetSRM(n_iter=n_iter, features=n_components, rand_seed=seed)
        srm.fit(X)

        # SRM gives a list of data projected in shared space
        # we get the shared response by averaging those
        concatenated_s = np.mean(srm.transform(X), axis=0).T

        # Let us return the shared response sliced by sessions
        return [concatenated_s[i] for i in timeframes_slices] def reduce_data(imgs, atlas, n_jobs=1, low_ram=False, temp_dir=None):
    """Reduce data using given atlas.
    Work done in parallel across subjects.

    Parameters
    ----------

    imgs : array of str, shape=[n_subjects, n_sessions]
        Element i, j of the array is a path to the data of subject i
        collected during session j.
        Data are loaded with numpy.load and expected shape is
        [n_timeframes, n_voxels]
        n_timeframes and n_voxels are assumed to be the same across subjects
        n_timeframes can vary across sessions
        Each voxel's timecourse is assumed to have mean 0 and variance 1

        imgs can also be a list of list of arrays where element i, j of
        the array is a numpy array of shape [n_voxels, n_timeframes] that
        contains the data of subject i collected during session j.

        imgs can also be a list of arrays where element i of the array is
        a numpy array of shape [n_voxels, n_timeframes] that contains the
        data of subject i (number of sessions is implicitly 1)

    atlas :  array, shape=[n_supervoxels, n_voxels] or array, shape=[n_voxels]
        or None
        Probabilistic or deterministic atlas on which to project the data
        Deterministic atlas is an array of shape [n_voxels,] where values
        range from 1 to n_supervoxels. Voxels labelled 0 will be ignored.

    n_jobs : integer, optional, default=1
        The number of CPUs to use to do the computation.
         -1 means all CPUs, -2 all CPUs but one, and so on.

    temp_dir : str or None
        path to dir where temporary results are stored
        if None temporary results will be stored in memory. This
        can results in memory errors when the number of subjects
        and / or sessions is large

    low_ram : bool
        if True and temp_dir is not None, reduced_data will be saved on disk
        this increases the number of IO but reduces memory complexity when
         the number of subject and/or sessions is large

    Returns
    -------

    reduced_data_list : array of str, shape=[n_subjects, n_sessions]
    or array, shape=[n_subjects, n_sessions, n_timeframes, n_supervoxels]
        Element i, j of the array is a path to the data of subject i collected
         during session j.
        Data are loaded with numpy.load and expected shape is
        [n_timeframes, n_supervoxels]
        or Element i, j of the array is the data in array of
        shape=[n_timeframes, n_supervoxels]
        n_timeframes and n_supervoxels
         are assumed to be the same across subjects
        n_timeframes can vary across sessions
        Each voxel's timecourse is assumed to have mean 0 and variance 1
    """
    if atlas is None:
        A = None
        A_inv = None
    else:
        loaded_atlas = safe_load(atlas)
        if len(loaded_atlas.shape) == 2:
            A = None
            A_inv = loaded_atlas.T.dot(
                np.linalg.inv(loaded_atlas.dot(loaded_atlas.T)))
        else:
            A = loaded_atlas
            A_inv = None

    n_subjects = len(imgs)
    n_sessions = len(imgs[0])

    reduced_data_list = Parallel(n_jobs=n_jobs)(
        delayed(reduce_data_single)(i,
                                    j,
                                    imgs[i][j],
                                    atlas=A,
                                    inv_atlas=A_inv,
                                    low_ram=low_ram,
                                    temp_dir=temp_dir)
        for i in range(n_subjects) for j in range(n_sessions))

    if low_ram:
        reduced_data_list = np.reshape(reduced_data_list,
                                       (n_subjects, n_sessions))
    else:
        if len(np.array(reduced_data_list).shape) == 1:
            reduced_data_list = np.reshape(reduced_data_list,
                                           (n_subjects, n_sessions))
        else:
            n_timeframes, n_supervoxels = np.array(reduced_data_list).shape[1:]
            reduced_data_list = np.reshape(
                reduced_data_list,
                (n_subjects, n_sessions, n_timeframes, n_supervoxels))

    return reduced_data_list def safe_load(data):
    """If data is an array returns data else returns np.load(data)"""
    if isinstance(data, np.ndarray):
        return data
    else:
        return np.load(data) class DetSRM(BaseEstimator, TransformerMixin):
    """Deterministic Shared Response Model (DetSRM)

    Given multi-subject data, factorize it as a shared response S among all
    subjects and an orthogonal transform W per subject:

    .. math:: X_i \\approx W_i S, \\forall i=1 \\dots N

    Parameters
    ----------

    n_iter : int, default: 10
        Number of iterations to run the algorithm.

    features : int, default: 50
        Number of features to compute.

    rand_seed : int, default: 0
        Seed for initializing the random number generator.


    Attributes
    ----------

    w_ : list of array, element i has shape=[voxels_i, features]
        The orthogonal transforms (mappings) for each subject.

    s_ : array, shape=[features, samples]
        The shared response.

    random_state_: `RandomState`
        Random number generator initialized using rand_seed

    Note
    ----

        The number of voxels may be different between subjects. However, the
        number of samples must be the same across subjects.

        The Deterministic Shared Response Model is approximated using the
        Block Coordinate Descent (BCD) algorithm proposed in [Chen2015]_.

        This is a single node version.

        The run-time complexity is :math:`O(I (V T K + V K^2))` and the memory
        complexity is :math:`O(V T)` with I - the number of iterations, V - the
        sum of voxels from all subjects, T - the number of samples, K - the
        number of features (typically, :math:`V \\gg T \\gg K`), and N - the
        number of subjects.
    """

    def __init__(self, n_iter=10, features=50, rand_seed=0):
        self.n_iter = n_iter
        self.features = features
        self.rand_seed = rand_seed
        return

    def fit(self, X, y=None):
        """Compute the Deterministic Shared Response Model

        Parameters
        ----------
        X : list of 2D arrays, element i has shape=[voxels_i, samples]
            Each element in the list contains the fMRI data of one subject.

        y : not used
        """
        logger.info('Starting Deterministic SRM')

        # Check the number of subjects
        if len(X) <= 1:
            raise ValueError("There are not enough subjects "
                             "({0:d}) to train the model.".format(len(X)))

        # Check for input data sizes
        if X[0].shape[1] < self.features:
            raise ValueError(
                "There are not enough samples to train the model with "
                "{0:d} features.".format(self.features))

        # Check if all subjects have same number of TRs
        number_trs = X[0].shape[1]
        number_subjects = len(X)
        for subject in range(number_subjects):
            assert_all_finite(X[subject])
            if X[subject].shape[1] != number_trs:
                raise ValueError("Different number of samples between subjects"
                                 ".")

        # Run SRM
        self.w_, self.s_ = self._srm(X)

        return self

    def transform(self, X, y=None):
        """Use the model to transform data to the Shared Response subspace

        Parameters
        ----------
        X : list of 2D arrays, element i has shape=[voxels_i, samples_i]
            Each element in the list contains the fMRI data of one subject.

        y : not used


        Returns
        -------
        s : list of 2D arrays, element i has shape=[features_i, samples_i]
            Shared responses from input data (X)
        """

        # Check if the model exist
        if hasattr(self, 'w_') is False:
            raise NotFittedError("The model fit has not been run yet.")

        # Check the number of subjects
        if len(X) != len(self.w_):
            raise ValueError("The number of subjects does not match the one"
                             " in the model.")

        s = [None] * len(X)
        for subject in range(len(X)):
            s[subject] = self.w_[subject].T.dot(X[subject])

        return s

    def _objective_function(self, data, w, s):
        """Calculate the objective function

        Parameters
        ----------

        data : list of 2D arrays, element i has shape=[voxels_i, samples]
            Each element in the list contains the fMRI data of one subject.

        w : list of 2D arrays, element i has shape=[voxels_i, features]
            The orthogonal transforms (mappings) :math:`W_i` for each subject.

        s : array, shape=[features, samples]
            The shared response

        Returns
        -------

        objective : float
            The objective function value.
        """
        subjects = len(data)
        objective = 0.0
        for m in range(subjects):
            objective += \
                np.linalg.norm(data[m] - w[m].dot(s), 'fro') ** 2

        return objective * 0.5 / data[0].shape[1]

    def _compute_shared_response(self, data, w):
        """ Compute the shared response S

        Parameters
        ----------

        data : list of 2D arrays, element i has shape=[voxels_i, samples]
            Each element in the list contains the fMRI data of one subject.

        w : list of 2D arrays, element i has shape=[voxels_i, features]
            The orthogonal transforms (mappings) :math:`W_i` for each subject.

        Returns
        -------

        s : array, shape=[features, samples]
            The shared response for the subjects data with the mappings in w.
        """
        s = np.zeros((w[0].shape[1], data[0].shape[1]))
        for m in range(len(w)):
            s = s + w[m].T.dot(data[m])
        s /= len(w)

        return s

    @staticmethod
    def _update_transform_subject(Xi, S):
        """Updates the mappings `W_i` for one subject.

        Parameters
        ----------

        Xi : array, shape=[voxels, timepoints]
            The fMRI data :math:`X_i` for aligning the subject.

        S : array, shape=[features, timepoints]
            The shared response.

        Returns
        -------

        Wi : array, shape=[voxels, features]
            The orthogonal transform (mapping) :math:`W_i` for the subject.
        """
        A = Xi.dot(S.T)
        # Solve the Procrustes problem
        U, _, V = np.linalg.svd(A, full_matrices=False)
        return U.dot(V)

    def transform_subject(self, X):
        """Transform a new subject using the existing model.
        The subject is assumed to have recieved equivalent stimulation

        Parameters
        ----------

        X : 2D array, shape=[voxels, timepoints]
            The fMRI data of the new subject.

        Returns
        -------

        w : 2D array, shape=[voxels, features]
            Orthogonal mapping `W_{new}` for new subject
        """
        # Check if the model exist
        if hasattr(self, 'w_') is False:
            raise NotFittedError("The model fit has not been run yet.")

        # Check the number of TRs in the subject
        if X.shape[1] != self.s_.shape[1]:
            raise ValueError("The number of timepoints(TRs) does not match the"
                             "one in the model.")

        w = self._update_transform_subject(X, self.s_)

        return w

    def _srm(self, data):
        """Expectation-Maximization algorithm for fitting the probabilistic SRM.

        Parameters
        ----------

        data : list of 2D arrays, element i has shape=[voxels_i, samples]
            Each element in the list contains the fMRI data of one subject.


        Returns
        -------

        w : list of array, element i has shape=[voxels_i, features]
            The orthogonal transforms (mappings) :math:`W_i` for each subject.

        s : array, shape=[features, samples]
            The shared response.
        """

        subjects = len(data)

        self.random_state_ = np.random.RandomState(self.rand_seed)
        random_states = [
            np.random.RandomState(self.random_state_.randint(2 ** 32))
            for i in range(len(data))]

        # Initialization step: initialize the outputs with initial values,
        # voxels with the number of voxels in each subject.
        w, _ = _init_w_transforms(data, self.features, random_states)
        shared_response = self._compute_shared_response(data, w)
        if logger.isEnabledFor(logging.INFO):
            # Calculate the current objective function value
            objective = self._objective_function(data, w, shared_response)
            logger.info('Objective function %f' % objective)

        # Main loop of the algorithm
        for iteration in range(self.n_iter):
            logger.info('Iteration %d' % (iteration + 1))

            # Update each subject's mapping transform W_i:
            for subject in range(subjects):
                a_subject = data[subject].dot(shared_response.T)
                perturbation = np.zeros(a_subject.shape)
                np.fill_diagonal(perturbation, 0.001)
                u_subject, _, v_subject = np.linalg.svd(
                    a_subject + perturbation, full_matrices=False)
                w[subject] = u_subject.dot(v_subject)

            # Update the shared response:
            shared_response = self._compute_shared_response(data, w)

            if logger.isEnabledFor(logging.INFO):
                # Calculate the current objective function value
                objective = self._objective_function(data, w, shared_response)
                logger.info('Objective function %f' % objective)

        return w, shared_response
import os
import tempfile
import numpy as np
import pytest
from numpy.testing import assert_array_almost_equal
from sklearn.exceptions import NotFittedError
from brainiak.funcalign.fastsrm import (
    FastSRM, _compute_and_save_corr_mat, _compute_and_save_subject_basis,
    _compute_basis_subject_online, _reduced_space_compute_shared_response,
    check_atlas, check_imgs, check_shared_response, create_temp_dir, fast_srm,
    reduce_data, safe_load)
from brainiak.funcalign.srm import DetSRM
            paths, _, _ = generate_data(n_voxels, n_timeframes, n_subjects,
                                        n_components, datadir)

            # Test if reduced data has the good shape
            # probabilistic atlas
            atlas = np.random.rand(n_supervoxels, n_voxels)
            data = reduce_data(paths,
                               atlas=atlas,
                               n_jobs=n_jobs,
                               low_ram=False,
                               temp_dir=None)

            for subject in range(n_subjects):
                for session in range(n_sessions):
                    assert data[subject, session].shape == (
                        n_timeframes[session], n_supervoxels)

            # deterministic atlas
            det_atlas = np.round(np.random.rand(n_voxels) * n_supervoxels)
            n_unique = len(np.unique(det_atlas)[1:])
            while n_unique != n_supervoxels:
                det_atlas = np.round(np.random.rand(n_voxels) * n_supervoxels)
                n_unique = len(np.unique(det_atlas)[1:])

            data = reduce_data(paths,
                               atlas=det_atlas,
                               n_jobs=n_jobs,
                               low_ram=True,
                               temp_dir=datadir)

            for subject in range(n_subjects):
                for session in range(n_sessions):
                    assert (np.load(data[subject, session]).shape == (
                        n_timeframes[session], n_supervoxels))


def test_reduced_data_srm():
    n_jobs = 1
    with tempfile.TemporaryDirectory() as datadir:
        np.random.seed(0)

        # We authorize different timeframes for different sessions but
        # they should be the same across subject
        n_voxels = 10
        n_timeframes = [25, 24]
        n_subjects = 5
        n_components = 3  # number of components used for SRM model
        n_sessions = len(n_timeframes)

        paths, W, S = generate_data(n_voxels, n_timeframes, n_subjects,
                                    n_components, datadir, 0)

        atlas = np.arange(1, n_voxels + 1)
        data = reduce_data(paths,
                           atlas=atlas,
                           n_jobs=n_jobs,
                           low_ram=False,
                           temp_dir=None)

        # Test if shared response has the good shape
        shared_response_list = \
            _reduced_space_compute_shared_response(
                data,
                reduced_basis_list=None,
                n_components=n_components
            )
        assert len(shared_response_list) == n_sessions

        for session in range(n_sessions):
            assert (shared_response_list[session].shape == (
                n_timeframes[session], n_components))

        # Test basis from shared response
        for i, sessions in enumerate(paths):
            basis = _compute_basis_subject_online(
                sessions, [S[k].T for k in range(len(S))])
            # test shape
            assert basis.shape == (n_components, n_voxels)
            # test orthogonality
            assert np.allclose(basis.dot(basis.T), np.eye(n_components))
            # test correctness
            assert_array_almost_equal(basis, W[i], 2)

        # Test reduced_data_shared_response
        shared_response_list = _reduced_space_compute_shared_response(
            data, reduced_basis_list=W, n_components=n_components)
        for session in range(n_sessions):
            S_real = np.mean(
                [data[i, session].dot(W[i].T) for i in range(n_subjects)],
                axis=0)
            assert_array_almost_equal(shared_response_list[session], S_real)
            assert_array_almost_equal(shared_response_list[session],
                                      S[session].T)

        shared_response_list = fast_srm(data, n_components=n_components)

        for i, sessions in enumerate(paths):
            basis = _compute_basis_subject_online(sessions,
                                                  shared_response_list)
            for j, session in enumerate(sessions):
                assert_array_almost_equal(shared_response_list[j].dot(basis),
                                          np.load(paths[i, j]).T, 3)


def test_compute_and_save():
    with tempfile.TemporaryDirectory() as datadir:
        np.random.seed(0)
        n_voxels = 10
        n_timeframes = [25, 24]
        n_subjects = 5
        n_components = 3  # number of components used for SRM model

        paths, W, S = generate_data(n_voxels, n_timeframes, n_subjects,
                                    n_components, datadir, 0)

        for m, subjects in enumerate(paths.T):
            for subject in subjects:
                _compute_and_save_corr_mat(subject, S[m].T, datadir)

        for i, sessions in enumerate(paths):

            basis = _compute_and_save_subject_basis(i, sessions, datadir)