diff --git a/brainiak/funcalign/srm.py b/brainiak/funcalign/srm.py index 023e71872..3ed56a1df 100644 --- 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 index 08a022dc4..b109a77bd 100644 --- 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