diff --git a/.gitignore b/.gitignore index 8f3224dc0..4147ffe5d 100644 --- a/.gitignore +++ b/.gitignore @@ -19,7 +19,7 @@ data brainiak.egg-info build -.cache +.*cache .coverage* coverage.xml dist diff --git a/brainiak/fcma/mvpa_voxelselector.py b/brainiak/fcma/mvpa_voxelselector.py index c6ecc7848..a308f9e2d 100644 --- a/brainiak/fcma/mvpa_voxelselector.py +++ b/brainiak/fcma/mvpa_voxelselector.py @@ -31,6 +31,24 @@ ] +def _sfn(l, mask, myrad, bcast_var): + """Score classifier on searchlight data using cross-validation. + + The classifier is in `bcast_var[2]`. The labels are in `bast_var[0]`. The + number of cross-validation folds is in `bast_var[1]. + """ + clf = bcast_var[2] + data = l[0][mask, :].T + # print(l[0].shape, mask.shape, data.shape) + skf = model_selection.StratifiedKFold(n_splits=bcast_var[1], + shuffle=False) + accuracy = np.mean(model_selection.cross_val_score(clf, data, + y=bcast_var[0], + cv=skf, + n_jobs=1)) + return accuracy + + class MVPAVoxelSelector: """Activity-based voxel selection component of FCMA @@ -94,23 +112,12 @@ def run(self, clf): 'running activity-based voxel selection via Searchlight' ) self.sl.distribute([self.data], self.mask) - self.sl.broadcast((self.labels, self.num_folds)) + self.sl.broadcast((self.labels, self.num_folds, clf)) if rank == 0: logger.info( 'data preparation done' ) - # Searchlight kernel function - def _sfn(l, mask, myrad, bcast_var): - data = l[0][mask, :].T - # print(l[0].shape, mask.shape, data.shape) - skf = model_selection.StratifiedKFold(n_splits=bcast_var[1], - shuffle=False) - accuracy = np.mean(model_selection.cross_val_score(clf, data, - y=bcast_var[0], - cv=skf, - n_jobs=1)) - return accuracy # obtain a 3D array with accuracy numbers result_volume = self.sl.run_searchlight(_sfn) # get result tuple list from the volume diff --git a/brainiak/fcma/voxelselector.py b/brainiak/fcma/voxelselector.py index 9cfd0693a..25adc5202 100644 --- a/brainiak/fcma/voxelselector.py +++ b/brainiak/fcma/voxelselector.py @@ -28,7 +28,7 @@ from . import fcma_extension # type: ignore from . import cython_blas as blas # type: ignore import logging -import multiprocess +import multiprocessing logger = logging.getLogger(__name__) @@ -37,6 +37,21 @@ ] +def _cross_validation_for_one_voxel(clf, vid, num_folds, subject_data, labels): + """Score classifier on data using cross validation.""" + # no shuffling in cv + skf = model_selection.StratifiedKFold(n_splits=num_folds, + shuffle=False) + scores = model_selection.cross_val_score(clf, subject_data, + y=labels, + cv=skf, n_jobs=1) + logger.debug( + 'cross validation for voxel %d is done' % + vid + ) + return (vid, scores.mean()) + + class VoxelSelector: """Correlation-based voxel selection component of FCMA. @@ -415,31 +430,17 @@ def _do_cross_validation(self, clf, data, task): """ time1 = time.time() - def _cross_validation_for_one_voxel(vid, num_folds, - subject_data, labels): - # no shuffling in cv - skf = model_selection.StratifiedKFold(n_splits=num_folds, - shuffle=False) - scores = model_selection.cross_val_score(clf, subject_data, - y=labels, - cv=skf, n_jobs=1) - logger.debug( - 'cross validation for voxel %d is done' % - vid - ) - return (vid, scores.mean()) - if isinstance(clf, sklearn.svm.SVC) and clf.kernel == 'precomputed': - inlist = [(i + task[0], self.num_folds, data[i, :, :], + inlist = [(clf, i + task[0], self.num_folds, data[i, :, :], self.labels) for i in range(task[1])] - with multiprocess.Pool(self.process_num) as pool: + with multiprocessing.Pool(self.process_num) as pool: results = list(pool.starmap(_cross_validation_for_one_voxel, inlist)) else: results = [] for i in range(task[1]): - result = _cross_validation_for_one_voxel(i + task[0], + result = _cross_validation_for_one_voxel(clf, i + task[0], self.num_folds, data[i, :, :], self.labels) diff --git a/brainiak/searchlight/searchlight.py b/brainiak/searchlight/searchlight.py index 0f009da44..034d838ba 100644 --- a/brainiak/searchlight/searchlight.py +++ b/brainiak/searchlight/searchlight.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from multiprocessing import Pool import numpy as np from mpi4py import MPI import sys @@ -337,7 +338,7 @@ def broadcast(self, bcast_var): self.bcast_var = self.comm.bcast(bcast_var) - def run_block_function(self, block_fn): + def run_block_function(self, block_fn, extra_block_fn_params=None): """Perform a function for each block in a volume. Parameters @@ -357,12 +358,17 @@ def run_block_function(self, block_fn): bcast_var: shared data which is broadcast to all processes + extra_params: extra parameters + Returns 3D array which is the same size as the mask input with padding removed + extra_block_fn_params: tuple + Extra parameters to pass to the block function + """ rank = self.comm.rank @@ -370,7 +376,8 @@ def run_block_function(self, block_fn): # Apply searchlight local_outputs = [(mypt[0], block_fn([d[idx] for d in self.subproblems], - self.submasks[idx], self.sl_rad, self.bcast_var)) + self.submasks[idx], self.sl_rad, self.bcast_var, + extra_block_fn_params)) for (idx, mypt) in enumerate(self.blocks)] # Collect results @@ -400,6 +407,9 @@ def run_searchlight(self, voxel_fn, pool_size=None): voxel_fn: function to apply at each voxel + Must be `serializeable using pickle + `_. + Parameters subj: list of 4D arrays containing subset of subject data @@ -430,40 +440,62 @@ def run_searchlight(self, voxel_fn, pool_size=None): """ - def _singlenode_searchlight(l, msk, mysl_rad, bcast_var): - - outmat = np.empty(msk.shape, dtype=np.object)[mysl_rad:-mysl_rad, - mysl_rad:-mysl_rad, - mysl_rad:-mysl_rad] - - import pathos.multiprocessing - inlist = [([ll[i:i+2*mysl_rad+1, - j:j+2*mysl_rad+1, - k:k+2*mysl_rad+1, - :] - for ll in l], - msk[i:i+2*mysl_rad+1, - j:j+2*mysl_rad+1, - k:k+2*mysl_rad+1] * self.shape, - mysl_rad, - bcast_var) - if msk[i+mysl_rad, j+mysl_rad, k+mysl_rad] else None - for i in range(0, outmat.shape[0]) - for j in range(0, outmat.shape[1]) - for k in range(0, outmat.shape[2])] - outlist = list(pathos.multiprocessing.ProcessingPool(pool_size) - .map(lambda x: - voxel_fn(x[0], x[1], x[2], x[3]) - if x is not None else None, inlist)) - - cnt = 0 - for i in range(0, outmat.shape[0]): - for j in range(0, outmat.shape[1]): - for k in range(0, outmat.shape[2]): - if outlist[cnt] is not None: - outmat[i, j, k] = outlist[cnt] - cnt = cnt + 1 - - return outmat - - return self.run_block_function(_singlenode_searchlight) + # Reuse a single `Pool` for all blocks of the MPI process. + with Pool(pool_size) as pool: + extra_block_fn_params = (pool, voxel_fn, self.shape) + block_fn_result = self.run_block_function(_singlenode_searchlight, + extra_block_fn_params) + return block_fn_result + + +def _singlenode_searchlight(l, msk, mysl_rad, bcast_var, extra_params): + """Run searchlight function on block data in parallel. + + `extra_params` contains: + + - `Pool`. + - Searchlight function. + - `Shape` mask. + """ + + pool = extra_params[0] + voxel_fn = extra_params[1] + shape_mask = extra_params[2] + outmat = np.empty(msk.shape, dtype=np.object)[mysl_rad:-mysl_rad, + mysl_rad:-mysl_rad, + mysl_rad:-mysl_rad] + + inlist = [] + inlist_where = [] + where_idx = 0 + for i in range(0, outmat.shape[0]): + for j in range(0, outmat.shape[1]): + for k in range(0, outmat.shape[2]): + if msk[i+mysl_rad, j+mysl_rad, k+mysl_rad]: + inlist.append(( + [ll[i:i+2*mysl_rad+1, + j:j+2*mysl_rad+1, + k:k+2*mysl_rad+1, + :] + for ll in l], + msk[i:i+2*mysl_rad+1, + j:j+2*mysl_rad+1, + k:k+2*mysl_rad+1 + ] * shape_mask, + mysl_rad, + bcast_var)) + inlist_where.append(where_idx) + where_idx += 1 + outlist_all = np.array([None] * where_idx) + outlist = list(pool.starmap(voxel_fn, inlist)) + outlist_all[inlist_where] = outlist + + cnt = 0 + for i in range(0, outmat.shape[0]): + for j in range(0, outmat.shape[1]): + for k in range(0, outmat.shape[2]): + if outlist_all[cnt] is not None: + outmat[i, j, k] = outlist_all[cnt] + cnt = cnt + 1 + + return outmat diff --git a/setup.py b/setup.py index 07fb33343..ab8fc3d0c 100644 --- a/setup.py +++ b/setup.py @@ -130,7 +130,6 @@ def finalize_options(self): 'pymanopt', 'theano', 'pybind11>=1.7', - 'pathos', 'nibabel', 'typing', ], diff --git a/tests/searchlight/test_searchlight.py b/tests/searchlight/test_searchlight.py index 8d673edc4..e4ca187ab 100644 --- a/tests/searchlight/test_searchlight.py +++ b/tests/searchlight/test_searchlight.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from collections import namedtuple + import numpy as np from mpi4py import MPI import sys @@ -22,6 +24,10 @@ """Distributed Searchlight Test """ +def cube_sfn(l, msk, myrad, bcast_var): + if np.all(msk) and np.any(msk): + return 1.0 + return None def test_searchlight_with_cube(): sl = Searchlight(sl_rad=3) @@ -37,13 +43,8 @@ def test_searchlight_with_cube(): # Put a spot in the mask mask[10:17,10:17,10:17] = True - def sfn(l, msk, myrad, bcast_var): - if np.all(msk) and np.any(msk): - return 1.0 - return None - sl.distribute(data, mask) - global_outputs = sl.run_searchlight(sfn) + global_outputs = sl.run_searchlight(cube_sfn) if rank == 0: assert global_outputs[13,13,13] == 1.0 @@ -54,6 +55,12 @@ def sfn(l, msk, myrad, bcast_var): for k in range(global_outputs.shape[2]): assert global_outputs[i,j,k] == None +def diamond_sfn(l, msk, myrad, bcast_var): + assert not np.any(msk[~Diamond(3).mask_]) + if np.all(msk[Diamond(3).mask_]): + return 1.0 + return None + def test_searchlight_with_diamond(): sl = Searchlight(sl_rad=3, shape=Diamond) comm = MPI.COMM_WORLD @@ -68,14 +75,8 @@ def test_searchlight_with_diamond(): # Put a spot in the mask mask[10:17,10:17,10:17] = Diamond(3).mask_ - def sfn(l, msk, myrad, bcast_var): - assert not np.any(msk[~Diamond(3).mask_]) - if np.all(msk[Diamond(3).mask_]): - return 1.0 - return None - sl.distribute(data, mask) - global_outputs = sl.run_searchlight(sfn) + global_outputs = sl.run_searchlight(diamond_sfn) if rank == 0: assert global_outputs[13,13,13] == 1.0 @@ -86,12 +87,43 @@ def sfn(l, msk, myrad, bcast_var): for k in range(global_outputs.shape[2]): assert global_outputs[i,j,k] == None +MaskRadBcast = namedtuple("MaskRadBcast", "mask rad") def test_instantiate(): sl = Searchlight(sl_rad=5, max_blk_edge=10) assert sl + +def voxel_test_sfn(l, msk, myrad, bcast): + rad = bcast.rad + # Check each point + for subj in l: + for _tr in range(subj.shape[3]): + tr = subj[:, :, :, _tr] + midpt = tr[rad, rad, rad] + for d0 in range(tr.shape[0]): + for d1 in range(tr.shape[1]): + for d2 in range(tr.shape[2]): + assert np.array_equal(tr[d0, d1, d2] - midpt, + np.array([d0-rad, d1-rad, + d2-rad, 0])) + + # Determine midpoint + midpt = l[0][rad, rad, rad, 0] + midpt = (midpt[0], midpt[1], midpt[2]) + + for d0 in range(msk.shape[0]): + for d1 in range(msk.shape[1]): + for d2 in range(msk.shape[2]): + pt = (midpt[0] - rad + d0, midpt[1] - rad + d1, + midpt[2] - rad + d2) + assert bcast.mask[pt] == msk[d0, d1, d2] + + # Return midpoint + return midpt + + def test_correctness(): def voxel_test(data, mask, max_blk_edge, rad): @@ -111,35 +143,10 @@ def voxel_test(data, mask, max_blk_edge, rad): for d3 in range(dim2): subj[d1,d2,d3,tr] = np.array([d1, d2, d3, tr]) - def sfn(l,msk,myrad,bcast_var): - # Check each point - for subj in l: - for _tr in range(subj.shape[3]): - tr = subj[:,:,:,_tr] - midpt = tr[rad,rad,rad] - for d0 in range(tr.shape[0]): - for d1 in range(tr.shape[1]): - for d2 in range(tr.shape[2]): - assert np.array_equal(tr[d0,d1,d2] - midpt, np.array([d0-rad,d1-rad,d2-rad,0])) - - # Determine midpoint - midpt = l[0][rad,rad,rad,0] - midpt = (midpt[0], midpt[1], midpt[2]) - - for d0 in range(msk.shape[0]): - for d1 in range(msk.shape[1]): - for d2 in range(msk.shape[2]): - pt = (midpt[0] - rad + d0, midpt[1] - rad + d1, midpt[2] - rad + d2) - assert bcast_var[pt] == msk[d0,d1,d2] - - # Return midpoint - return midpt - - sl = Searchlight(sl_rad=rad, max_blk_edge=max_blk_edge) sl.distribute(data, mask) - sl.broadcast(mask) - global_outputs = sl.run_searchlight(sfn) + sl.broadcast(MaskRadBcast(mask, rad)) + global_outputs = sl.run_searchlight(voxel_test_sfn) if rank == 0: for d0 in range(rad, global_outputs.shape[0]-rad): @@ -167,7 +174,7 @@ def block_test(data, mask, max_blk_edge, rad): for d3 in range(dim2): subj[d1,d2,d3,tr] = np.array([ d1, d2, d3, tr]) - def sfn(l,msk,myrad,bcast_var): + def sfn(l,msk,myrad,bcast_var, extra_params): outmat = l[0][:,:,:,0] outmat[~msk] = None return outmat[rad:-rad,rad:-rad,rad:-rad]