diff --git a/brainiak/funcalign/srm.py b/brainiak/funcalign/srm.py index a2ec01ac6..668085a81 100644 --- a/brainiak/funcalign/srm.py +++ b/brainiak/funcalign/srm.py @@ -38,7 +38,9 @@ import scipy from sklearn.base import BaseEstimator, TransformerMixin from sklearn.utils import assert_all_finite -from sklearn.utils.validation import NotFittedError +from sklearn.exceptions import NotFittedError +from mpi4py import MPI +import sys __all__ = [ "SRM", "DetSRM" @@ -47,7 +49,7 @@ logger = logging.getLogger(__name__) -def _init_w_transforms(data, features, random_states): +def _init_w_transforms(data, features, random_states, comm=MPI.COMM_SELF): """Initialize the mappings (Wi) for the SRM with random orthogonal matrices. Parameters @@ -62,6 +64,8 @@ def _init_w_transforms(data, features, random_states): random_states : list of `RandomState`s One `RandomState` instance per subject. + comm : mpi4py.MPI.Intracomm + The MPI communicator containing the data Returns ------- @@ -88,12 +92,16 @@ def _init_w_transforms(data, features, random_states): # Set Wi to a random orthogonal voxels by features matrix for subject in range(subjects): - voxels[subject] = data[subject].shape[0] - rnd_matrix = random_states[subject].random_sample((voxels[subject], - features)) - q, r = np.linalg.qr(rnd_matrix) - w.append(q) - + if data[subject] is not None: + voxels[subject] = data[subject].shape[0] + rnd_matrix = random_states[subject].random_sample(( + voxels[subject], features)) + q, r = np.linalg.qr(rnd_matrix) + w.append(q) + else: + voxels[subject] = 0 + w.append(None) + voxels = comm.allreduce(voxels, op=MPI.SUM) return w, voxels @@ -117,6 +125,8 @@ class SRM(BaseEstimator, TransformerMixin): rand_seed : int, default: 0 Seed for initializing the random number generator. + comm : mpi4py.MPI.Intracomm + The MPI communicator containing the data Attributes ---------- @@ -136,10 +146,12 @@ class SRM(BaseEstimator, TransformerMixin): rho2_ : array, shape=[subjects] The estimated noise variance :math:`\\rho_i^2` for each subject + comm : mpi4py.MPI.Intracomm + The MPI communicator containing the data + random_state_: `RandomState` Random number generator initialized using rand_seed - Note ---- @@ -158,10 +170,12 @@ class SRM(BaseEstimator, TransformerMixin): K - the number of features (typically, :math:`V \\gg T \\gg K`). """ - def __init__(self, n_iter=10, features=50, rand_seed=0): + def __init__(self, n_iter=10, features=50, rand_seed=0, + comm=MPI.COMM_SELF): self.n_iter = n_iter self.features = features self.rand_seed = rand_seed + self.comm = comm return def fit(self, X, y=None): @@ -182,20 +196,36 @@ def fit(self, X, y=None): "({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)) + number_subjects = len(X) + number_subjects_vec = self.comm.allgather(number_subjects) + for rank in range(self.comm.Get_size()): + if number_subjects_vec[rank] != number_subjects: + raise ValueError( + "Not all ranks have same number of subjects") + + # Collect size information + shape0 = np.zeros((number_subjects,), dtype=np.int) + shape1 = np.zeros((number_subjects,), dtype=np.int) + + for subject in range(number_subjects): + if X[subject] is not None: + assert_all_finite(X[subject]) + shape0[subject] = X[subject].shape[0] + shape1[subject] = X[subject].shape[1] + + shape0 = self.comm.allreduce(shape0, op=MPI.SUM) + shape1 = self.comm.allreduce(shape1, op=MPI.SUM) # Check if all subjects have same number of TRs - number_trs = X[0].shape[1] - number_subjects = len(X) + number_trs = np.min(shape1) for subject in range(number_subjects): - assert_all_finite(X[subject]) - if X[subject].shape[1] != number_trs: + if shape1[subject] < self.features: + raise ValueError( + "There are not enough samples to train the model with " + "{0:d} features.".format(self.features)) + if shape1[subject] != number_trs: raise ValueError("Different number of samples between subjects" ".") - # Run SRM self.sigma_s_, self.w_, self.mu_, self.rho2_, self.s_ = self._srm(X) @@ -229,7 +259,8 @@ def transform(self, X, y=None): s = [None] * len(X) for subject in range(len(X)): - s[subject] = self.w_[subject].T.dot(X[subject]) + if X[subject] is not None: + s[subject] = self.w_[subject].T.dot(X[subject]) return s @@ -265,11 +296,17 @@ def _init_structures(self, data, subjects): rho2 = np.zeros(subjects) trace_xtx = np.zeros(subjects) + for subject in range(subjects): - mu.append(np.mean(data[subject], 1)) rho2[subject] = 1 - trace_xtx[subject] = np.sum(data[subject] ** 2) - x.append(data[subject] - mu[subject][:, np.newaxis]) + if data[subject] is not None: + mu.append(np.mean(data[subject], 1)) + trace_xtx[subject] = np.sum(data[subject] ** 2) + x.append(data[subject] - mu[subject][:, np.newaxis]) + else: + mu.append(None) + trace_xtx[subject] = 0 + x.append(None) return x, mu, rho2, trace_xtx @@ -350,9 +387,10 @@ def _srm(self, data): The shared response. """ - samples = data[0].shape[1] + 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)) @@ -361,11 +399,14 @@ def _srm(self, 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) + 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)) @@ -373,68 +414,87 @@ def _srm(self, data): # E-step: # Sum the inverted the rho2 elements for computing W^T * Psi^-1 * W - 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) + 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): - wt_invpsi_x += (w[subject].T.dot(x[subject])) / rho2[subject] - trace_xt_invsigma2_x += trace_xtx[subject] / rho2[subject] + 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] - log_det_psi = np.sum(np.log(rho2) * voxels) + 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) + # Update the shared response + shared_response = sigma_s.dot( + np.identity(self.features) - rho0 * inv_sigma_s_rhos).dot( + wt_invpsi_x) - # M-step + # 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) + # 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): - 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] - - 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) - + 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 diff --git a/examples/funcalign/srm_image_prediction_example_distributed.py b/examples/funcalign/srm_image_prediction_example_distributed.py new file mode 100644 index 000000000..2c164ba4f --- /dev/null +++ b/examples/funcalign/srm_image_prediction_example_distributed.py @@ -0,0 +1,147 @@ +# Copyright 2016 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import scipy.io +from scipy.stats import stats +from sklearn.metrics import confusion_matrix +from sklearn.svm import NuSVC +import numpy as np +import brainiak.funcalign.srm +from mpi4py import MPI + +comm = MPI.COMM_WORLD +rank = comm.Get_rank() +size = comm.Get_size() + +# Load the input data that contains the movie stimuli for unsupervised training with SRM +movie_file = scipy.io.loadmat('data/movie_data.mat') +movie_data_left = movie_file['movie_data_lh'] +movie_data_right = movie_file['movie_data_rh'] +subjects = movie_data_left.shape[2] + +# Convert data to a list of arrays matching SRM input. +# Each element is a matrix of voxels by TRs. +# Also, concatenate data from both hemispheres in the brain. +movie_data = [] +for s in range(subjects): + if s % size == rank: + movie_data.append(np.concatenate([movie_data_left[:, :, s], movie_data_right[:, :, s]], axis=0)) + else: + movie_data.append(None) + +# Z-score the data +for subject in range(subjects): + if movie_data[subject] is not None: + movie_data[subject] = stats.zscore(movie_data[subject], axis=1, ddof=1) + +# Run SRM with the movie data +srm = brainiak.funcalign.srm.SRM(n_iter=10, features=50, comm=comm) +srm.fit(movie_data) + +# We define a function to present the output of the experiment. +def plot_confusion_matrix(cm, title="Confusion Matrix"): + """Plots a confusion matrix for each subject + """ + import matplotlib.pyplot as plt + import math + plt.figure() + subjects = len(cm) + root_subjects = math.sqrt(subjects) + cols = math.ceil(root_subjects) + rows = math.ceil(subjects/cols) + classes = cm[0].shape[0] + for subject in range(subjects): + plt.subplot(rows, cols, subject+1) + plt.imshow(cm[subject], interpolation='nearest', cmap=plt.cm.bone) + plt.xticks(np.arange(classes), range(1, classes+1)) + plt.yticks(np.arange(classes), range(1, classes+1)) + cbar = plt.colorbar(ticks=[0.0, 1.0], shrink=0.6) + cbar.set_clim(0.0, 1.0) + plt.xlabel("Predicted") + plt.ylabel("True label") + plt.title("{0:d}".format(subject + 1)) + plt.suptitle(title) + plt.tight_layout() + plt.show() + +# Load the input data that contains the image stimuli and its labels for training a classifier +image_file = scipy.io.loadmat('data/image_data.mat') +image_data_left = image_file['image_data_lh'] +image_data_right = image_file['image_data_rh'] + +# Convert data to a list of arrays matching SRM input. +# Each element is a matrix of voxels by TRs. +# Also, concatenate data from both hemispheres in the brain. +image_data = [] +for s in range(subjects): + if s % size == rank: + image_data.append(np.concatenate([image_data_left[:, :, s], image_data_right[:, :, s]], axis=0)) + else: + image_data.append(None) + +# Z-score the image data +for subject in range(subjects): + if subject % size == rank: + image_data[subject] = stats.zscore(image_data[subject], axis=1, ddof=1) + else: + image_data[subject] = None + +# Z-score the shared response data +image_data_shared = srm.transform(image_data) + +for subject in range(subjects): + image_data_shared[subject] = comm.bcast(image_data_shared[subject], + root=(subject % size)) + +if rank == 0: + for subject in range(subjects): + image_data_shared[subject] = stats.zscore(image_data_shared[subject], axis=1, ddof=1) + + # Read the labels of the image data for training the classifier. + labels = scipy.io.loadmat('data/label.mat') + labels = np.squeeze(labels['label']) + + # Run a leave-one-out cross validation with the subjects + train_labels = np.tile(labels, subjects-1) + test_labels = labels + accuracy = np.zeros((subjects,)) + cm = [None] * subjects + for subject in range(subjects): + # Concatenate the subjects' data for training into one matrix + train_subjects = list(range(subjects)) + train_subjects.remove(subject) + TRs = image_data_shared[0].shape[1] + train_data = np.zeros((image_data_shared[0].shape[0], len(train_labels))) + for train_subject in range(len(train_subjects)): + start_index = train_subject*TRs + end_index = start_index+TRs + train_data[:, start_index:end_index] = image_data_shared[train_subjects[train_subject]] + + # Train a Nu-SVM classifier using scikit learn + classifier = NuSVC(nu=0.5, kernel='linear') + classifier = classifier.fit(train_data.T, train_labels) + + # Predict on the test data + predicted_labels = classifier.predict(image_data_shared[subject].T) + accuracy[subject] = sum(predicted_labels == test_labels)/float(len(predicted_labels)) + + # Create a confusion matrix to see the accuracy of each class + cm[subject] = confusion_matrix(test_labels, predicted_labels) + + # Normalize the confusion matrix + cm[subject] = cm[subject].astype('float') / cm[subject].sum(axis=1)[:, np.newaxis] + + + # Plot and print the results + plot_confusion_matrix(cm, title="Confusion matrices for different test subjects with Probabilistic SRM") + print("SRM: The average accuracy among all subjects is {0:f} +/- {1:f}".format(np.mean(accuracy), np.std(accuracy))) diff --git a/tests/funcalign/sr_v0_4.npz b/tests/funcalign/sr_v0_4.npz new file mode 100644 index 000000000..699080221 Binary files /dev/null and b/tests/funcalign/sr_v0_4.npz differ diff --git a/tests/funcalign/test_srm.py b/tests/funcalign/test_srm.py index 5f9694260..d73a0db74 100644 --- a/tests/funcalign/test_srm.py +++ b/tests/funcalign/test_srm.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from sklearn.utils.validation import NotFittedError +from sklearn.exceptions import NotFittedError import pytest def test_can_instantiate(): @@ -20,6 +20,7 @@ def test_can_instantiate(): assert s, "Invalid SRM instance!" import numpy as np + np.random.seed(0) voxels = 100 samples = 500 @@ -62,6 +63,9 @@ def test_can_instantiate(): # Check that runs with 2 subject s.fit(X) + from pathlib import Path + sr_v0_4 = np.load(Path(__file__).parent / "sr_v0_4.npz")['sr'] + assert(np.allclose(sr_v0_4, s.s_)) assert len(s.w_) == subjects, "Invalid computation of SRM! (wrong # subjects in W)" for subject in range(subjects): @@ -190,4 +194,3 @@ def test_det_srm(): model.fit(X) print("Test: different number of samples per subject") - diff --git a/tests/funcalign/test_srm_distributed.py b/tests/funcalign/test_srm_distributed.py new file mode 100644 index 000000000..bc921dd86 --- /dev/null +++ b/tests/funcalign/test_srm_distributed.py @@ -0,0 +1,140 @@ +# Copyright 2016 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from sklearn.exceptions import NotFittedError +import pytest +from mpi4py import MPI + +def test_distributed_srm(): + import brainiak.funcalign.srm + s = brainiak.funcalign.srm.SRM() + assert s, "Invalid SRM instance!" + + import numpy as np + np.random.seed(0) + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + nrank = comm.Get_size() + + voxels = 100 + samples = 500 + subjects = 2 + features = 3 + + s = brainiak.funcalign.srm.SRM(n_iter=5, features=features, comm=comm) + assert s, "Invalid SRM instance!" + + # Create a Shared response S with K = 3 + theta = np.linspace(-4 * np.pi, 4 * np.pi, samples) + z = np.linspace(-2, 2, samples) + r = z**2 + 1 + x = r * np.sin(theta) + y = r * np.cos(theta) + + S = np.vstack((x, y, z)) + + # DSRM: broadcast S + S = comm.bcast(S) + + X = [] + W = [] + # DSRM: only append on rank 0 + Q, R = np.linalg.qr(np.random.random((voxels, features))) + tmp_noise = 0.1*np.random.random((voxels, samples)) + if rank == 0: + W.append(Q) + X.append(Q.dot(S) + tmp_noise) + else: + W.append(None) + X.append(None) + + # Check that transform does NOT run before fitting the model + with pytest.raises(NotFittedError) as excinfo: + s.transform(X) + if rank == 0: + print("Test: transforming before fitting the model") + + # Check that it does NOT run with 1 subject + with pytest.raises(ValueError) as excinfo: + s.fit(X) + if rank == 0: + print("Test: running SRM with 1 subject") + + # DSRM: cyclic distribution of subject data, otherwise None + for subject in range(1, subjects): + Q, R = np.linalg.qr(np.random.random((voxels, features))) + tmp_noise = 0.1*np.random.random((voxels, samples)) + if subject % nrank == rank: + W.append(Q) + X.append(Q.dot(S) + tmp_noise) + else: + W.append(None) + X.append(None) + + # Check that runs with 2 subject + s.fit(X) + from pathlib import Path + sr_v0_4 = np.load(Path(__file__).parent / "sr_v0_4.npz")['sr'] + assert(np.allclose(sr_v0_4, s.s_)) + + assert len(s.w_) == subjects, "Invalid computation of SRM! (wrong # subjects in W)" + for subject in range(subjects): + if s.w_[subject] is not None: + assert s.w_[subject].shape[0] == voxels, "Invalid computation of SRM! (wrong # voxels in W)" + assert s.w_[subject].shape[1] == features, "Invalid computation of SRM! (wrong # features in W)" + ortho = np.linalg.norm(s.w_[subject].T.dot(s.w_[subject]) - np.eye(s.w_[subject].shape[1]), 'fro') + assert ortho < 1e-7, "A Wi mapping is not orthonormal in SRM." + difference = np.linalg.norm(X[subject] - s.w_[subject].dot(s.s_), 'fro') + datanorm = np.linalg.norm(X[subject], 'fro') + assert difference/datanorm < 1.0, "Model seems incorrectly computed." + + assert s.s_.shape[0] == features, "Invalid computation of SRM! (wrong # features in S)" + assert s.s_.shape[1] == samples, "Invalid computation of SRM! (wrong # samples in S)" + + # Check that it does run to compute the shared response after the model computation + new_s = s.transform(X) + + assert len(new_s) == subjects, "Invalid computation of SRM! (wrong # subjects after transform)" + for subject in range(subjects): + if new_s[subject] is not None: + assert new_s[subject].shape[0] == features, "Invalid computation of SRM! (wrong # features after transform)" + assert new_s[subject].shape[1] == samples, "Invalid computation of SRM! (wrong # samples after transform)" + + # Check that it does NOT run with non-matching number of subjects + with pytest.raises(ValueError) as excinfo: + s.transform([X[1]]) + if rank == 0: + print("Test: transforming with non-matching number of subjects") + + # Check that it does not run without enough samples (TRs). + with pytest.raises(ValueError) as excinfo: + s.set_params(features=(samples+1)) + s.fit(X) + if rank == 0: + print("Test: not enough samples") + + + # Check that it does not run with different number of samples (TRs) + if rank == 0: + S2 = S[:,:-2] + X.append(Q.dot(S2)) + else: + X.append(None) + + + with pytest.raises(ValueError) as excinfo: + s.fit(X) + if rank == 0: + print("Test: different number of samples per subject") + +test_distributed_srm()