brainiak · CameronTEllis · May 15, 2018 · May 4, 2018 · May 4, 2018 · May 4, 2018
diff --git 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
@@ -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])