brainiak · mihaic · Oct 14, 2020 · Oct 10, 2020 · Oct 10, 2020 · Oct 13, 2020
diff --git a/brainiak/funcalign/srm.py b/brainiak/funcalign/srm.py
@@ -106,6 +106,41 @@ def _init_w_transforms(data, features, random_states, comm=MPI.COMM_SELF):
     return w, voxels
 
 
+def load(file):
+    """Load fitted SRM from .npz file.
+
+    Parameters
+    ----------
+
+    file : str, file-like object, or pathlib.Path
+        The .npz file to read containing fitted SRM saved using srm.save
+
+    Returns
+    --------
+
+    srm : fitted SRM model
+    """
+
+    # Load file and extract SRM attributes
+    loaded = np.load(file)
+    w_ = [s for s in loaded['w_']]
+    s_ = loaded['s_']
+    sigma_s_ = loaded['sigma_s_']
+    mu_ = [s for s in loaded['mu_']]
+    rho2_ = loaded['rho2_']
+    features, n_iter, rand_seed = loaded['kwargs']
+
+    # Initialize new SRM object and attach loaded attributes
+    srm = SRM(n_iter=n_iter, features=features, rand_seed=rand_seed)
+    srm.w_ = w_
+    srm.s_ = s_
+    srm.sigma_s_ = sigma_s_
+    srm.mu_ = mu_
+    srm.rho2_ = rho2_
+
+    return srm
+
+
 class SRM(BaseEstimator, TransformerMixin):
     """Probabilistic Shared Response Model (SRM)
 
@@ -412,6 +447,38 @@ def transform_subject(self, X):
 
         return w
 
+    def save(self, file):
+        """Save fitted SRM to .npz file.
+
+        Parameters
+        ----------
+
+        file : str, file-like object, or pathlib.Path
+            Filename (string), open file (file-like object) or pathlib.Path
+            where the fitted SRM will be saved. If file is a string or a Path,
+            the .npz extension will be appended to the filename if it is not
+            already there.
+
+        Returns
+        -------
+
+        None
+        """
+
+        # Check if the model has been estimated
+        if hasattr(self, 'w_') is False:
+            raise NotFittedError("The model fit has not been run yet.")
+
+        np.savez_compressed(
+            file,
+            w_=self.w_,
+            s_=self.s_,
+            sigma_s_=self.sigma_s_,
+            mu_=self.mu_,
+            rho2_=self.rho2_,
+            kwargs=np.array([self.features, self.n_iter, self.rand_seed])
+        )
+
     def _srm(self, data):
         """Expectation-Maximization algorithm for fitting the probabilistic SRM.
 

diff --git a/tests/funcalign/test_srm.py b/tests/funcalign/test_srm.py
@@ -15,7 +15,7 @@
 import pytest
 
 
-def test_can_instantiate():
+def test_can_instantiate(tmp_path):
     import brainiak.funcalign.srm
     s = brainiak.funcalign.srm.SRM()
     assert s, "Invalid SRM instance!"
@@ -117,6 +117,21 @@ def test_can_instantiate():
         s.fit(X)
     print("Test: different number of samples per subject")
 
+    # Check save/load functionality for fitted SRM
+    srm_path = tmp_path / 'srm.npz'
+    s.save(srm_path)
+    s_load = brainiak.funcalign.srm.load(srm_path)
+    assert np.array_equal(s.s_, s_load.s_)
+    for w, wl in zip(s.w_, s_load.w_):
+        assert np.array_equal(w, wl)
+    assert np.array_equal(s.sigma_s_, s_load.sigma_s_)
+    assert np.array_equal(s.mu_, s_load.mu_)
+    assert np.array_equal(s.rho2_, s_load.rho2_)
+    assert s.features == s_load.features
+    assert s.n_iter == s_load.n_iter
+    assert s.rand_seed == s_load.rand_seed
+    print("Test: save/load functionality")
+
 
 def test_new_subject():
     import brainiak.funcalign.srm