diff --git a/brainiak/searchlight/searchlight.py b/brainiak/searchlight/searchlight.py index e5f7fdf9d..0f009da44 100644 --- a/brainiak/searchlight/searchlight.py +++ b/brainiak/searchlight/searchlight.py @@ -15,6 +15,7 @@ import numpy as np from mpi4py import MPI import sys +from scipy.spatial.distance import cityblock """Distributed Searchlight """ @@ -24,6 +25,80 @@ ] +class Shape: + """Shape + + Searchlight shape which is contained in a cube sized + (2*rad+1,2*rad+1,2*rad+1) + + Attributes + ---------- + + mask_ : a 3D boolean numpy array of size (2*rad+1,2*rad+1,2*rad+1) + which is set to True within the boundaries of the desired shape + """ + + def __init__(self, rad): + """Constructor + + Parameters + ---------- + + rad: radius, in voxels, of the sphere inscribed in the + searchlight cube, not counting the center voxel + + """ + self.rad = rad + + +class Cube(Shape): + """Cube + + Searchlight shape which is a cube of size (2*rad+1,2*rad+1,2*rad+1) + """ + def __init__(self, rad): + """Constructor + + Parameters + ---------- + + rad: radius, in voxels, of the sphere inscribed in the + searchlight cube, not counting the center voxel + + """ + super().__init__(rad) + self.rad = rad + self.mask_ = np.ones((2*rad+1, 2*rad+1, 2*rad+1), dtype=np.bool) + + +class Diamond(Shape): + """Diamond + + Searchlight shape which is a diamond + inscribed in a cube of size (2*rad+1,2*rad+1,2*rad+1). + Any location in the cube which has a Manhattan distance of less than rad + from the center point is set to True. + """ + def __init__(self, rad): + """Constructor + + Parameters + ---------- + + rad: radius, in voxels, of the sphere inscribed in the + searchlight cube, not counting the center voxel + + """ + super().__init__(rad) + self.mask_ = np.zeros((2*rad+1, 2*rad+1, 2*rad+1), dtype=np.bool) + for r1 in range(2*self.rad+1): + for r2 in range(2*self.rad+1): + for r3 in range(2*self.rad+1): + if(cityblock((r1, r2, r3), + (self.rad, self.rad, self.rad)) <= self.rad): + self.mask_[r1, r2, r3] = True + + class Searchlight: """Distributed Searchlight @@ -33,7 +108,7 @@ class Searchlight: Optionally, users can define a block function which runs over larger portions of the volume called blocks. """ - def __init__(self, sl_rad=1, max_blk_edge=10): + def __init__(self, sl_rad=1, max_blk_edge=10, shape=Cube): """Constructor Parameters @@ -44,10 +119,15 @@ def __init__(self, sl_rad=1, max_blk_edge=10): max_blk_edge: max edge length, in voxels, of the 3D block + shape: brainiak.searchlight.searchlight.Shape indicating the + shape in voxels of the searchlight region + """ self.sl_rad = sl_rad self.max_blk_edge = max_blk_edge self.comm = MPI.COMM_WORLD + self.shape = shape(sl_rad).mask_ + self.bcast_var = None def _get_ownership(self, data): """Determine on which rank each subject currently resides @@ -311,7 +391,9 @@ def run_block_function(self, block_fn): return outmat def run_searchlight(self, voxel_fn, pool_size=None): - """Perform a function at each active voxel + """Perform a function at each voxel which is set to True in the + user-provided mask. The mask passed to the searchlight function will be + further masked by the user-provided searchlight shape. Parameters ---------- @@ -337,6 +419,15 @@ def run_searchlight(self, voxel_fn, pool_size=None): pool_size: Number of parallel processes in shared memory process pool + Returns + ------- + + A volume which is the same size as the mask, however a number of voxels + equal to the searchlight radius has been removed from each border of + the volume. This volume contains the values returned from the + searchlight function at each voxel which was set to True in the mask, + and None elsewhere. + """ def _singlenode_searchlight(l, msk, mysl_rad, bcast_var): @@ -353,7 +444,7 @@ def _singlenode_searchlight(l, msk, mysl_rad, bcast_var): for ll in l], msk[i:i+2*mysl_rad+1, j:j+2*mysl_rad+1, - k:k+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 diff --git a/examples/searchlight/example_searchlight.py b/examples/searchlight/example_searchlight.py index 2da316b4f..ca29435f5 100644 --- a/examples/searchlight/example_searchlight.py +++ b/examples/searchlight/example_searchlight.py @@ -17,6 +17,7 @@ import sys from brainiak.searchlight.searchlight import Searchlight +from brainiak.searchlight.searchlight import Diamond """Distributed Searchlight Example example usage: mpirun -n 4 python3 example_searchlight.py @@ -66,7 +67,7 @@ data[pt[0]:pt[0]+kernel_dim,pt[1]:pt[1]+kernel_dim,pt[2]:pt[2]+kernel_dim,idx] -= kernel * weight # Create searchlight object -sl = Searchlight(sl_rad=1, max_blk_edge=5) +sl = Searchlight(sl_rad=1, max_blk_edge=5, shape=Diamond) # Distribute data to processes sl.distribute([data], mask) diff --git a/tests/searchlight/test_searchlight.py b/tests/searchlight/test_searchlight.py index 85a6e26ae..8d673edc4 100644 --- a/tests/searchlight/test_searchlight.py +++ b/tests/searchlight/test_searchlight.py @@ -17,10 +17,77 @@ import sys from brainiak.searchlight.searchlight import Searchlight +from brainiak.searchlight.searchlight import Diamond """Distributed Searchlight Test """ + +def test_searchlight_with_cube(): + sl = Searchlight(sl_rad=3) + comm = MPI.COMM_WORLD + rank = comm.rank + size = comm.size + dim0, dim1, dim2 = (50,50,50) + ntr = 30 + nsubj = 3 + mask = np.zeros((dim0,dim1,dim2), dtype=np.bool) + data = [np.empty((dim0,dim1,dim2,ntr), dtype=np.object) if i % size == rank else None for i in range(0, nsubj)] + + # 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) + + if rank == 0: + assert global_outputs[13,13,13] == 1.0 + global_outputs[13,13,13] = None + + for i in range(global_outputs.shape[0]): + for j in range(global_outputs.shape[1]): + for k in range(global_outputs.shape[2]): + assert global_outputs[i,j,k] == None + +def test_searchlight_with_diamond(): + sl = Searchlight(sl_rad=3, shape=Diamond) + comm = MPI.COMM_WORLD + rank = comm.rank + size = comm.size + dim0, dim1, dim2 = (50,50,50) + ntr = 30 + nsubj = 3 + mask = np.zeros((dim0,dim1,dim2), dtype=np.bool) + data = [np.empty((dim0,dim1,dim2,ntr), dtype=np.object) if i % size == rank else None for i in range(0, nsubj)] + + # 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) + + if rank == 0: + assert global_outputs[13,13,13] == 1.0 + global_outputs[13,13,13] = None + + for i in range(global_outputs.shape[0]): + for j in range(global_outputs.shape[1]): + for k in range(global_outputs.shape[2]): + assert global_outputs[i,j,k] == None + + + def test_instantiate(): sl = Searchlight(sl_rad=5, max_blk_edge=10) assert sl @@ -78,8 +145,8 @@ def sfn(l,msk,myrad,bcast_var): for d0 in range(rad, global_outputs.shape[0]-rad): for d1 in range(rad, global_outputs.shape[1]-rad): for d2 in range(rad, global_outputs.shape[2]-rad): - if mask[d0, d1, d2]: - assert np.array_equal(np.array(global_outputs[d0,d1,d2]), np.array([d0,d1,d2])) + if mask[d0, d1, d2]: + assert np.array_equal(np.array(global_outputs[d0,d1,d2]), np.array([d0,d1,d2])) def block_test(data, mask, max_blk_edge, rad): @@ -135,4 +202,3 @@ def do_test(dim0, dim1, dim2, ntr, nsubj, max_blk_edge, rad): do_test(dim0=7, dim1=5, dim2=9, ntr=5, nsubj=1, max_blk_edge=4, rad=2) do_test(dim0=7, dim1=5, dim2=9, ntr=5, nsubj=1, max_blk_edge=4, rad=3) -test_correctness()