diff --git a/brainiak/funcalign/srm.py b/brainiak/funcalign/srm.py index 668085a81..1eef5f79c 100644 --- a/brainiak/funcalign/srm.py +++ b/brainiak/funcalign/srm.py @@ -357,6 +357,60 @@ def _likelihood(self, chol_sigma_s_rhos, log_det_psi, chol_sigma_s, return loglikehood + @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. @@ -393,7 +447,7 @@ def _srm(self, data): subjects = len(data) self.random_state_ = np.random.RandomState(self.rand_seed) random_states = [ - np.random.RandomState(self.random_state_.randint(2**32)) + np.random.RandomState(self.random_state_.randint(2 ** 32)) for i in range(len(data))] # Initialization step: initialize the outputs with initial values, @@ -453,7 +507,7 @@ def _srm(self, data): # Update the shared response shared_response = sigma_s.dot( np.identity(self.features) - rho0 * inv_sigma_s_rhos).dot( - wt_invpsi_x) + wt_invpsi_x) # M-step @@ -649,7 +703,7 @@ def _objective_function(self, data, w, s): objective = 0.0 for m in range(subjects): objective += \ - np.linalg.norm(data[m] - w[m].dot(s), 'fro')**2 + np.linalg.norm(data[m] - w[m].dot(s), 'fro') ** 2 return objective * 0.5 / data[0].shape[1] @@ -678,6 +732,59 @@ def _compute_shared_response(self, data, 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. @@ -702,7 +809,7 @@ def _srm(self, data): self.random_state_ = np.random.RandomState(self.rand_seed) random_states = [ - np.random.RandomState(self.random_state_.randint(2**32)) + np.random.RandomState(self.random_state_.randint(2 ** 32)) for i in range(len(data))] # Initialization step: initialize the outputs with initial values, diff --git a/tests/funcalign/test_srm.py b/tests/funcalign/test_srm.py index c07aa42bf..08a022dc4 100644 --- a/tests/funcalign/test_srm.py +++ b/tests/funcalign/test_srm.py @@ -118,6 +118,84 @@ def test_can_instantiate(): print("Test: different number of samples per subject") +def test_new_subject(): + import brainiak.funcalign.srm + s = brainiak.funcalign.srm.SRM() + assert s, "Invalid SRM instance!" + + import numpy as np + np.random.seed(0) + + voxels = 100 + samples = 500 + subjects = 3 + features = 3 + + s = brainiak.funcalign.srm.SRM(n_iter=5, features=features) + 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)) + + X = [] + W = [] + Q, R = np.linalg.qr(np.random.random((voxels, features))) + W.append(Q) + X.append(Q.dot(S) + 0.1*np.random.random((voxels, samples))) + + for subject in range(1, subjects): + Q, R = np.linalg.qr(np.random.random((voxels, features))) + W.append(Q) + X.append(Q.dot(S) + 0.1*np.random.random((voxels, samples))) + + # Check that transform does NOT run before fitting the model + with pytest.raises(NotFittedError): + s.transform_subject(X) + print("Test: transforming before fitting the model") + + # Check that runs with 3 subject + s.fit(X) + + # Check that you get an error when the data is the wrong shape + with pytest.raises(ValueError): + s.transform_subject(X[0].T) + + # Check that it does run to compute a new subject + new_w = s.transform_subject(X[0]) + assert new_w.shape[1] == features, ( + "Invalid computation of SRM! (wrong # features for new subject)") + assert new_w.shape[0] == voxels, ( + "Invalid computation of SRM! (wrong # voxels for new subject)") + + # Check that these analyses work with the deterministic SRM too + ds = brainiak.funcalign.srm.DetSRM(n_iter=5, features=features) + + # Check that transform does NOT run before fitting the model + with pytest.raises(NotFittedError): + ds.transform_subject(X) + print("Test: transforming before fitting the model") + + # Check that runs with 3 subject + ds.fit(X) + + # Check that you get an error when the data is the wrong shape + with pytest.raises(ValueError): + ds.transform_subject(X[0].T) + + # Check that it does run to compute a new subject + new_w = ds.transform_subject(X[0]) + assert new_w.shape[1] == features, ( + "Invalid computation of SRM! (wrong # features for new subject)") + assert new_w.shape[0] == voxels, ( + "Invalid computation of SRM! (wrong # voxels for new subject)") + + def test_det_srm(): import brainiak.funcalign.srm model = brainiak.funcalign.srm.DetSRM() @@ -200,6 +278,13 @@ def test_det_srm(): assert new_s[subject].shape[1] == samples, ( "Invalid computation of DetSRM! (wrong # samples after transform)") + # Check that it does run to compute a new subject + new_w = model.transform_subject(X[0]) + assert new_w.shape[1] == features, ( + "Invalid computation of SRM! (wrong # features for new subject)") + assert new_w.shape[0] == voxels, ( + "Invalid computation of SRM! (wrong # voxels for new subject)") + # Check that it does NOT run with non-matching number of subjects with pytest.raises(ValueError): model.transform(X[1])