mihaic · May 10, 2017 · Apr 27, 2017 · Apr 27, 2017 · Apr 28, 2017 · Apr 28, 2017
diff --git 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
@@ -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
@@ -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()