brainiak · mihaic · Mar 23, 2018 · Mar 22, 2018 · Mar 22, 2018 · Mar 22, 2018
diff --git a/brainiak/searchlight/searchlight.py b/brainiak/searchlight/searchlight.py
@@ -15,15 +15,15 @@
 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
 
 """Distributed Searchlight
 """
 
 __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
@@ -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")