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
        -------

        sigma_s : array, shape=[features, features]
            The covariance :math:`\\Sigma_s` of the shared response Normal
            distribution.

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

        mu : list of array, element i has shape=[voxels_i]
            The voxel means :math:`\\mu_i` over the samples for each subject.

        rho2 : array, shape=[subjects]
            The estimated noise variance :math:`\\rho_i^2` for each subject

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

        local_min = min([d.shape[1] for d in data if d is not None],
                        default=sys.maxsize)
        samples = self.comm.allreduce(local_min, op=MPI.MIN)
        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, and trace_xtx with
        # the ||X_i||_F^2 of each subject.
        w, voxels = _init_w_transforms(data, self.features, random_states,
                                       self.comm)
        x, mu, rho2, trace_xtx = self._init_structures(data, subjects)
        shared_response = np.zeros((self.features, samples))
        sigma_s = np.identity(self.features)

        rank = self.comm.Get_rank()

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

            # E-step:

            # Sum the inverted the rho2 elements for computing W^T * Psi^-1 * W
            if rank == 0:
                rho0 = (1 / rho2).sum()

                # Invert Sigma_s using Cholesky factorization
                (chol_sigma_s, lower_sigma_s) = scipy.linalg.cho_factor(
                    sigma_s, check_finite=False)
                inv_sigma_s = scipy.linalg.cho_solve(
                    (chol_sigma_s, lower_sigma_s), np.identity(self.features),
                    check_finite=False)

                # Invert (Sigma_s + rho_0 * I) using Cholesky factorization
                sigma_s_rhos = inv_sigma_s + np.identity(self.features) * rho0
                chol_sigma_s_rhos, lower_sigma_s_rhos = \
                    scipy.linalg.cho_factor(sigma_s_rhos,
                                            check_finite=False)
                inv_sigma_s_rhos = scipy.linalg.cho_solve(
                    (chol_sigma_s_rhos, lower_sigma_s_rhos),
                    np.identity(self.features), check_finite=False)

            # Compute the sum of W_i^T * rho_i^-2 * X_i, and the sum of traces
            # of X_i^T * rho_i^-2 * X_i
            wt_invpsi_x = np.zeros((self.features, samples))
            trace_xt_invsigma2_x = 0.0
            for subject in range(subjects):
                if data[subject] is not None:
                    wt_invpsi_x += (w[subject].T.dot(x[subject])) \
                                   / rho2[subject]
                    trace_xt_invsigma2_x += trace_xtx[subject] / rho2[subject]

            wt_invpsi_x = self.comm.reduce(wt_invpsi_x, op=MPI.SUM)
            trace_xt_invsigma2_x = self.comm.reduce(trace_xt_invsigma2_x,
                                                    op=MPI.SUM)
            trace_sigma_s = None
            if rank == 0:
                log_det_psi = np.sum(np.log(rho2) * voxels)

                # Update the shared response
                shared_response = sigma_s.dot(
                    np.identity(self.features) - rho0 * inv_sigma_s_rhos).dot(
                    wt_invpsi_x)

                # M-step

                # Update Sigma_s and compute its trace
                sigma_s = (inv_sigma_s_rhos
                           + shared_response.dot(shared_response.T) / samples)
                trace_sigma_s = samples * np.trace(sigma_s)

            shared_response = self.comm.bcast(shared_response)
            trace_sigma_s = self.comm.bcast(trace_sigma_s)

            # Update each subject's mapping transform W_i and error variance
            # rho_i^2
            for subject in range(subjects):
                if x[subject] is not None:
                    a_subject = x[subject].dot(shared_response.T)
                    perturbation = np.zeros(a_subject.shape)
                    np.fill_diagonal(perturbation, 0.001)
                    u_subject, s_subject, v_subject = np.linalg.svd(
                        a_subject + perturbation, full_matrices=False)
                    w[subject] = u_subject.dot(v_subject)
                    rho2[subject] = trace_xtx[subject]
                    rho2[subject] += -2 * np.sum(w[subject] * a_subject).sum()
                    rho2[subject] += trace_sigma_s
                    rho2[subject] /= samples * voxels[subject]
                else:
                    rho2[subject] = 0

            rho2 = self.comm.allreduce(rho2, op=MPI.SUM)

            if rank == 0:
                if logger.isEnabledFor(logging.INFO):
                    # Calculate and log the current log-likelihood for checking
                    # convergence
                    loglike = self._likelihood(
                        chol_sigma_s_rhos, log_det_psi, chol_sigma_s,
                        trace_xt_invsigma2_x, inv_sigma_s_rhos, wt_invpsi_x,
                        samples)
                    logger.info('Objective function %f' % loglike)

        sigma_s = self.comm.bcast(sigma_s)
        return sigma_s, w, mu, rho2, shared_response
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
        -------

        sigma_s : array, shape=[features, features]
            The covariance :math:`\\Sigma_s` of the shared response Normal
            distribution.

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

        mu : list of array, element i has shape=[voxels_i]
            The voxel means :math:`\\mu_i` over the samples for each subject.

        rho2 : array, shape=[subjects]
            The estimated noise variance :math:`\\rho_i^2` for each subject

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