From 0708bbbe07be310788ff8d45c1dd1c5b8d4400af Mon Sep 17 00:00:00 2001
From: Michael Anderson <michael.j.anderson@intel.com>
Date: Sun, 20 Nov 2016 12:38:13 -0800
Subject: [PATCH 01/11] Added searchlight

---
 brainiak/searchlight/__init__.py            |  14 ++
 brainiak/searchlight/example_fn.cpp         | 119 +++++++++
 brainiak/searchlight/searchlight.py         | 263 ++++++++++++++++++++
 examples/searchlight/example_searchlight.py |  98 ++++++++
 setup.py                                    |   5 +
 tests/searchlight/test_searchlight.py       | 135 ++++++++++
 6 files changed, 634 insertions(+)
 create mode 100644 brainiak/searchlight/__init__.py
 create mode 100644 brainiak/searchlight/example_fn.cpp
 create mode 100644 brainiak/searchlight/searchlight.py
 create mode 100644 examples/searchlight/example_searchlight.py
 create mode 100644 tests/searchlight/test_searchlight.py

diff --git a/brainiak/searchlight/__init__.py b/brainiak/searchlight/__init__.py
new file mode 100644
index 000000000..c06a3e14c
--- /dev/null
+++ b/brainiak/searchlight/__init__.py
@@ -0,0 +1,14 @@
+#  Copyright 2016 Intel Corporation
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+"""Distributed Searchlight"""
diff --git a/brainiak/searchlight/example_fn.cpp b/brainiak/searchlight/example_fn.cpp
new file mode 100644
index 000000000..7ec2f8ead
--- /dev/null
+++ b/brainiak/searchlight/example_fn.cpp
@@ -0,0 +1,119 @@
+// Copyright Intel Corporation 2016
+
+#include <iostream>
+#include <vector>
+#include <list>
+#include <pybind11/pybind11.h>
+#include <pybind11/numpy.h>
+#include <pybind11/stl.h>
+namespace py = pybind11;
+
+inline size_t getIdx3(const std::vector<size_t> & pt,
+                  const std::vector<size_t> & strides)
+{
+  return pt[0]*strides[0] + pt[1]*strides[1] + pt[2]*strides[2];
+}
+
+inline size_t getIdx4(const std::vector<size_t> & pt,
+                  const std::vector<size_t> & strides)
+{
+  return pt[0]*strides[0] + pt[1]*strides[1] + pt[2]*strides[2] + pt[3]*strides[3];
+}
+
+py::array_t<double> yah(std::vector< py::array_t<double, py::array::c_style | py::array::forcecast> > data,
+         py::array_t<double, py::array::c_style | py::array::forcecast> & _mask,
+         size_t rad,
+         py::object bcast_var)
+{
+  auto _mask_buf = _mask.request();
+  if(_mask_buf.ndim != 3)
+  {
+    throw std::runtime_error("Mask must be a 3D numpy array");
+  }
+  double * mask = (double*) _mask_buf.ptr;
+
+  // If region is too small, return empty array
+  if(_mask_buf.shape[0] <= 2*rad ||
+     _mask_buf.shape[1] <= 2*rad ||
+     _mask_buf.shape[2] <= 2*rad)
+  {
+    std::vector<size_t> outshape = {0,0,0};
+    std::vector<size_t> outstrides = {sizeof(double),sizeof(double),sizeof(double)};
+    return py::array(py::buffer_info(NULL, sizeof(double), py::format_descriptor<double>::value, 3, outshape, outstrides));
+  }
+
+  // Create output array
+  std::vector<size_t> outshape = {_mask_buf.shape[0]-2*rad, 
+                                             _mask_buf.shape[1]-2*rad,
+                                             _mask_buf.shape[2]-2*rad};
+
+  std::vector<size_t> outstrides = {outshape[2]*outshape[1]*sizeof(double),
+                                               outshape[2]*sizeof(double),
+                                               sizeof(double)};
+  std::vector<double> outmat(outshape[0]*outshape[1]*outshape[2]);
+
+  std::vector<double*> subj_data;
+  std::vector<std::vector<size_t> > subj_strides;
+  std::vector<std::vector<size_t> > subj_shapes;
+
+  for(size_t subj = 0 ; subj < data.size() ; subj++)
+  {
+    auto subj_buf = data[subj].request();
+    subj_data.push_back((double*)subj_buf.ptr);
+    subj_strides.push_back(subj_buf.strides);
+    subj_shapes.push_back(subj_buf.shape);
+  }
+  #pragma omp parallel for schedule(dynamic) collapse(3)
+  for(size_t i = rad ; i < _mask_buf.shape[0]-rad ; i++)
+  {
+    for(size_t j = rad ; j < _mask_buf.shape[1]-rad ; j++)
+    {
+      for(size_t k = rad ; k < _mask_buf.shape[2]-rad ; k++)
+      {
+        if(mask[getIdx3({i,j,k},_mask_buf.strides)/sizeof(double)] > 0) 
+        {
+          double sum = 0.0;
+          size_t niter = 0;
+          for(size_t subj = 0 ; subj < data.size() ; subj++)
+          {
+            size_t base_pt = getIdx4({0,i,j,k},subj_strides[subj])/sizeof(double);
+            size_t cnt = 0;
+            for(size_t zz = 0 ; zz < subj_shapes[subj][0] ; zz++)
+            {
+              for(size_t ii = i-rad ; ii < i+rad+1 ; ii++)
+              {
+                for(size_t jj = j-rad ; jj < j+rad+1 ; jj++)
+                {
+                  for(size_t kk = k-rad ; kk < k+rad+1 ; kk++)
+                  {
+                    sum += subj_data[subj][(zz*subj_strides[subj][0] +
+                                           ii*subj_strides[subj][1] + 
+                                           jj*subj_strides[subj][2] + 
+                                           kk*subj_strides[subj][3]) / sizeof(double)];
+                    niter++;
+                  }
+                }
+              }
+            }
+          }
+          outmat[getIdx3({i-rad,j-rad,k-rad},outstrides)/sizeof(double)] = sum;
+        }
+      }
+    }
+  }
+  return py::array(py::buffer_info(outmat.data(), sizeof(double), py::format_descriptor<double>::value, 3, outshape, outstrides));
+}
+
+
+PYBIND11_PLUGIN(example_fn) {
+    py::module m("example_fn");
+    m.def("yah", &yah, "yah");
+
+#ifdef VERSION_INFO
+    m.attr("__version__") = py::str(VERSION_INFO);
+#else
+    m.attr("__version__") = py::str("dev");
+#endif
+
+    return m.ptr();
+}
diff --git a/brainiak/searchlight/searchlight.py b/brainiak/searchlight/searchlight.py
new file mode 100644
index 000000000..be0123481
--- /dev/null
+++ b/brainiak/searchlight/searchlight.py
@@ -0,0 +1,263 @@
+#    Copyright 2016 Intel Corporation
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#             http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+import numpy as np
+from mpi4py import MPI
+import sys
+
+"""Distributed Searchlight
+"""
+
+__all__ = [
+        "Searchlight",
+]
+
+
+class Searchlight:
+    def __init__(self, rad=1, region_size=10):
+        """Constructor
+
+        Parameters
+        ----------
+
+        rad: number of voxels in border region
+
+        region_size: max dimension of a 3D region
+
+        """
+        self.rad = rad
+        self.region_size = region_size
+
+    def _get_ownership(self, l):
+        comm = MPI.COMM_WORLD
+        rank = comm.rank
+
+        B = [(rank, idx) for (idx, c) in enumerate(l) if c is not None]
+        C = comm.allreduce(B)
+        ownership = [None] * len(l)
+        for c in C:
+            ownership[c[1]] = c[0]
+        return ownership
+
+    def _get_points(self, mask):
+        outerblk = self.region_size + 2*self.rad
+        return [((i, j, k),
+                mask[i:i+outerblk,
+                j:j+outerblk,
+                k:k+outerblk].shape)
+                for i in range(0, mask.shape[0], self.region_size)
+                for j in range(0, mask.shape[1], self.region_size)
+                for k in range(0, mask.shape[2], self.region_size)
+                if np.sum(
+                mask[i+self.rad:i+self.rad+self.region_size,
+                     j+self.rad:j+self.rad+self.region_size,
+                     k+self.rad:k+self.rad+self.region_size]) > 0]
+
+    def _subarray(self, mat3d, pts):
+        (pt, sz) = pts
+        if len(mat3d.shape) == 3:
+            return mat3d[pt[0]:pt[0]+sz[0],
+                         pt[1]:pt[1]+sz[1],
+                         pt[2]:pt[2]+sz[2]].copy()
+        elif len(mat3d.shape) == 4:
+            return mat3d[:,
+                         pt[0]:pt[0]+sz[0],
+                         pt[1]:pt[1]+sz[1],
+                         pt[2]:pt[2]+sz[2]].copy()
+        else:
+            print('Error num dims')
+            sys.exit(0)
+
+    def _split_volume(self, mat3d, pts):
+        return [self._subarray(mat3d, pt) for pt in pts]
+
+    def _scatter_list(self, insubj, pts, owner):
+        comm = MPI.COMM_WORLD
+        size = comm.size
+        subject_submatrices = []
+
+        # For each submatrix
+        for idx in range(0, len(pts), size):
+            padded = None
+            extra = max(0, idx+size - len(pts))
+
+            # Pad with "None" so scatter can go to all processes
+            if insubj is not None:
+                padded = insubj[idx:idx+size]
+                if extra > 0:
+                    padded = padded + [None]*extra
+
+            # Scatter submatrices to all processes
+            mytrans = comm.scatter(padded, root=owner)
+
+            # Contribute submatrix to subject list
+            if mytrans is not None:
+                subject_submatrices += [mytrans]
+
+        return subject_submatrices
+
+    def distribute(self, subj, mask):
+        """Distribute data to processes
+
+        Parameters
+        ----------
+
+        subj: list of 4D arrays containing subject data
+
+        mask: 3D array with nonzero values at active vertices
+
+        """
+        self.mask = mask
+        comm = MPI.COMM_WORLD
+        rank = comm.rank
+
+        # Get/set ownership
+        ownership = self._get_ownership(subj)
+        full_pts = self._get_points(mask) if rank == 0 else None
+        full_pts = comm.bcast(full_pts)
+
+        # Divide data and mask
+        splitsubj = [self._split_volume(s, full_pts)
+                     if s is not None else None
+                     for s in subj]
+        submasks = self._split_volume(mask, full_pts)
+
+        # Scatter points, data, and mask
+        self.pts = self._scatter_list(full_pts, full_pts, 0)
+        self.submasks = self._scatter_list(submasks, full_pts, 0)
+        self.subproblems = [self._scatter_list(s, full_pts, ownership[s_idx])
+                            for (s_idx, s) in enumerate(splitsubj)]
+
+    def searchlight_region(self, region_fn, bcast_var):
+        """Perform a function for each region in a volume.
+
+        Parameters
+        ----------
+
+        region_fn: function to apply to each region:
+
+                Paramters
+                ---------
+
+                subj: list of 4D arrays containing subset of subject data
+
+                mask: 3D array containing subset of mask data
+
+                rad: number of voxels in border region
+
+                bcast_var: shared data which is broadcast to all processes
+
+
+                Returns
+                -------
+
+                3D array which is the same size as the mask
+                input with padding removed
+
+        bcast_var:    shared data which is broadcast to all processes
+
+        """
+        comm = MPI.COMM_WORLD
+        rank = comm.rank
+        bcast_var = comm.bcast(bcast_var)
+
+        # Apply searchlight
+        local_outputs = [(mypt[0],
+                          region_fn([d[idx] for d in self.subproblems],
+                          self.submasks[idx], self.rad, bcast_var))
+                         for (idx, mypt) in enumerate(self.pts)]
+
+        # Collect results
+        global_outputs = comm.gather(local_outputs)
+
+        # Coalesce results
+        outmat = np.empty(self.mask.shape, dtype=np.object)
+        if rank == 0:
+            for go_rank in global_outputs:
+                for (pt, mat) in go_rank:
+                    for i in range(0, mat.shape[0]):
+                        for j in range(0, mat.shape[1]):
+                            for k in range(0, mat.shape[2]):
+                                outmat[pt[0]+self.rad+i,
+                                       pt[1]+self.rad+j,
+                                       pt[2]+self.rad+k] = mat[i, j, k]
+
+        return outmat
+
+    def searchlight_voxel(self, voxel_fn, bcast_var):
+        """Perform a function at each active voxel
+
+        Parameters
+        ----------
+
+        voxel_fn: function to apply at each voxel
+
+                Paramters
+                ---------
+
+                subj: list of 4D arrays containing subset of subject data
+
+                mask: 3D array containing subset of mask data
+
+                rad: number of voxels in border region
+
+                bcast_var: shared data which is broadcast to all processes
+
+
+                Returns
+                -------
+
+                Value of any pickle-able type
+
+        bcast_var:    shared data which is broadcast to all processes
+
+        """
+
+        def _singlenode_searchlight(l, msk, myrad, bcast_var):
+
+            outmat = np.empty(msk.shape, dtype=np.object)[myrad:-myrad,
+                                                          myrad:-myrad,
+                                                          myrad:-myrad]
+
+            import pathos.multiprocessing
+            inlist = [([ll[:,
+                           i:i+2*myrad+1,
+                           j:j+2*myrad+1,
+                           k:k+2*myrad+1]
+                        for ll in l],
+                       msk[i:i+2*myrad+1,
+                           j:j+2*myrad+1,
+                           k:k+2*myrad+1],
+                       myrad,
+                       bcast_var)
+                      if msk[i+myrad, j+myrad, k+myrad] 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()
+                           .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.searchlight_region(_singlenode_searchlight, bcast_var)
diff --git a/examples/searchlight/example_searchlight.py b/examples/searchlight/example_searchlight.py
new file mode 100644
index 000000000..2bf598e77
--- /dev/null
+++ b/examples/searchlight/example_searchlight.py
@@ -0,0 +1,98 @@
+#  Copyright 2016 Intel Corporation
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+import numpy as np
+from mpi4py import MPI
+import sys
+
+from brainiak.searchlight import example_fn
+from brainiak.searchlight.searchlight import Searchlight
+
+"""Distributed Searchlight Example
+"""
+
+comm = MPI.COMM_WORLD
+rank = comm.rank
+size = comm.size
+
+# Dataset size parameters
+dim = 40
+ntr = 400
+maskrad = 15
+
+# Predictive point parameters
+pt = (23,23,23)
+kernel_dim = 5
+weight = 1
+
+# Generate data
+data = np.random.random((ntr,dim,dim,dim)) if rank == 0 else None
+mask = np.zeros((dim,dim,dim), dtype=np.bool)
+for i in range(dim):
+  for j in range(dim):
+    for k in range(dim):
+      dist = np.sqrt(((dim/2)-i)**2 + ((dim/2)-j)**2 + ((dim/2)-k)**2)
+      if(dist < maskrad):
+        mask[i,j,k] = 1
+
+# Generate labels
+labels = np.random.choice([True, False], (ntr,)) if rank == 0 else None
+
+# Inject predictive region in random data
+if rank == 0:
+  kernel = np.zeros((kernel_dim,kernel_dim,kernel_dim))
+  for i in range(kernel_dim):
+    for j in range(kernel_dim):
+      for k in range(kernel_dim):
+        arr = np.array([i-(kernel_dim/2),j-(kernel_dim/2),k-(kernel_dim/2)])
+        kernel [i,j,k] = np.exp(-np.dot(arr.T,arr))
+  kernel = kernel / np.sum(kernel)
+
+  for (idx, l) in enumerate(labels):
+    if l:
+      data[idx][pt[0]:pt[0]+kernel_dim,pt[1]:pt[1]+kernel_dim,pt[2]:pt[2]+kernel_dim] += kernel * weight
+    else:
+      data[idx][pt[0]:pt[0]+kernel_dim,pt[1]:pt[1]+kernel_dim,pt[2]:pt[2]+kernel_dim] -= kernel * weight
+
+# Create searchlight object
+sl = Searchlight(rad=1, region_size=5)
+
+# Distribute data to processes
+sl.distribute([data], mask)
+
+# Define voxel function
+def sfn(l, msk, myrad, bcast_var):
+  import sklearn.svm
+  import sklearn.model_selection
+  classifier = sklearn.svm.SVC()
+  data = l[0][:,msk]
+  return np.mean(sklearn.model_selection.cross_val_score(classifier, data, bcast_var,n_jobs=1))
+
+# Run searchlight
+global_outputs = sl.searchlight_voxel(sfn, labels)
+
+# Visualize result
+if rank == 0:
+  print(global_outputs)
+  global_outputs = np.array(global_outputs, dtype=np.float)
+  import matplotlib.pyplot as plt
+  for (cnt, img) in enumerate(global_outputs):
+    plt.imshow(img,cmap='hot',vmin=0,vmax=1)
+    plt.savefig('img' + str(cnt) + '.png')
+    plt.clf()
+
+
+
+
+
diff --git a/setup.py b/setup.py
index 02cb347b7..6dbe00541 100644
--- a/setup.py
+++ b/setup.py
@@ -24,6 +24,10 @@
         'brainiak.factoranalysis.tfa_extension',
         ['brainiak/factoranalysis/tfa_extension.cpp'],
     ),
+    Extension(
+        'brainiak.searchlight.example_fn',
+        ['brainiak/searchlight/example_fn.cpp'],
+    ),
     Extension(
         'brainiak.fcma.fcma_extension',
         ['brainiak/fcma/src/fcma_extension.cc'],
@@ -129,6 +133,7 @@ def finalize_options(self):
         'pymanopt',
         'theano',
         'pybind11>=1.7',
+        'pathos',
     ],
     author='Princeton Neuroscience Institute and Intel Corporation',
     author_email='bryn.keller@intel.com',
diff --git a/tests/searchlight/test_searchlight.py b/tests/searchlight/test_searchlight.py
new file mode 100644
index 000000000..5a7299fc1
--- /dev/null
+++ b/tests/searchlight/test_searchlight.py
@@ -0,0 +1,135 @@
+#  Copyright 2016 Intel Corporation
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+import numpy as np
+from mpi4py import MPI
+import sys
+
+from brainiak.searchlight import example_fn
+from brainiak.searchlight.searchlight import Searchlight
+
+"""Distributed Searchlight Test
+"""
+
+def test_instantiate():
+  sl = Searchlight(rad=5, region_size=10)
+  assert sl
+
+def test_correctness():
+  def voxel_test(data, mask, region_size, rad):
+    
+    comm = MPI.COMM_WORLD
+    rank = comm.rank
+    size = comm.size
+    
+    nsubj = len(data)
+    (dim0, dim1, dim2) = mask.shape
+    
+    # Initialize dataset with known pattern
+    for subj in data:
+      if subj is not None:
+        for tr in range(subj.shape[0]):
+          for d1 in range(dim0):
+            for d2 in range(dim1):
+              for d3 in range(dim2): 
+                subj[tr,d1,d2,d3] = np.array([tr, d1, d2, d3])
+    
+    def sfn(l,msk,myrad,bcast_var):
+      # Check each point
+      for subj in l:
+        for tr in subj:
+          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([0,d0-rad,d1-rad,d2-rad]))
+    
+      # Determine midpoint
+      midpt = l[0][0,rad,rad,rad]
+      midpt = (midpt[1], midpt[2], midpt[3])
+    
+      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(rad=rad, region_size=region_size)
+    sl.distribute(data, mask)
+    global_outputs = sl.searchlight_voxel(sfn, mask)
+  
+    if rank == 0:
+      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]))
+  
+  
+  def region_test(data, mask, region_size, rad):
+    
+    comm = MPI.COMM_WORLD
+    rank = comm.rank
+    size = comm.size
+    
+    nsubj = len(data)
+    (dim0, dim1, dim2) = mask.shape
+    
+    # Initialize dataset with known pattern
+    for subj in data:
+      if subj is not None:
+        for tr in range(subj.shape[0]):
+          for d1 in range(dim0):
+            for d2 in range(dim1):
+              for d3 in range(dim2): 
+                subj[tr,d1,d2,d3] = np.array([tr, d1, d2, d3])
+    
+    def sfn(l,msk,myrad,bcast_var):
+      outmat = l[0][0,:,:,:] 
+      outmat[~msk] = None
+      return outmat[rad:-rad,rad:-rad,rad:-rad] 
+    
+    sl = Searchlight(rad=rad, region_size=region_size)
+    sl.distribute(data, mask)
+    global_outputs = sl.searchlight_region(sfn, mask)
+  
+    if rank == 0:
+      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([0,d0,d1,d2]))
+  
+  # Create dataset
+  def do_test(dim0, dim1, dim2, ntr, nsubj, region_size, rad):
+    comm = MPI.COMM_WORLD
+    rank = comm.rank
+    size = comm.size
+    mask = np.random.choice([True, False], (dim0,dim1,dim2))
+    data = [np.empty((ntr,dim0,dim1,dim2), dtype=np.object) if i % size == rank else None for i in range(0, nsubj)]
+    voxel_test(data, mask, region_size, rad)
+    region_test(data, mask, region_size, rad)
+  
+  
+  do_test(dim0=7, dim1=5, dim2=9, ntr=5, nsubj=1, region_size=4, rad=1)
+  do_test(dim0=7, dim1=5, dim2=9, ntr=5, nsubj=5, region_size=4, rad=1)
+  do_test(dim0=1, dim1=5, dim2=9, ntr=5, nsubj=5, region_size=4, rad=1)
+  do_test(dim0=0, dim1=10, dim2=8, ntr=5, nsubj=5, region_size=4, rad=1)
+  do_test(dim0=7, dim1=5, dim2=9, ntr=5, nsubj=1, region_size=4, rad=2)
+  do_test(dim0=7, dim1=5, dim2=9, ntr=5, nsubj=1, region_size=4, rad=3)
+

From bb7523559590402a138b9176f2feaaa6729f3a39 Mon Sep 17 00:00:00 2001
From: Michael Anderson <michael.j.anderson@intel.com>
Date: Sun, 20 Nov 2016 12:39:58 -0800
Subject: [PATCH 02/11] removed cpp function

---
 brainiak/searchlight/example_fn.cpp | 119 ----------------------------
 1 file changed, 119 deletions(-)
 delete mode 100644 brainiak/searchlight/example_fn.cpp

diff --git a/brainiak/searchlight/example_fn.cpp b/brainiak/searchlight/example_fn.cpp
deleted file mode 100644
index 7ec2f8ead..000000000
--- a/brainiak/searchlight/example_fn.cpp
+++ /dev/null
@@ -1,119 +0,0 @@
-// Copyright Intel Corporation 2016
-
-#include <iostream>
-#include <vector>
-#include <list>
-#include <pybind11/pybind11.h>
-#include <pybind11/numpy.h>
-#include <pybind11/stl.h>
-namespace py = pybind11;
-
-inline size_t getIdx3(const std::vector<size_t> & pt,
-                  const std::vector<size_t> & strides)
-{
-  return pt[0]*strides[0] + pt[1]*strides[1] + pt[2]*strides[2];
-}
-
-inline size_t getIdx4(const std::vector<size_t> & pt,
-                  const std::vector<size_t> & strides)
-{
-  return pt[0]*strides[0] + pt[1]*strides[1] + pt[2]*strides[2] + pt[3]*strides[3];
-}
-
-py::array_t<double> yah(std::vector< py::array_t<double, py::array::c_style | py::array::forcecast> > data,
-         py::array_t<double, py::array::c_style | py::array::forcecast> & _mask,
-         size_t rad,
-         py::object bcast_var)
-{
-  auto _mask_buf = _mask.request();
-  if(_mask_buf.ndim != 3)
-  {
-    throw std::runtime_error("Mask must be a 3D numpy array");
-  }
-  double * mask = (double*) _mask_buf.ptr;
-
-  // If region is too small, return empty array
-  if(_mask_buf.shape[0] <= 2*rad ||
-     _mask_buf.shape[1] <= 2*rad ||
-     _mask_buf.shape[2] <= 2*rad)
-  {
-    std::vector<size_t> outshape = {0,0,0};
-    std::vector<size_t> outstrides = {sizeof(double),sizeof(double),sizeof(double)};
-    return py::array(py::buffer_info(NULL, sizeof(double), py::format_descriptor<double>::value, 3, outshape, outstrides));
-  }
-
-  // Create output array
-  std::vector<size_t> outshape = {_mask_buf.shape[0]-2*rad, 
-                                             _mask_buf.shape[1]-2*rad,
-                                             _mask_buf.shape[2]-2*rad};
-
-  std::vector<size_t> outstrides = {outshape[2]*outshape[1]*sizeof(double),
-                                               outshape[2]*sizeof(double),
-                                               sizeof(double)};
-  std::vector<double> outmat(outshape[0]*outshape[1]*outshape[2]);
-
-  std::vector<double*> subj_data;
-  std::vector<std::vector<size_t> > subj_strides;
-  std::vector<std::vector<size_t> > subj_shapes;
-
-  for(size_t subj = 0 ; subj < data.size() ; subj++)
-  {
-    auto subj_buf = data[subj].request();
-    subj_data.push_back((double*)subj_buf.ptr);
-    subj_strides.push_back(subj_buf.strides);
-    subj_shapes.push_back(subj_buf.shape);
-  }
-  #pragma omp parallel for schedule(dynamic) collapse(3)
-  for(size_t i = rad ; i < _mask_buf.shape[0]-rad ; i++)
-  {
-    for(size_t j = rad ; j < _mask_buf.shape[1]-rad ; j++)
-    {
-      for(size_t k = rad ; k < _mask_buf.shape[2]-rad ; k++)
-      {
-        if(mask[getIdx3({i,j,k},_mask_buf.strides)/sizeof(double)] > 0) 
-        {
-          double sum = 0.0;
-          size_t niter = 0;
-          for(size_t subj = 0 ; subj < data.size() ; subj++)
-          {
-            size_t base_pt = getIdx4({0,i,j,k},subj_strides[subj])/sizeof(double);
-            size_t cnt = 0;
-            for(size_t zz = 0 ; zz < subj_shapes[subj][0] ; zz++)
-            {
-              for(size_t ii = i-rad ; ii < i+rad+1 ; ii++)
-              {
-                for(size_t jj = j-rad ; jj < j+rad+1 ; jj++)
-                {
-                  for(size_t kk = k-rad ; kk < k+rad+1 ; kk++)
-                  {
-                    sum += subj_data[subj][(zz*subj_strides[subj][0] +
-                                           ii*subj_strides[subj][1] + 
-                                           jj*subj_strides[subj][2] + 
-                                           kk*subj_strides[subj][3]) / sizeof(double)];
-                    niter++;
-                  }
-                }
-              }
-            }
-          }
-          outmat[getIdx3({i-rad,j-rad,k-rad},outstrides)/sizeof(double)] = sum;
-        }
-      }
-    }
-  }
-  return py::array(py::buffer_info(outmat.data(), sizeof(double), py::format_descriptor<double>::value, 3, outshape, outstrides));
-}
-
-
-PYBIND11_PLUGIN(example_fn) {
-    py::module m("example_fn");
-    m.def("yah", &yah, "yah");
-
-#ifdef VERSION_INFO
-    m.attr("__version__") = py::str(VERSION_INFO);
-#else
-    m.attr("__version__") = py::str("dev");
-#endif
-
-    return m.ptr();
-}

From 5954c59a478edb61a1bfae9b5a56082c62c08a6b Mon Sep 17 00:00:00 2001
From: Michael Anderson <michael.j.anderson@intel.com>
Date: Sun, 20 Nov 2016 12:42:45 -0800
Subject: [PATCH 03/11] removed example cpp function

---
 setup.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/setup.py b/setup.py
index 6dbe00541..33552e9c9 100644
--- a/setup.py
+++ b/setup.py
@@ -24,10 +24,6 @@
         'brainiak.factoranalysis.tfa_extension',
         ['brainiak/factoranalysis/tfa_extension.cpp'],
     ),
-    Extension(
-        'brainiak.searchlight.example_fn',
-        ['brainiak/searchlight/example_fn.cpp'],
-    ),
     Extension(
         'brainiak.fcma.fcma_extension',
         ['brainiak/fcma/src/fcma_extension.cc'],

From 38c618c9933e1bc72d560151af7b06bd8a1dea2f Mon Sep 17 00:00:00 2001
From: Michael Anderson <michael.j.anderson@intel.com>
Date: Sun, 20 Nov 2016 12:57:48 -0800
Subject: [PATCH 04/11] fixed comment for pr-check

---
 brainiak/searchlight/searchlight.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/brainiak/searchlight/searchlight.py b/brainiak/searchlight/searchlight.py
index be0123481..e20822c2e 100644
--- a/brainiak/searchlight/searchlight.py
+++ b/brainiak/searchlight/searchlight.py
@@ -147,8 +147,7 @@ def searchlight_region(self, region_fn, bcast_var):
 
         region_fn: function to apply to each region:
 
-                Paramters
-                ---------
+                Parameters
 
                 subj: list of 4D arrays containing subset of subject data
 
@@ -160,7 +159,6 @@ def searchlight_region(self, region_fn, bcast_var):
 
 
                 Returns
-                -------
 
                 3D array which is the same size as the mask
                 input with padding removed
@@ -203,8 +201,7 @@ def searchlight_voxel(self, voxel_fn, bcast_var):
 
         voxel_fn: function to apply at each voxel
 
-                Paramters
-                ---------
+                Parameters
 
                 subj: list of 4D arrays containing subset of subject data
 
@@ -216,7 +213,6 @@ def searchlight_voxel(self, voxel_fn, bcast_var):
 
 
                 Returns
-                -------
 
                 Value of any pickle-able type
 

From 8276e8ca6406ef5da39f8980a962d8ec1846771d Mon Sep 17 00:00:00 2001
From: Michael Anderson <michael.j.anderson@intel.com>
Date: Tue, 22 Nov 2016 07:49:37 -0800
Subject: [PATCH 05/11] added documentation

---
 brainiak/searchlight/searchlight.py         | 239 ++++++++++++++------
 examples/searchlight/example_searchlight.py |   5 +-
 tests/searchlight/test_searchlight.py       |  35 +--
 3 files changed, 197 insertions(+), 82 deletions(-)

diff --git a/brainiak/searchlight/searchlight.py b/brainiak/searchlight/searchlight.py
index e20822c2e..6a81336b9 100644
--- a/brainiak/searchlight/searchlight.py
+++ b/brainiak/searchlight/searchlight.py
@@ -25,53 +25,110 @@
 
 
 class Searchlight:
-    def __init__(self, rad=1, region_size=10):
+    """Distributed Searchlight
+    
+    Run a user-defined function over each voxel in a multi-subject 
+    dataset.
+
+    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):
         """Constructor
 
         Parameters
         ----------
 
-        rad: number of voxels in border region
+        sl_rad: radius, in voxels, of the sphere inscribed in the 
+                   searchlight cube
 
-        region_size: max dimension of a 3D region
+        max_blk_edge: max edge length, in voxels, of the 3D block
 
         """
-        self.rad = rad
-        self.region_size = region_size
+        self.sl_rad = sl_rad
+        self.max_blk_edge = max_blk_edge
 
-    def _get_ownership(self, l):
+    def _get_ownership(self, data):
+        """Determine which processes own each subject
+
+        Parameters
+        ----------
+
+        data: list of 4D arrays with subject data
+
+        Returns
+        -------
+
+        list of ranks indicating the owner of each subject
+        """
         comm = MPI.COMM_WORLD
         rank = comm.rank
 
-        B = [(rank, idx) for (idx, c) in enumerate(l) if c is not None]
+        B = [(rank, idx) for (idx, c) in enumerate(data) if c is not None]
         C = comm.allreduce(B)
-        ownership = [None] * len(l)
+        ownership = [None] * len(data)
         for c in C:
             ownership[c[1]] = c[0]
         return ownership
 
-    def _get_points(self, mask):
-        outerblk = self.region_size + 2*self.rad
+    def _get_blocks(self, mask):
+        """Divide the volume into a set of blocks
+
+        Ignore blocks that have no active voxels in the mask
+
+        Parameters
+        ----------
+
+        mask: a boolean 3D array which is true at every active voxel
+
+        Returns
+        -------
+
+        list of tuples containing block information:
+           - a triple containing top left point of the block and
+           - a triple containing the size in voxels of the block
+
+        """
+        outerblk = self.max_blk_edge + 2*self.sl_rad
         return [((i, j, k),
                 mask[i:i+outerblk,
                 j:j+outerblk,
                 k:k+outerblk].shape)
-                for i in range(0, mask.shape[0], self.region_size)
-                for j in range(0, mask.shape[1], self.region_size)
-                for k in range(0, mask.shape[2], self.region_size)
-                if np.sum(
-                mask[i+self.rad:i+self.rad+self.region_size,
-                     j+self.rad:j+self.rad+self.region_size,
-                     k+self.rad:k+self.rad+self.region_size]) > 0]
-
-    def _subarray(self, mat3d, pts):
-        (pt, sz) = pts
-        if len(mat3d.shape) == 3:
-            return mat3d[pt[0]:pt[0]+sz[0],
+                for i in range(0, mask.shape[0], self.max_blk_edge)
+                for j in range(0, mask.shape[1], self.max_blk_edge)
+                for k in range(0, mask.shape[2], self.max_blk_edge)
+                if np.any( 
+                mask[i+self.sl_rad:i+self.sl_rad+self.max_blk_edge,
+                     j+self.sl_rad:j+self.sl_rad+self.max_blk_edge,
+                     k+self.sl_rad:k+self.sl_rad+self.max_blk_edge])]
+
+    def _get_block_data(self, mat, block):
+        """Retrieve a block from a 3D or 4D volume
+
+        Parameters
+        ----------
+
+        mat: a 3D or 4D volume
+
+        block: a tuple containing block information:
+          - a triple containing the top left point of the block and
+          - a triple containing the size in voxels of the block
+
+        Returns
+        -------
+
+        In the case of a 3D array, a 3D subarray at the block location
+
+        In the case of a 4D array, a 4D subarray at the block location,
+        including the entire 0th dimension.
+        """
+        (pt, sz) = block
+        if len(mat.shape) == 3:
+            return mat[pt[0]:pt[0]+sz[0],
                          pt[1]:pt[1]+sz[1],
                          pt[2]:pt[2]+sz[2]].copy()
-        elif len(mat3d.shape) == 4:
-            return mat3d[:,
+        elif len(mat.shape) == 4:
+            return mat[:,
                          pt[0]:pt[0]+sz[0],
                          pt[1]:pt[1]+sz[1],
                          pt[2]:pt[2]+sz[2]].copy()
@@ -79,22 +136,59 @@ def _subarray(self, mat3d, pts):
             print('Error num dims')
             sys.exit(0)
 
-    def _split_volume(self, mat3d, pts):
-        return [self._subarray(mat3d, pt) for pt in pts]
+    def _split_volume(self, mat, blocks):
+        """Convert a volume into a list of block data
+
+        Parameters
+        ----------
+
+        mat: A 3D or 4D array to be split
+
+        blocks: a list of tuples containing block information:
+          - a triple containing the top left point of the block and
+          - a triple containing the size in voxels of the block
+
+
+        Returns
+        -------
+
+        A list of the subarrays corresponding to each block
+          
+        """
+        return [self._get_block_data(mat, block) for block in blocks]
+
+    def _scatter_list(self, data, owner):
+        """Distribute a list from one rank to other ranks in a cyclic manner
+
+        Parameters
+        ----------
+
+        data: list of pickle-able data
+
+        owner: rank that owns the data
+
+        Returns
+        -------
+        
+        A list containing the data in a cyclic layout across ranks
+          
+        """
+
 
-    def _scatter_list(self, insubj, pts, owner):
         comm = MPI.COMM_WORLD
+        rank = comm.rank
         size = comm.size
         subject_submatrices = []
+        nblocks = comm.bcast(len(data) if rank == owner else None, root=owner)
 
         # For each submatrix
-        for idx in range(0, len(pts), size):
+        for idx in range(0, nblocks, size):
             padded = None
-            extra = max(0, idx+size - len(pts))
+            extra = max(0, idx+size - nblocks)
 
             # Pad with "None" so scatter can go to all processes
-            if insubj is not None:
-                padded = insubj[idx:idx+size]
+            if data is not None:
+                padded = data[idx:idx+size]
                 if extra > 0:
                     padded = padded + [None]*extra
 
@@ -113,7 +207,9 @@ def distribute(self, subj, mask):
         Parameters
         ----------
 
-        subj: list of 4D arrays containing subject data
+        subj: list of 4D arrays containing data for one or more subjects.
+              Each entry of the list must contaion a 4D array on only one
+              MPI rank. 
 
         mask: 3D array with nonzero values at active vertices
 
@@ -124,7 +220,7 @@ def distribute(self, subj, mask):
 
         # Get/set ownership
         ownership = self._get_ownership(subj)
-        full_pts = self._get_points(mask) if rank == 0 else None
+        full_pts = self._get_blocks(mask) if rank == 0 else None
         full_pts = comm.bcast(full_pts)
 
         # Divide data and mask
@@ -134,18 +230,32 @@ def distribute(self, subj, mask):
         submasks = self._split_volume(mask, full_pts)
 
         # Scatter points, data, and mask
-        self.pts = self._scatter_list(full_pts, full_pts, 0)
-        self.submasks = self._scatter_list(submasks, full_pts, 0)
-        self.subproblems = [self._scatter_list(s, full_pts, ownership[s_idx])
+        self.pts = self._scatter_list(full_pts, 0)
+        self.submasks = self._scatter_list(submasks,  0)
+        self.subproblems = [self._scatter_list(s, ownership[s_idx])
                             for (s_idx, s) in enumerate(splitsubj)]
 
-    def searchlight_region(self, region_fn, bcast_var):
-        """Perform a function for each region in a volume.
+    def broadcast(self, bcast_var):
+        """Distribute data to processes
+
+        Parameters
+        ----------
+
+        bcast_var:    shared data which is broadcast to all processes
+
+        """
+
+        comm = MPI.COMM_WORLD
+        rank = comm.rank
+        self.bcast_var = comm.bcast(bcast_var)
+
+    def searchlight_block(self, block_fn):
+        """Perform a function for each block in a volume.
 
         Parameters
         ----------
 
-        region_fn: function to apply to each region:
+        block_fn: function to apply to each block:
 
                 Parameters
 
@@ -153,7 +263,8 @@ def searchlight_region(self, region_fn, bcast_var):
 
                 mask: 3D array containing subset of mask data
 
-                rad: number of voxels in border region
+                sl_rad: radius, in voxels, of the sphere inscribed in the 
+                cube
 
                 bcast_var: shared data which is broadcast to all processes
 
@@ -163,17 +274,15 @@ def searchlight_region(self, region_fn, bcast_var):
                 3D array which is the same size as the mask
                 input with padding removed
 
-        bcast_var:    shared data which is broadcast to all processes
 
         """
         comm = MPI.COMM_WORLD
         rank = comm.rank
-        bcast_var = comm.bcast(bcast_var)
 
         # Apply searchlight
         local_outputs = [(mypt[0],
-                          region_fn([d[idx] for d in self.subproblems],
-                          self.submasks[idx], self.rad, bcast_var))
+                          block_fn([d[idx] for d in self.subproblems],
+                          self.submasks[idx], self.sl_rad, self.bcast_var))
                          for (idx, mypt) in enumerate(self.pts)]
 
         # Collect results
@@ -187,13 +296,13 @@ def searchlight_region(self, region_fn, bcast_var):
                     for i in range(0, mat.shape[0]):
                         for j in range(0, mat.shape[1]):
                             for k in range(0, mat.shape[2]):
-                                outmat[pt[0]+self.rad+i,
-                                       pt[1]+self.rad+j,
-                                       pt[2]+self.rad+k] = mat[i, j, k]
+                                outmat[pt[0]+self.sl_rad+i,
+                                       pt[1]+self.sl_rad+j,
+                                       pt[2]+self.sl_rad+k] = mat[i, j, k]
 
         return outmat
 
-    def searchlight_voxel(self, voxel_fn, bcast_var):
+    def searchlight_voxel(self, voxel_fn, pool_size=None):
         """Perform a function at each active voxel
 
         Parameters
@@ -207,7 +316,8 @@ def searchlight_voxel(self, voxel_fn, bcast_var):
 
                 mask: 3D array containing subset of mask data
 
-                rad: number of voxels in border region
+                sl_rad: radius, in voxels, of the sphere inscribed in the 
+                cube
 
                 bcast_var: shared data which is broadcast to all processes
 
@@ -216,32 +326,33 @@ def searchlight_voxel(self, voxel_fn, bcast_var):
 
                 Value of any pickle-able type
 
-        bcast_var:    shared data which is broadcast to all processes
+        pool_size:    Number of parallel processes in shared memory 
+                      process pool
 
         """
 
-        def _singlenode_searchlight(l, msk, myrad, bcast_var):
+        def _singlenode_searchlight(l, msk, mysl_rad, bcast_var):
 
-            outmat = np.empty(msk.shape, dtype=np.object)[myrad:-myrad,
-                                                          myrad:-myrad,
-                                                          myrad:-myrad]
+            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*myrad+1,
-                           j:j+2*myrad+1,
-                           k:k+2*myrad+1]
+                           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*myrad+1,
-                           j:j+2*myrad+1,
-                           k:k+2*myrad+1],
-                       myrad,
+                       msk[i:i+2*mysl_rad+1,
+                           j:j+2*mysl_rad+1,
+                           k:k+2*mysl_rad+1],
+                       mysl_rad,
                        bcast_var)
-                      if msk[i+myrad, j+myrad, k+myrad] else None
+                      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()
+            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))
@@ -256,4 +367,4 @@ def _singlenode_searchlight(l, msk, myrad, bcast_var):
 
             return outmat
 
-        return self.searchlight_region(_singlenode_searchlight, bcast_var)
+        return self.searchlight_block(_singlenode_searchlight)
diff --git a/examples/searchlight/example_searchlight.py b/examples/searchlight/example_searchlight.py
index 2bf598e77..40a0a077b 100644
--- a/examples/searchlight/example_searchlight.py
+++ b/examples/searchlight/example_searchlight.py
@@ -66,10 +66,11 @@
       data[idx][pt[0]:pt[0]+kernel_dim,pt[1]:pt[1]+kernel_dim,pt[2]:pt[2]+kernel_dim] -= kernel * weight
 
 # Create searchlight object
-sl = Searchlight(rad=1, region_size=5)
+sl = Searchlight(sl_rad=1, max_blk_edge=5)
 
 # Distribute data to processes
 sl.distribute([data], mask)
+sl.broadcast(labels)
 
 # Define voxel function
 def sfn(l, msk, myrad, bcast_var):
@@ -80,7 +81,7 @@ def sfn(l, msk, myrad, bcast_var):
   return np.mean(sklearn.model_selection.cross_val_score(classifier, data, bcast_var,n_jobs=1))
 
 # Run searchlight
-global_outputs = sl.searchlight_voxel(sfn, labels)
+global_outputs = sl.searchlight_voxel(sfn)
 
 # Visualize result
 if rank == 0:
diff --git a/tests/searchlight/test_searchlight.py b/tests/searchlight/test_searchlight.py
index 5a7299fc1..3ea7ae67a 100644
--- a/tests/searchlight/test_searchlight.py
+++ b/tests/searchlight/test_searchlight.py
@@ -23,11 +23,11 @@
 """
 
 def test_instantiate():
-  sl = Searchlight(rad=5, region_size=10)
+  sl = Searchlight(sl_rad=5, max_blk_edge=10)
   assert sl
 
 def test_correctness():
-  def voxel_test(data, mask, region_size, rad):
+  def voxel_test(data, mask, max_blk_edge, rad):
     
     comm = MPI.COMM_WORLD
     rank = comm.rank
@@ -69,9 +69,10 @@ def sfn(l,msk,myrad,bcast_var):
       return midpt
     
   
-    sl = Searchlight(rad=rad, region_size=region_size)
+    sl = Searchlight(sl_rad=rad, max_blk_edge=max_blk_edge)
     sl.distribute(data, mask)
-    global_outputs = sl.searchlight_voxel(sfn, mask)
+    sl.broadcast(mask)
+    global_outputs = sl.searchlight_voxel(sfn)
   
     if rank == 0:
       for d0 in range(rad, global_outputs.shape[0]-rad):
@@ -81,7 +82,7 @@ def sfn(l,msk,myrad,bcast_var):
               assert np.array_equal(np.array(global_outputs[d0,d1,d2]), np.array([d0,d1,d2]))
   
   
-  def region_test(data, mask, region_size, rad):
+  def block_test(data, mask, max_blk_edge, rad):
     
     comm = MPI.COMM_WORLD
     rank = comm.rank
@@ -104,9 +105,10 @@ def sfn(l,msk,myrad,bcast_var):
       outmat[~msk] = None
       return outmat[rad:-rad,rad:-rad,rad:-rad] 
     
-    sl = Searchlight(rad=rad, region_size=region_size)
+    sl = Searchlight(sl_rad=rad, max_blk_edge=max_blk_edge)
     sl.distribute(data, mask)
-    global_outputs = sl.searchlight_region(sfn, mask)
+    sl.broadcast(mask)
+    global_outputs = sl.searchlight_block(sfn)
   
     if rank == 0:
       for d0 in range(rad, global_outputs.shape[0]-rad):
@@ -116,20 +118,21 @@ def sfn(l,msk,myrad,bcast_var):
               assert np.array_equal(np.array(global_outputs[d0,d1,d2]), np.array([0,d0,d1,d2]))
   
   # Create dataset
-  def do_test(dim0, dim1, dim2, ntr, nsubj, region_size, rad):
+  def do_test(dim0, dim1, dim2, ntr, nsubj, max_blk_edge, rad):
     comm = MPI.COMM_WORLD
     rank = comm.rank
     size = comm.size
     mask = np.random.choice([True, False], (dim0,dim1,dim2))
     data = [np.empty((ntr,dim0,dim1,dim2), dtype=np.object) if i % size == rank else None for i in range(0, nsubj)]
-    voxel_test(data, mask, region_size, rad)
-    region_test(data, mask, region_size, rad)
+    voxel_test(data, mask, max_blk_edge, rad)
+    block_test(data, mask, max_blk_edge, rad)
   
   
-  do_test(dim0=7, dim1=5, dim2=9, ntr=5, nsubj=1, region_size=4, rad=1)
-  do_test(dim0=7, dim1=5, dim2=9, ntr=5, nsubj=5, region_size=4, rad=1)
-  do_test(dim0=1, dim1=5, dim2=9, ntr=5, nsubj=5, region_size=4, rad=1)
-  do_test(dim0=0, dim1=10, dim2=8, ntr=5, nsubj=5, region_size=4, rad=1)
-  do_test(dim0=7, dim1=5, dim2=9, ntr=5, nsubj=1, region_size=4, rad=2)
-  do_test(dim0=7, dim1=5, dim2=9, ntr=5, nsubj=1, region_size=4, rad=3)
+  do_test(dim0=7, dim1=5, dim2=9, ntr=5, nsubj=1, max_blk_edge=4, rad=1)
+  do_test(dim0=7, dim1=5, dim2=9, ntr=5, nsubj=5, max_blk_edge=4, rad=1)
+  do_test(dim0=1, dim1=5, dim2=9, ntr=5, nsubj=5, max_blk_edge=4, rad=1)
+  do_test(dim0=0, dim1=10, dim2=8, ntr=5, nsubj=5, max_blk_edge=4, rad=1)
+  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()

From 515afa146d88b6d0c59fd2ddd8397cb7e05ba974 Mon Sep 17 00:00:00 2001
From: Michael Anderson <michael.j.anderson@intel.com>
Date: Tue, 22 Nov 2016 07:51:50 -0800
Subject: [PATCH 06/11] fixed import error

---
 examples/searchlight/example_searchlight.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/examples/searchlight/example_searchlight.py b/examples/searchlight/example_searchlight.py
index 40a0a077b..e8a6b708f 100644
--- a/examples/searchlight/example_searchlight.py
+++ b/examples/searchlight/example_searchlight.py
@@ -16,7 +16,6 @@
 from mpi4py import MPI
 import sys
 
-from brainiak.searchlight import example_fn
 from brainiak.searchlight.searchlight import Searchlight
 
 """Distributed Searchlight Example

From 08d7a95a1848c52010b95e597a51cb9840085f18 Mon Sep 17 00:00:00 2001
From: Michael Anderson <michael.j.anderson@intel.com>
Date: Tue, 22 Nov 2016 08:02:36 -0800
Subject: [PATCH 07/11] addressing pr commentS

---
 brainiak/searchlight/searchlight.py         | 2 +-
 examples/searchlight/example_searchlight.py | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/brainiak/searchlight/searchlight.py b/brainiak/searchlight/searchlight.py
index 6a81336b9..402c8963a 100644
--- a/brainiak/searchlight/searchlight.py
+++ b/brainiak/searchlight/searchlight.py
@@ -209,7 +209,7 @@ def distribute(self, subj, mask):
 
         subj: list of 4D arrays containing data for one or more subjects.
               Each entry of the list must contaion a 4D array on only one
-              MPI rank. 
+              MPI rank, and "None" on all other ranks.
 
         mask: 3D array with nonzero values at active vertices
 
diff --git a/examples/searchlight/example_searchlight.py b/examples/searchlight/example_searchlight.py
index e8a6b708f..2f8e32ee5 100644
--- a/examples/searchlight/example_searchlight.py
+++ b/examples/searchlight/example_searchlight.py
@@ -19,6 +19,7 @@
 from brainiak.searchlight.searchlight import Searchlight
 
 """Distributed Searchlight Example
+example usage: mpirun -n 4 python3 example_searchlight.py
 """
 
 comm = MPI.COMM_WORLD

From 7ef8437617a67178c8eab4c41c532b10a3db94c3 Mon Sep 17 00:00:00 2001
From: Michael Anderson <michael.j.anderson@intel.com>
Date: Tue, 22 Nov 2016 08:08:03 -0800
Subject: [PATCH 08/11] fixed formatting

---
 brainiak/searchlight/searchlight.py   | 33 ++++++++++++---------------
 tests/searchlight/test_searchlight.py |  1 -
 2 files changed, 15 insertions(+), 19 deletions(-)

diff --git a/brainiak/searchlight/searchlight.py b/brainiak/searchlight/searchlight.py
index 402c8963a..58e888703 100644
--- a/brainiak/searchlight/searchlight.py
+++ b/brainiak/searchlight/searchlight.py
@@ -26,8 +26,8 @@
 
 class Searchlight:
     """Distributed Searchlight
-    
-    Run a user-defined function over each voxel in a multi-subject 
+
+    Run a user-defined function over each voxel in a multi-subject
     dataset.
 
     Optionally, users can define a block function which runs over
@@ -39,7 +39,7 @@ def __init__(self, sl_rad=1, max_blk_edge=10):
         Parameters
         ----------
 
-        sl_rad: radius, in voxels, of the sphere inscribed in the 
+        sl_rad: radius, in voxels, of the sphere inscribed in the
                    searchlight cube
 
         max_blk_edge: max edge length, in voxels, of the 3D block
@@ -97,7 +97,7 @@ def _get_blocks(self, mask):
                 for i in range(0, mask.shape[0], self.max_blk_edge)
                 for j in range(0, mask.shape[1], self.max_blk_edge)
                 for k in range(0, mask.shape[2], self.max_blk_edge)
-                if np.any( 
+                if np.any(
                 mask[i+self.sl_rad:i+self.sl_rad+self.max_blk_edge,
                      j+self.sl_rad:j+self.sl_rad+self.max_blk_edge,
                      k+self.sl_rad:k+self.sl_rad+self.max_blk_edge])]
@@ -125,15 +125,14 @@ def _get_block_data(self, mat, block):
         (pt, sz) = block
         if len(mat.shape) == 3:
             return mat[pt[0]:pt[0]+sz[0],
-                         pt[1]:pt[1]+sz[1],
-                         pt[2]:pt[2]+sz[2]].copy()
+                       pt[1]:pt[1]+sz[1],
+                       pt[2]:pt[2]+sz[2]].copy()
         elif len(mat.shape) == 4:
             return mat[:,
-                         pt[0]:pt[0]+sz[0],
-                         pt[1]:pt[1]+sz[1],
-                         pt[2]:pt[2]+sz[2]].copy()
+                       pt[0]:pt[0]+sz[0],
+                       pt[1]:pt[1]+sz[1],
+                       pt[2]:pt[2]+sz[2]].copy()
         else:
-            print('Error num dims')
             sys.exit(0)
 
     def _split_volume(self, mat, blocks):
@@ -153,7 +152,7 @@ def _split_volume(self, mat, blocks):
         -------
 
         A list of the subarrays corresponding to each block
-          
+
         """
         return [self._get_block_data(mat, block) for block in blocks]
 
@@ -169,11 +168,10 @@ def _scatter_list(self, data, owner):
 
         Returns
         -------
-        
+
         A list containing the data in a cyclic layout across ranks
-          
-        """
 
+        """
 
         comm = MPI.COMM_WORLD
         rank = comm.rank
@@ -246,7 +244,6 @@ def broadcast(self, bcast_var):
         """
 
         comm = MPI.COMM_WORLD
-        rank = comm.rank
         self.bcast_var = comm.bcast(bcast_var)
 
     def searchlight_block(self, block_fn):
@@ -263,7 +260,7 @@ def searchlight_block(self, block_fn):
 
                 mask: 3D array containing subset of mask data
 
-                sl_rad: radius, in voxels, of the sphere inscribed in the 
+                sl_rad: radius, in voxels, of the sphere inscribed in the
                 cube
 
                 bcast_var: shared data which is broadcast to all processes
@@ -316,7 +313,7 @@ def searchlight_voxel(self, voxel_fn, pool_size=None):
 
                 mask: 3D array containing subset of mask data
 
-                sl_rad: radius, in voxels, of the sphere inscribed in the 
+                sl_rad: radius, in voxels, of the sphere inscribed in the
                 cube
 
                 bcast_var: shared data which is broadcast to all processes
@@ -326,7 +323,7 @@ def searchlight_voxel(self, voxel_fn, pool_size=None):
 
                 Value of any pickle-able type
 
-        pool_size:    Number of parallel processes in shared memory 
+        pool_size:    Number of parallel processes in shared memory
                       process pool
 
         """
diff --git a/tests/searchlight/test_searchlight.py b/tests/searchlight/test_searchlight.py
index 3ea7ae67a..785658504 100644
--- a/tests/searchlight/test_searchlight.py
+++ b/tests/searchlight/test_searchlight.py
@@ -16,7 +16,6 @@
 from mpi4py import MPI
 import sys
 
-from brainiak.searchlight import example_fn
 from brainiak.searchlight.searchlight import Searchlight
 
 """Distributed Searchlight Test

From b571b5857bce909e3ed907bb48881f8ae90fb168 Mon Sep 17 00:00:00 2001
From: Michael Anderson <michael.j.anderson@intel.com>
Date: Tue, 22 Nov 2016 08:39:34 -0800
Subject: [PATCH 09/11] removed last line

---
 tests/searchlight/test_searchlight.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tests/searchlight/test_searchlight.py b/tests/searchlight/test_searchlight.py
index 785658504..712b1c97c 100644
--- a/tests/searchlight/test_searchlight.py
+++ b/tests/searchlight/test_searchlight.py
@@ -134,4 +134,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()

From d5cc3b845b68cd708a9e4c360733a11dc1f52a7e Mon Sep 17 00:00:00 2001
From: Michael Anderson <michael.j.anderson@intel.com>
Date: Tue, 22 Nov 2016 09:41:56 -0800
Subject: [PATCH 10/11] pr fix 2

---
 brainiak/searchlight/searchlight.py         | 71 ++++++++++++---------
 examples/searchlight/example_searchlight.py |  2 +-
 tests/searchlight/test_searchlight.py       |  4 +-
 3 files changed, 44 insertions(+), 33 deletions(-)

diff --git a/brainiak/searchlight/searchlight.py b/brainiak/searchlight/searchlight.py
index 58e888703..24e8e1f6c 100644
--- a/brainiak/searchlight/searchlight.py
+++ b/brainiak/searchlight/searchlight.py
@@ -40,16 +40,17 @@ def __init__(self, sl_rad=1, max_blk_edge=10):
         ----------
 
         sl_rad: radius, in voxels, of the sphere inscribed in the
-                   searchlight cube
+                   searchlight cube, not counting the center voxel
 
         max_blk_edge: max edge length, in voxels, of the 3D block
 
         """
         self.sl_rad = sl_rad
         self.max_blk_edge = max_blk_edge
+        self.comm = MPI.COMM_WORLD
 
     def _get_ownership(self, data):
-        """Determine which processes own each subject
+        """Determine on which rank each subject currently resides
 
         Parameters
         ----------
@@ -61,11 +62,10 @@ def _get_ownership(self, data):
 
         list of ranks indicating the owner of each subject
         """
-        comm = MPI.COMM_WORLD
-        rank = comm.rank
+        rank = self.comm.rank
 
         B = [(rank, idx) for (idx, c) in enumerate(data) if c is not None]
-        C = comm.allreduce(B)
+        C = self.comm.allreduce(B)
         ownership = [None] * len(data)
         for c in C:
             ownership[c[1]] = c[0]
@@ -120,7 +120,7 @@ def _get_block_data(self, mat, block):
         In the case of a 3D array, a 3D subarray at the block location
 
         In the case of a 4D array, a 4D subarray at the block location,
-        including the entire 0th dimension.
+        including the entire fourth dimension.
         """
         (pt, sz) = block
         if len(mat.shape) == 3:
@@ -173,11 +173,11 @@ def _scatter_list(self, data, owner):
 
         """
 
-        comm = MPI.COMM_WORLD
-        rank = comm.rank
-        size = comm.size
+        rank = self.comm.rank
+        size = self.comm.size
         subject_submatrices = []
-        nblocks = comm.bcast(len(data) if rank == owner else None, root=owner)
+        nblocks = self.comm.bcast(len(data)
+                                  if rank == owner else None, root=owner)
 
         # For each submatrix
         for idx in range(0, nblocks, size):
@@ -191,7 +191,7 @@ def _scatter_list(self, data, owner):
                     padded = padded + [None]*extra
 
             # Scatter submatrices to all processes
-            mytrans = comm.scatter(padded, root=owner)
+            mytrans = self.comm.scatter(padded, root=owner)
 
             # Contribute submatrix to subject list
             if mytrans is not None:
@@ -206,29 +206,41 @@ def distribute(self, subj, mask):
         ----------
 
         subj: list of 4D arrays containing data for one or more subjects.
-              Each entry of the list must contaion a 4D array on only one
-              MPI rank, and "None" on all other ranks.
+              Each entry of the list must be present on at most one rank,
+              and the other ranks contain a "None" at this list location.
 
-        mask: 3D array with nonzero values at active vertices
+              For example, for 3 ranks you may lay out the data in the
+              following manner:
+
+              Rank 0: [Subj0, None, None]
+              Rank 1: [None, Subj1, None]
+              Rank 2: [None, None, Subj2]
+
+              Or alternatively, you may lay out the data in this manner:
+
+              Rank 0: [Subj0, Subj1, Subj2]
+              Rank 1: [None, None, None]
+              Rank 2: [None, None, None]
+
+        mask: 3D array with "True" entries at active vertices
 
         """
         self.mask = mask
-        comm = MPI.COMM_WORLD
-        rank = comm.rank
+        rank = self.comm.rank
 
         # Get/set ownership
         ownership = self._get_ownership(subj)
-        full_pts = self._get_blocks(mask) if rank == 0 else None
-        full_pts = comm.bcast(full_pts)
+        all_blocks = self._get_blocks(mask) if rank == 0 else None
+        all_blocks = self.comm.bcast(all_blocks)
 
         # Divide data and mask
-        splitsubj = [self._split_volume(s, full_pts)
+        splitsubj = [self._split_volume(s, all_blocks)
                      if s is not None else None
                      for s in subj]
-        submasks = self._split_volume(mask, full_pts)
+        submasks = self._split_volume(mask, all_blocks)
 
         # Scatter points, data, and mask
-        self.pts = self._scatter_list(full_pts, 0)
+        self.pts = self._scatter_list(all_blocks, 0)
         self.submasks = self._scatter_list(submasks,  0)
         self.subproblems = [self._scatter_list(s, ownership[s_idx])
                             for (s_idx, s) in enumerate(splitsubj)]
@@ -243,10 +255,9 @@ def broadcast(self, bcast_var):
 
         """
 
-        comm = MPI.COMM_WORLD
-        self.bcast_var = comm.bcast(bcast_var)
+        self.bcast_var = self.comm.bcast(bcast_var)
 
-    def searchlight_block(self, block_fn):
+    def run_block_function(self, block_fn):
         """Perform a function for each block in a volume.
 
         Parameters
@@ -256,7 +267,8 @@ def searchlight_block(self, block_fn):
 
                 Parameters
 
-                subj: list of 4D arrays containing subset of subject data
+                data: list of 4D arrays containing subset of subject data,
+                      which is padded with sl_rad voxels.
 
                 mask: 3D array containing subset of mask data
 
@@ -273,8 +285,7 @@ def searchlight_block(self, block_fn):
 
 
         """
-        comm = MPI.COMM_WORLD
-        rank = comm.rank
+        rank = self.comm.rank
 
         # Apply searchlight
         local_outputs = [(mypt[0],
@@ -283,7 +294,7 @@ def searchlight_block(self, block_fn):
                          for (idx, mypt) in enumerate(self.pts)]
 
         # Collect results
-        global_outputs = comm.gather(local_outputs)
+        global_outputs = self.comm.gather(local_outputs)
 
         # Coalesce results
         outmat = np.empty(self.mask.shape, dtype=np.object)
@@ -299,7 +310,7 @@ def searchlight_block(self, block_fn):
 
         return outmat
 
-    def searchlight_voxel(self, voxel_fn, pool_size=None):
+    def run_searchlight(self, voxel_fn, pool_size=None):
         """Perform a function at each active voxel
 
         Parameters
@@ -364,4 +375,4 @@ def _singlenode_searchlight(l, msk, mysl_rad, bcast_var):
 
             return outmat
 
-        return self.searchlight_block(_singlenode_searchlight)
+        return self.run_block_function(_singlenode_searchlight)
diff --git a/examples/searchlight/example_searchlight.py b/examples/searchlight/example_searchlight.py
index 2f8e32ee5..28901e354 100644
--- a/examples/searchlight/example_searchlight.py
+++ b/examples/searchlight/example_searchlight.py
@@ -81,7 +81,7 @@ def sfn(l, msk, myrad, bcast_var):
   return np.mean(sklearn.model_selection.cross_val_score(classifier, data, bcast_var,n_jobs=1))
 
 # Run searchlight
-global_outputs = sl.searchlight_voxel(sfn)
+global_outputs = sl.run_searchlight(sfn)
 
 # Visualize result
 if rank == 0:
diff --git a/tests/searchlight/test_searchlight.py b/tests/searchlight/test_searchlight.py
index 712b1c97c..8dca187d1 100644
--- a/tests/searchlight/test_searchlight.py
+++ b/tests/searchlight/test_searchlight.py
@@ -71,7 +71,7 @@ def sfn(l,msk,myrad,bcast_var):
     sl = Searchlight(sl_rad=rad, max_blk_edge=max_blk_edge)
     sl.distribute(data, mask)
     sl.broadcast(mask)
-    global_outputs = sl.searchlight_voxel(sfn)
+    global_outputs = sl.run_searchlight(sfn)
   
     if rank == 0:
       for d0 in range(rad, global_outputs.shape[0]-rad):
@@ -107,7 +107,7 @@ def sfn(l,msk,myrad,bcast_var):
     sl = Searchlight(sl_rad=rad, max_blk_edge=max_blk_edge)
     sl.distribute(data, mask)
     sl.broadcast(mask)
-    global_outputs = sl.searchlight_block(sfn)
+    global_outputs = sl.run_block_function(sfn)
   
     if rank == 0:
       for d0 in range(rad, global_outputs.shape[0]-rad):

From e0f7e9d1c6848d152a7c5afc0c6822504921f9e7 Mon Sep 17 00:00:00 2001
From: Michael Anderson <michael.j.anderson@intel.com>
Date: Tue, 22 Nov 2016 11:00:06 -0800
Subject: [PATCH 11/11] changes for pr 3

---
 brainiak/searchlight/searchlight.py         | 28 ++++++++++-----------
 examples/searchlight/example_searchlight.py |  8 +++---
 tests/searchlight/test_searchlight.py       | 24 ++++++++++--------
 3 files changed, 31 insertions(+), 29 deletions(-)

diff --git a/brainiak/searchlight/searchlight.py b/brainiak/searchlight/searchlight.py
index 24e8e1f6c..e5f7fdf9d 100644
--- a/brainiak/searchlight/searchlight.py
+++ b/brainiak/searchlight/searchlight.py
@@ -111,7 +111,7 @@ def _get_block_data(self, mat, block):
         mat: a 3D or 4D volume
 
         block: a tuple containing block information:
-          - a triple containing the top left point of the block and
+          - a triple containing the lowest-coordinate voxel in the block
           - a triple containing the size in voxels of the block
 
         Returns
@@ -128,10 +128,10 @@ def _get_block_data(self, mat, block):
                        pt[1]:pt[1]+sz[1],
                        pt[2]:pt[2]+sz[2]].copy()
         elif len(mat.shape) == 4:
-            return mat[:,
-                       pt[0]:pt[0]+sz[0],
+            return mat[pt[0]:pt[0]+sz[0],
                        pt[1]:pt[1]+sz[1],
-                       pt[2]:pt[2]+sz[2]].copy()
+                       pt[2]:pt[2]+sz[2],
+                       :].copy()
         else:
             sys.exit(0)
 
@@ -199,13 +199,13 @@ def _scatter_list(self, data, owner):
 
         return subject_submatrices
 
-    def distribute(self, subj, mask):
-        """Distribute data to processes
+    def distribute(self, subjects, mask):
+        """Distribute data to MPI ranks
 
         Parameters
         ----------
 
-        subj: list of 4D arrays containing data for one or more subjects.
+        subjects : list of 4D arrays containing data for one or more subjects.
               Each entry of the list must be present on at most one rank,
               and the other ranks contain a "None" at this list location.
 
@@ -229,18 +229,18 @@ def distribute(self, subj, mask):
         rank = self.comm.rank
 
         # Get/set ownership
-        ownership = self._get_ownership(subj)
+        ownership = self._get_ownership(subjects)
         all_blocks = self._get_blocks(mask) if rank == 0 else None
         all_blocks = self.comm.bcast(all_blocks)
 
         # Divide data and mask
         splitsubj = [self._split_volume(s, all_blocks)
                      if s is not None else None
-                     for s in subj]
+                     for s in subjects]
         submasks = self._split_volume(mask, all_blocks)
 
         # Scatter points, data, and mask
-        self.pts = self._scatter_list(all_blocks, 0)
+        self.blocks = self._scatter_list(all_blocks, 0)
         self.submasks = self._scatter_list(submasks,  0)
         self.subproblems = [self._scatter_list(s, ownership[s_idx])
                             for (s_idx, s) in enumerate(splitsubj)]
@@ -291,7 +291,7 @@ def run_block_function(self, block_fn):
         local_outputs = [(mypt[0],
                           block_fn([d[idx] for d in self.subproblems],
                           self.submasks[idx], self.sl_rad, self.bcast_var))
-                         for (idx, mypt) in enumerate(self.pts)]
+                         for (idx, mypt) in enumerate(self.blocks)]
 
         # Collect results
         global_outputs = self.comm.gather(local_outputs)
@@ -346,10 +346,10 @@ def _singlenode_searchlight(l, msk, mysl_rad, bcast_var):
                                                           mysl_rad:-mysl_rad]
 
             import pathos.multiprocessing
-            inlist = [([ll[:,
-                           i:i+2*mysl_rad+1,
+            inlist = [([ll[i:i+2*mysl_rad+1,
                            j:j+2*mysl_rad+1,
-                           k:k+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,
diff --git a/examples/searchlight/example_searchlight.py b/examples/searchlight/example_searchlight.py
index 28901e354..2da316b4f 100644
--- a/examples/searchlight/example_searchlight.py
+++ b/examples/searchlight/example_searchlight.py
@@ -37,7 +37,7 @@
 weight = 1
 
 # Generate data
-data = np.random.random((ntr,dim,dim,dim)) if rank == 0 else None
+data = np.random.random((dim,dim,dim,ntr)) if rank == 0 else None
 mask = np.zeros((dim,dim,dim), dtype=np.bool)
 for i in range(dim):
   for j in range(dim):
@@ -61,9 +61,9 @@
 
   for (idx, l) in enumerate(labels):
     if l:
-      data[idx][pt[0]:pt[0]+kernel_dim,pt[1]:pt[1]+kernel_dim,pt[2]:pt[2]+kernel_dim] += kernel * weight
+      data[pt[0]:pt[0]+kernel_dim,pt[1]:pt[1]+kernel_dim,pt[2]:pt[2]+kernel_dim,idx] += kernel * weight
     else:
-      data[idx][pt[0]:pt[0]+kernel_dim,pt[1]:pt[1]+kernel_dim,pt[2]:pt[2]+kernel_dim] -= kernel * weight
+      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)
@@ -77,7 +77,7 @@ def sfn(l, msk, myrad, bcast_var):
   import sklearn.svm
   import sklearn.model_selection
   classifier = sklearn.svm.SVC()
-  data = l[0][:,msk]
+  data = l[0][msk,:].T
   return np.mean(sklearn.model_selection.cross_val_score(classifier, data, bcast_var,n_jobs=1))
 
 # Run searchlight
diff --git a/tests/searchlight/test_searchlight.py b/tests/searchlight/test_searchlight.py
index 8dca187d1..85a6e26ae 100644
--- a/tests/searchlight/test_searchlight.py
+++ b/tests/searchlight/test_searchlight.py
@@ -38,25 +38,26 @@ def voxel_test(data, mask, max_blk_edge, rad):
     # Initialize dataset with known pattern
     for subj in data:
       if subj is not None:
-        for tr in range(subj.shape[0]):
+        for tr in range(subj.shape[3]):
           for d1 in range(dim0):
             for d2 in range(dim1):
               for d3 in range(dim2): 
-                subj[tr,d1,d2,d3] = np.array([tr, d1, d2, d3])
+                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 subj:
+        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([0,d0-rad,d1-rad,d2-rad]))
+                assert np.array_equal(tr[d0,d1,d2] - midpt, np.array([d0-rad,d1-rad,d2-rad,0]))
     
       # Determine midpoint
-      midpt = l[0][0,rad,rad,rad]
-      midpt = (midpt[1], midpt[2], midpt[3])
+      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]):
@@ -93,14 +94,14 @@ def block_test(data, mask, max_blk_edge, rad):
     # Initialize dataset with known pattern
     for subj in data:
       if subj is not None:
-        for tr in range(subj.shape[0]):
+        for tr in range(subj.shape[3]):
           for d1 in range(dim0):
             for d2 in range(dim1):
               for d3 in range(dim2): 
-                subj[tr,d1,d2,d3] = np.array([tr, d1, d2, d3])
+                subj[d1,d2,d3,tr] = np.array([ d1, d2, d3, tr])
     
     def sfn(l,msk,myrad,bcast_var):
-      outmat = l[0][0,:,:,:] 
+      outmat = l[0][:,:,:,0] 
       outmat[~msk] = None
       return outmat[rad:-rad,rad:-rad,rad:-rad] 
     
@@ -114,7 +115,7 @@ def sfn(l,msk,myrad,bcast_var):
         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([0,d0,d1,d2]))
+              assert np.array_equal(np.array(global_outputs[d0,d1,d2]), np.array([d0,d1,d2,0]))
   
   # Create dataset
   def do_test(dim0, dim1, dim2, ntr, nsubj, max_blk_edge, rad):
@@ -122,7 +123,7 @@ def do_test(dim0, dim1, dim2, ntr, nsubj, max_blk_edge, rad):
     rank = comm.rank
     size = comm.size
     mask = np.random.choice([True, False], (dim0,dim1,dim2))
-    data = [np.empty((ntr,dim0,dim1,dim2), dtype=np.object) if i % size == rank else None for i in range(0, nsubj)]
+    data = [np.empty((dim0,dim1,dim2,ntr), dtype=np.object) if i % size == rank else None for i in range(0, nsubj)]
     voxel_test(data, mask, max_blk_edge, rad)
     block_test(data, mask, max_blk_edge, rad)
   
@@ -134,3 +135,4 @@ 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()