diff --git a/brainiak/searchlight/searchlight.py b/brainiak/searchlight/searchlight.py index 651c9c988..c829311e9 100644 --- a/brainiak/searchlight/searchlight.py +++ b/brainiak/searchlight/searchlight.py @@ -15,7 +15,7 @@ from multiprocessing import Pool import numpy as np from mpi4py import MPI -from scipy.spatial.distance import cityblock +from scipy.spatial.distance import cityblock, euclidean from ..utils.utils import usable_cpu_count @@ -23,7 +23,7 @@ """ __all__ = [ - "Searchlight", "Shape", "Cube", "Diamond" + "Searchlight", "Shape", "Cube", "Diamond", "Ball" ] @@ -74,8 +74,8 @@ class Diamond(Shape): 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. + Any location in the cube which has a Manhattan distance of equal to or + less than rad from the center point is set to True. Parameters ---------- @@ -95,6 +95,32 @@ def __init__(self, rad): self.mask_[r1, r2, r3] = True +class Ball(Shape): + """Ball + + Searchlight shape which is a ball + inscribed in a cube of size (2*rad+1,2*rad+1,2*rad+1). + Any location in the cube which has a Euclidean distance of equal to or + less than rad from the center point is set to True. + + Parameters + ---------- + + rad: radius, in voxels, of the sphere inscribed in the + searchlight cube, not counting the center voxel + + """ + def __init__(self, rad): + 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(euclidean((r1, r2, r3), + (self.rad, self.rad, self.rad)) <= self.rad): + self.mask_[r1, r2, r3] = True + + class Searchlight: """Distributed Searchlight diff --git a/tests/searchlight/test_searchlight.py b/tests/searchlight/test_searchlight.py index 316456830..f6da405fc 100644 --- a/tests/searchlight/test_searchlight.py +++ b/tests/searchlight/test_searchlight.py @@ -18,7 +18,7 @@ from mpi4py import MPI from brainiak.searchlight.searchlight import Searchlight -from brainiak.searchlight.searchlight import Diamond +from brainiak.searchlight.searchlight import Diamond, Ball """Distributed Searchlight Test """ @@ -97,6 +97,45 @@ def test_searchlight_with_diamond(): assert global_outputs[i, j, k] is None +def ball_sfn(l, msk, myrad, bcast_var): + x, y, z = np.mgrid[-myrad:myrad+1, -myrad:myrad+1, -myrad:myrad+1] + correct_mask = np.square(x) + np.square(y) + np.square(z) <= myrad ** 2 + assert not np.any(msk[~Ball(3).mask_]) + if np.all(correct_mask == msk): + return 1.0 + return None + + +def test_searchlight_with_ball(): + sl = Searchlight(sl_rad=3, shape=Ball) + 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] = Ball(3).mask_ + + sl.distribute(data, mask) + global_outputs = sl.run_searchlight(ball_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] is None + + MaskRadBcast = namedtuple("MaskRadBcast", "mask rad")