Searchlights

V.0.1 - Alpha testing, contributions

In univariate analyses there is a distinction between whole-brain and ROI-based analyses. When done correctly, whole-brain univariate analyses allow you to identify, in an unbiased and data-driven way, regions of the brain that differ between experimental conditions. Compare this to what we have done with multivariate classification: up to this point we have taken ROIs (e.g., FFA, PPA, even the whole brain) and looked at the pattern of activity across voxels. Using these procedures we have been able to determine whether these voxels collectively contain information about different conditions. However, it has been unclear whether a subset of these voxels have been driving classification accuracy, partly because it is tricky to interpret the resulting weights. In other words, we have been unable to say with certainty where in the brain there are representations that distinguish between conditions. For more information on the searchlight technique, reade this comprehensive NeuroImage Comment by Etzel et al. (2013), that also includes citations to landmark papers on the topic.

A searchlight is a spherical or cubic moving window that allows us to identify and localize such representations. Running a searchlight is computationally intensive because it involves running a separate multivariate analysis for every voxel in the brain. Imagine doing your feature selection, hyper-parameter optimization, and cross-validation 100,000 times or more. Fortunately, BrainIAK contains a procedure that efficiently carves up brain data into appropriate chunks and distributes them across the computational resources available.

The validity of using searchlights is explained in the following article: Kriegeskorte, N., Goebel, R., & Bandettini, P. (2006). Information-based functional brain mapping. Proceedings of the National Academy of Sciences of the United States of America, 103, 3863–3868. https://doi.org/10.1073/pnas.0600244103

For this script we will use the localizer dataset from Kim et al. (2017) again.

Goal of this script

  1. Learn how to perform a whole-brain searchlight
  2. Learn how to replace the kernel used inside the searchlight

Table of contents

  1. Searchlights
  2. Searchlight Workflow

    2.1 Data Preparation
    2.2 Executing the Searchlight Workflow

  3. Running Searchlights on High Performance Clusters

Exercises:

Exercise 1
Exercise 2
Exercise 3
Exercise 4
Exercise 5
Exercise 6
Exercise 7
Exercise 8
Exercise 9
Exercise 10

2. Searchlight Workflow

Running a searchlight is computationally intensive, complex and involves many steps. To show how the searchlight functionality in BrainIAK works, we will walk through the steps needed to perform a simple searchlight that can be run in this notebook before moving onto a real example. This simple workflow takes the following steps:

  1. Data Preparation
  2. Set the searchlight parameters
  3. Create the searchlight object
  4. Create our classification kernel
  5. Execute the searchlight

However, before we start, there are a few things that you should now about searchlights. Think of a searchlight as a processing step in which you pull out a certain size chunk of your data and then perform a kernel operation (specified by a function you write). You can use any kernel on this chunk of data interchangeably. Critically, the searchlight function does not specify the analysis you want to perform, all it does is carve up the data.

Computational demand and parallelization

As mentioned before, searchlights are computationally intensive and so we need to be aware of what kind of burden they might impose. In the participant you loaded, we want to perform the operation about 177,000 times (once for each of their brain voxels). If we were to run this serially, even if each operation only took 1 s to complete, the analysis would take 50 hours(!), and only for this one participant. With nested cross-validation or other recursive steps, the full analysis could take months or years (and you'd lose a lot of points on your assignment).

Fortunately, the searchlight function in BrainIAK intelligently parallelizes your data to give you a considerable and scalable speed boost. Parallelization is the idea that when two or more computational tasks can be completed independently (because they don't interact in any way) then it is possible to run these tasks simultaneously on different cores. Note that we refer to cores as our unit of serial processing, although there are other parallelizations available within-core, such as threads or even multiple instructions within thread. The nice thing about parallelization is that it is scalable: in general, a job parallelized across 10 cores will run 10 times faster than on one core. For reference, your computer likely has 2, 4, or more cores, so you could speed up processing if you recruited all of these resources (and shut down all other types of background processing). Milgram, the cluster you are running on, has about 1,500 cores, meaning in theory if you could run a job simultaneously across all of these cores then the searchlight we talked about before would be taken down from 50 hours to 2 mins. In practice, there is often some start up cost when calling a function, such as loading the data, which could reduce this benefit a little (maybe to 5 mins).

So remember, the two main things that determine the speed of a searchlight: the **kernel algorithm** and the **amount of parallelization**.

How does the BrainIAK searchlight work?

To analyze data efficiently, it is important to understand how the Searchlight function in BrainIAK works. The searchlight code simply searches for every voxel in the mask that is equal to 1 and then centers a searchlight around this. This means that for every value of 1 in your mask, the searchlight function will apply your kernel function. Hence when writing the kernel you need to keep in mind that the input data that the function receives is not the size of the original data but is instead the size of the searchlight radius. In other words, you might give the searchlight function a brain and mask volume that is 64x64x64 but each kernel operation only runs on a data and mask volume that is 3x3x3 (or whatever your searchlight radius is). As per the logic of searchlights, the center of each mask in the searchlight will be equal to 1 (otherwise the searchlight wouldn't have selected it).

How to start using a searchlight? Small!

When getting used to searchlights it is encouraged that you scale up your code gently. This is to prevent the possibility that you run a searchlight that takes a very long time to finish only to find there was a small error (which happens a lot). Instead it is better to write a simple function that runs quickly so you can check if your code works properly before scaling up. A simple workflow for testing your searchlight (and estimating how long your searchlight would take on a bigger dataset) would be this:

  1. Create a mask of one voxel and run the searchlight interactively (like we are doing now, using a notebook) to check whether the code works.
  2. Use Timestamps to extract the execution time
  3. Print the number of voxels that are passed to the searchlight function
  4. Run the searchlight as a job on the smallest unit of real data you have (a single run or single participant)
  5. Check the runtime and memory usage of this searchlight (e.g. on slurm: sacct -j $JID --format=jobid,maxvmsize,elapsed).

Taking our own advice, we are going to write a searchlight script below to perform a very, very simple searchlight. In fact we are only going to run this searchlight on one voxel in the brain so that we can see whether our code is working and how long each searchlight would take. After we have familiarized ourselves with the way searchlights work on a small scale, we will then graduate to full scale analyses using batch scripts and cluster computing.

2.1 Data Preparation

Prepare a single participant using a similar pipeline to what we have developed previously. The critical change needed here is the shape of the data: In the past we wanted a trial x voxel matrix, but the input to a searchlight is a 3D volume, hence we need to add a step, as below.

Exercise 1: Why does the input to a searchlight analysis need to be 3D rather than 2D?
A:

In [ ]:
# Import libraries
import nibabel as nib
import numpy as np
import matplotlib.pyplot as plt
%matplotlib notebook

from utils import load_data, load_labels, label2TR, shift_timing, reshape_data, blockwise_sampling

import time
from nilearn import plotting
from brainiak.searchlight.searchlight import Searchlight
from scipy.sparse import random

# Import machine learning libraries
from sklearn.model_selection import StratifiedShuffleSplit, GridSearchCV, cross_val_score
from sklearn.svm import SVC
In [ ]:
# Preset variables
# Make sure you edit the following line to reflect the directory where you are accessing the VDC dataset
dir = '/gpfs/milgram/data/cmhn-s18/datasets/vdc/'
# dir = '/jukebox/pniintel/brainiak_edu/datasets/vdc/' #Princeton

num_runs=3
TR=1.5
hrf_lag = 4.5  # In seconds what is the lag between a stimulus onset and the peak bold response
shift_size = int(hrf_lag / TR) # Convert the shift into TRs

sub_id = 1

# Convert the number into a participant folder name
if (sub_id < 10):
    sids = '0' + str(sub_id)
else:
    sids = str(sub_id)   

# Specify the subject name
sub = 'sub-' + sids

# Load subject labels
stim_label_allruns = load_labels(dir, sub)

# Load the fMRI data
epi_mask_data_all, whole_brain_mask = load_data(directory=dir, subject_name=sub, mask_name='', zscore_data=True)

# How many events are there on this run
_, events = stim_label_allruns.shape
events_run = int(events / num_runs)

# This can differ per participant
print(sub, '= TRs: ', epi_mask_data_all.shape[1], '; Voxels: ', epi_mask_data_all.shape[0])
TRs_run = int(epi_mask_data_all.shape[1] / num_runs)

# Convert the timing into TR indexes
stim_label_TR = label2TR(stim_label_allruns, num_runs, TR, TRs_run)

# Shift the data some amount
stim_label_TR_shifted = shift_timing(stim_label_TR, shift_size)

# Perform the reshaping of the data
bold_data, labels = reshape_data(stim_label_TR_shifted, epi_mask_data_all)

# Down sample the data to be blockwise rather than trialwise
bold_data, labels = blockwise_sampling(bold_data, labels)
In [ ]:
# Reshape the data so that they are a 3D volume stacked in time (4D array). Basically, this reverses a step that we did when using the "load_data()" function from utils.py

bold_vol = np.zeros((whole_brain_mask.shape[0], whole_brain_mask.shape[1], whole_brain_mask.shape[2], bold_data.shape[0]))  # Preset the shape of the output
coords = np.where(whole_brain_mask)  # Where are the non zero values
bold_vol[coords[0], coords[1], coords[2], :] = bold_data.T  # Insert the bold data

Easier Data Processing: Save the formatted Data

Now we save the data you created. We want to do this so that it will be easier to read the data into the searchlight function we will use later. However, be aware that saving many copies of your data is wasteful and unnecessary. In fact, it is the nilearn philosophy to save as little data as possible, ideally only saving your raw data and analysis code -- which are sufficient to replicate the analysis -- along with any outputs needed for visualizing the results. To save data with nibabel we need information about the voxel size and orientation of the brain in the data matrix.

In [ ]:
# Get the nifti parameters (would have been better to get these when we loaded in the data)
nii = nib.load(dir + sub + "/preprocessed/loc/%s_filtered2_d1_firstExampleFunc_r%d.nii" % (sub, 1))
affine_mat = nii.affine  # What is the orientation of the data
dimsize = nii.header.get_zooms()  # How big is each dimension of the data (first three are voxel size in mm, last is TR duration in s)

# Save the volume
output_name = ('%s_input.nii.gz' % (sub))
bold_nii = nib.Nifti1Image(bold_vol, affine_mat)
hdr = bold_nii.header  # get a handle for the .nii file's header
hdr.set_zooms((dimsize[0], dimsize[1], dimsize[2], dimsize[3]))
nib.save(bold_nii, output_name)

# Also save the mask data
output_name = ('%s_mask.nii.gz' % (sub))
mask_nii = nib.Nifti1Image(whole_brain_mask, affine_mat)
hdr = mask_nii.header  # get a handle for the .nii file's header
hdr.set_zooms((dimsize[0], dimsize[1], dimsize[2]))
nib.save(mask_nii, output_name)

# Save the label data
output_name = ('%s_labels.npy' % (sub))
np.save(output_name, labels)

2.3 Executing the Searchlight Workflow

2.3.1 Set Searchlight Parameters

To run the searchlight function in BrainIAK you need the following parameters:

  1. data = The brain data as a 4D volume.
  2. mask = A binary mask specifying the "center" voxels in the brain around which you want to perform searchlight analyses. A searchlight will be drawn around every voxel with the value of 1. Hence, if you chose to use the wholebrain mask as the mask for the searchlight procedure, the searchlight may include voxels outside of your mask when the "center" voxel is at the border of the mask. It is up to you to decide whether then to include these results.
  3. bcvar = An additional variable which can be a list, numpy array, dictionary, etc. you want to use in your searchlight kernel. For instance you might want the condition labels so that you can determine to which condition each 3D volume corresponds. If you don't need to broadcast anything, e.g, when doing RSA, set this to 'None'.
  4. sl_rad = The size of the searchlight's radius, excluding the center voxel. This means the total volume size of the searchlight, if using a cube, is defined as: ((2 * sl_rad) + 1) ^ 3.
  5. max_blk_edge = When the searchlight function carves the data up into chunks, it doesn't distribute only a single searchlight's worth of data. Instead, it creates a block of data, with the edge length specified by this variable, which determines the number of searchlights to run within a job.
  6. pool_size = Maximum number of cores running on a block (typically 1).

Exercise 2: Searchlights don't need to be cubes. What other shape(s) does BrainIAK support and how do you specify this?

In [ ]:
# Make a mask of only one arbitrary voxel
small_mask = np.zeros(whole_brain_mask.shape)
small_mask[80, 54, 9] = 1

# Preset the variables
data = bold_vol
mask = small_mask
bcvar = labels
sl_rad = 1
max_blk_edge = 5
pool_size = 1

# We will get back to these commands once we finish running a simple searchlight. 
#comm = MPI.COMM_WORLD
#rank = comm.rank
#size = comm.size

# Start the clock to time searchlight
begin_time = time.time()

2.3.2 Create Searchlight Object

In [ ]:
# Create the searchlight object
sl = Searchlight(sl_rad=sl_rad,max_blk_edge=max_blk_edge)
print("Setup searchlight inputs")
print("Input data shape: " + str(data.shape))
print("Input mask shape: " + str(mask.shape) + "\n")

# Distribute the information to the searchlights (preparing it to run)
sl.distribute([data], mask)

# Data that is needed for all searchlights is sent to all cores via the sl.broadcast function. In this example, we are sending the labels for classification to all searchlights.
sl.broadcast(bcvar)

2.3.3 Define the function (a.k.a kernel) that needs to be computed

This is the function that you want to measure/classify your data with. This could perform classification, RSA, or any other computation that you wish.

NB. The work kernel is used in multiple mathematical/computational contexts.

In [ ]:
# Set up the kernel, in this case an SVM
def calc_svm(data, sl_mask, myrad, bcvar):
    
    # Pull out the data
    data4D = data[0]
    labels = bcvar
    
    bolddata_sl = data4D.reshape(sl_mask.shape[0] * sl_mask.shape[1] * sl_mask.shape[2], data[0].shape[3]).T

    # Check if the number of voxels is what you expect.
    print("Searchlight data shape: " + str(data[0].shape))
    print("Searchlight data shape after reshaping: " + str(bolddata_sl.shape))
    print("Searchlight mask shape:" + str(sl_mask.shape) + "\n")
    print("Searchlight mask (note that the center equals 1):\n" + str(sl_mask) + "\n")
    
    t1 = time.time()
    clf = SVC(kernel='linear', C=1)
    scores = cross_val_score(clf, bolddata_sl, labels, cv=3)
    accuracy = scores.mean()
    t2 = time.time()
    
    print('Kernel duration: ' + str(t2 - t1) + "\n\n")
    
    return accuracy

2.3.4 Execute the Searchlight

In [ ]:
print("Begin SearchLight\n")
sl_result = sl.run_searchlight(calc_svm, pool_size=pool_size)
print("End SearchLight\n")

end_time = time.time()

# Print outputs
print("Summarize searchlight results")
print("Number of searchlights run: " + str(len(sl_result[mask==1])))
print("Accuracy for each kernel function: " + str(sl_result[mask==1]))
print('Total searchlight duration (including start up time): ' + str(end_time - begin_time))

# Save the results to a .nii file
output_name = ('%s_SL_result.nii.gz' % (sub))
sl_result = sl_result.astype('double')  # Convert the output into a precision format that can be used by other applications
sl_result[np.isnan(sl_result)] = 0  # Exchange nans with zero to ensure compatibility with other applications
sl_nii = nib.Nifti1Image(sl_result, affine_mat)  # create the volume image
hdr = sl_nii.header  # get a handle of the .nii file's header
hdr.set_zooms((dimsize[0], dimsize[1], dimsize[2]))
nib.save(sl_nii, output_name)  # Save the volume

For some of the following exercises you want to copy and paste the code from sections 2.3.1 - 2.3.5 into the cell below in order to print out all of the required outputs. Also, before finalizing your answers, please remove the comments from within the kernel that print the data shape and mask variable, to keep your output clean.

Exercise 3: From running the code in one single chunck in a new cell below, make an estimate of how long it would take to run the searchlight on the whole brain, using the same parameters.

A:

Exercise 4: What happens if you set sl_rad to 3?

A:

Exercise 5: The first input to the Searchlight.distribute is a list. Why would we have multiple entries? What could this be useful for?

A:

Exercise 6: Modify the small_mask variable to select a 4x4x4 sub-volume of 64 voxels around the voxel in the example above.

Exercise 7: Modify the calc_svm function to run only if there are at least 14 voxels from the mask in the 27-voxel searchlight centered on each of these 64 voxels, otherwise return -1. This could be useful if you wanted to ignore searchlights near the edge of the brain.

Exercise 8: Modify the calc_svm function to perform nested cross-validation (use the code you made for the notebook on classification) over different cost parameters of the linear kernel.

In [ ]:
# Paste code here for the excercises

Visualizing the data

As you have seen, the output of the searchlight analysis is a 3D volume with the outputs of your kernel for each voxel it was centered on. The nilearn package has multiple plotting options to allow you to view your data within python. This is good for quick visualizations but can be buggy and is not great for exploration. You can also view the data in other packages such as FSL or AFNI. For instance on Milgram, you can run module load Apps/FSL/5.0.6; fslview sub-01_SL_result.nii.gz to view the results. If you followed the steps above, you likely won't see many bright spots because we only ran the searchlight in a tiny part of the brain. But if you look at the data in fslview and change the coordinates to X=80, Y=54, Z=9 you should see a non-zero value.

3. Running Searchlights on High Performance Clusters

Running searchlight analyses through notebooks or interactive sessions isn't tractable for real studies. Although the example above ran quickly and without parallelization, we only performed 64 analyses. Thus, we are now going to write a script to run a searchlight as a "batch" job. To learn how to submit jobs, you need to know a bit about slurm, the scheduling system used on Milgram. But don't worry, you have been submitting batch jobs via slurm since the first class -- every time you use ./launch_jupyter.sh, you submit a job.

To run a job, a good work flow is to have two scripts: One script that actually does the computation you care about (e.g., a python script like utils.py) and a bash script that sets up the environment and specifies the job parameters. The environment refers to the modules and packages you need to run your job. The job parameters refer to the partition you are going to use (-p), the number of cores (-n), the amount of memory (-m) and required time (-t). To run your job you then call the bash script with something like: 'sbatch script.sh'

Lucky for you we have already written the script needed here, called run_searchlight.sh. This script is written in the bash command language. Please explore this script to get familar with submitting jobs. You won't need to edit this script for the exercises below, but it will be very useful for future analyses to customize it for your needs (using a text editor like nano or nedit).

In [ ]:
# Print the bash script for running searchlights
!cat 07-searchlight/run_searchlight.sh

When a job starts, a log file is created that you can look at. You can see any output from your script printed into the log file. To check the status of your job you should use "squeue -u $USER". Sometimes your jobs won't run immediately because they are waiting in the queue. Conversely, don't submit too many jobs or your classmates won't be able to run theirs.

When you parallelize an analysis across multiple cores, this typically means that each core will run the code independently (though there are ways for cores to communicate directly with each other). This means you will load in all of the data on each core. The searchlight function then notices that there are multiple cores running the same job and assigns different pieces of the data to each core for analysis. The message passing interface (MPI), the parallelizing framework used by BrainIAK and most high-performance computing applications, keeps track of each core by assigning it a rank, starting at 0. In other words, if you are running an analysis across 2 cores then some of your computations will be run with rank=0 and some with rank=1. After the searchlights have all finished, you need to save the data. MPI makes the results available on all of the cores, but we typically save the results from the rank-0 core.

Self-study: Check that MPI is handling ranks and blocks as you expected. To do this create a searchlight that outputs the rank (doesn't do anything with the input data). To set this up you probably want to pull out the rank from the global variable MPI ('rank=MPI.COMM_WORLD.rank') on each searchlight and return the rank as an output. You should get different voxels with different rank values in the range of the number of cores you ran and this should be blocky such that each block has the side length of max_blk_edge.

We are now going to write the python script for running searchlights. The file searchlight.py has been started with some of the code needed. However, it is missing the meat and potatos.

In [ ]:
# Print the python script for running searchlights
!cat 07-searchlight/searchlight.py

Exercise 9: Finish the searchlight.py script so that you can run an SVM searchlight (linear, no grid search) like above but over the whole brain. To run the searchlight command as a batch job, enter into your terminal "sbatch run_searchlight.sh". Please make sure no comments are printed out in your kernel

Note that if you run the scripts on milgram but do not participate in the course, you will run into this error at some point: "sbatch: error: Batch job submission failed: Invalid account or account/partition combination specified". This happens because the .sh-scripts are set up to run on a dedicated partition for the course. To solve the error, change the .sh-script file (e.g. using the nano command) and change the line #SBATCH --partition cmhn-s18 to #SBATCH --partition short and delete the line #SBATCH -A cmhn-s18 before running the scripts.

Self-study: Another way to check the searchlight function is to give it noise and see whether you get null results. Try rerunning the searchlight but shuffle the label assignment between conditions.

Exercise 10: Once your analysis has finished (should take ~5 mins), load in the searchlight results and plot them with the imported function plotting from nilearn. However, be aware that some of these functions assume the brain is MNI space. If you think this data is not in MNI space then do not use those plotting tools.

Novel contribution: be creative and make one new discovery by adding an analysis, visualization, or optimization.

Contributions

M. Kumar, C. Ellis and N. Turk-Browne produced the initial notebook
T. Meissner minor edits