NeMoS Advanced: Cross-Validation and Model Selection

Newer versions of Jupyter Lab have addressed an issue with skipping around the notebook while scrolling. To make sure this fix is enabled, in the Jupyter Lab GUI, navigate to Settings > Settings Editor > Notebook and scroll down to the Windowing mode setting and make sure it is set to contentVisibility.

Also reminder to presenter: Go to View > Appearance, select Simple Interface and turn off everything else to hide as many bars as possible. And maybe activate Presentation Mode.

And turn on View > Render side-by-side (shortcut Shift+R).

NeMoS Advanced: Cross-Validation and Model Selection#

This notebook has had all its explanatory text removed and has not been run. It is intended to be downloaded and run locally (or on the provided binder) while listening to the presenter’s explanation. In order to see the fully rendered of this notebook, go here

Learning Objectives#

In this tutorial we will keep working on the hippocampal place field recordings with the goal of learning how to combine NeMoS and scikit-learn to perform cross-validation and model selection. In particular we will:

Learn how to use NeMoS objects with scikit-learn for cross-validation
Learn how to use NeMoS objects with scikit-learn pipelines
Learn how to use cross-validation to perform model and feature selection. More specifically, we will compare models including position and speed as predictors with model including only speed or only position.

Pre-Processing#

Let’s first load and wrangle the data with pynapple and NeMoS. You can run the following cells for preparing the variables that we are going to use in the notebook and recapitulate the content of this dataset with a few visualizations.

import workshop_utils
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pynapple as nap

import nemos as nmo

# some helper plotting functions
from nemos import _documentation_utils as doc_plots

# configure plots some
plt.style.use(nmo.styles.plot_style)

import workshop_utils

from sklearn import model_selection
from sklearn import pipeline

# shut down jax to numpy conversion warning
nap.nap_config.suppress_conversion_warnings = True

Load the data using pynapple.

path = workshop_utils.fetch_data("Achilles_10252013_EEG.nwb")
data = nap.load_file(path)
data

Extract the spike times and mouse position.

spikes = data["units"]
position = data["position"]

Restrict data to when animal was traversing the linear track.

position = position.restrict(data["forward_ep"])
spikes = spikes.restrict(data["forward_ep"])

Restrict neurons to only excitatory neurons, discarding neurons with a low-firing rate.

spikes = spikes.getby_category("cell_type")["pE"]
spikes = spikes.getby_threshold("rate", 0.3)

Place fields#

Visualize the place fields: neuronal firing rate as a function of position.

place_fields = nap.compute_tuning_curves(spikes, position, bins=50, epochs=position.time_support, feature_names=["distance"])
workshop_utils.plot_place_fields(place_fields)

For speed, we’re only going to investigate the three neurons highlighted above.
Bin spikes to counts at 100 Hz.
Interpolate position to match spike resolution.

neurons = [82, 92, 220]
place_fields = place_fields.sel(unit=neurons)
spikes = spikes[neurons]
bin_size = .01
count = spikes.count(bin_size, ep=position.time_support)
position = position.interpolate(count, ep=count.time_support)
print(count.shape)
print(position.shape)

Extract Speed per Epoch#

Compute the animal’s speed.
Visualize tuning curves to speed and position.

speed = position.derivative()
print(speed.shape)

# utility function to visualize predictions
tc_speed = nap.compute_tuning_curves(spikes, speed, bins=20, epochs=speed.time_support, feature_names=["speed"])
fig = workshop_utils.plot_position_speed(position, speed, place_fields, tc_speed, neurons);

def visualize_model_predictions(glm, X):
    # predict the model's firing rate
    predicted_rate = glm.predict(X) / bin_size

    # compute the position and speed tuning curves using the predicted firing rate.
    glm_pos = nap.compute_tuning_curves(predicted_rate, position, bins=50, epochs=position.time_support, feature_names=["position"])
    glm_speed = nap.compute_tuning_curves(predicted_rate, speed, bins=30, epochs=position.time_support, feature_names=["speed"])

    workshop_utils.plot_position_speed_tuning(place_fields, tc_speed, glm_pos, glm_speed);

Define 1D NeMoS Bases#

Define the position and speed bases, and visualize them.

position_basis = nmo.basis.MSplineEval(n_basis_funcs=10, label="position")
speed_basis = nmo.basis.MSplineEval(n_basis_funcs=15, label="speed")
workshop_utils.plot_pos_speed_bases(position_basis, speed_basis);

Basis Composition#

Adding the position and speed bases together defines a 2D basis.
Call compute_features to define a design matrix that concatenates both features.

basis = position_basis + speed_basis

X = basis.compute_features(position, speed)
X_numpy = np.concatenate(
    [
        position_basis.compute_features(position),
        speed_basis.compute_features(speed),
    ],
    axis=1
)

print("Are the design matrices equivalent?", np.all(X.d == X_numpy.d))

Scikit-learn#

How to know when to regularize?#

How do we decide when to use regularization?
Cross-validation allows you to fairly compare different models on the same dataset.
NeMoS makes use of scikit-learn, the standard machine learning library in python.
Define parameter grid to search over.
Anything not specified in grid will be kept constant.

# define a Ridge PopulationGLM
glm = nmo.glm.PopulationGLM(
    regularizer="Ridge",
    solver_kwargs={"tol": 1e-12},
    solver_name="LBFGS",
)
param_grid = {
    "regularizer_strength": [0.0001, 1.],
}

Initialize scikit-learn’s model_selection.GridSearchCV object.

cv_folds = 5
cv = model_selection.GridSearchCV(glm, param_grid, cv=cv_folds)
cv

We interact with this in a very similar way to the glm object.
In particular, call fit with same arguments:

cv.fit(X, count)

Let’s investigate results:

pd.DataFrame(cv.cv_results_)

Find the best regularization strength!

As an exercise, spend 10 minutes trying to find the best regularization strength!

You should use the glm model we defined in this section.
You will need to redefine the param_grid dictionary, selecting different values for "regularizer_strength":

param_grid = {
    "regularizer_strength": ...,
}

After defining param_grid, reinitialize cv (you can do so with the same arguments).
Then call cv.fit and re-run pd.DataFrame(cv.cv_results_) to summarize the results.

Who can find the best regularization strength?

If you finish early, try out different regularizers and try to find the best regularization strength for each of them.

Select basis#

You can (and should) do something similar to determine how many basis functions you need for each input.
NeMoS basis objects are not scikit-learn-compatible right out of the box.
But we have provided a simple method to make them so:

position_basis = nmo.basis.MSplineEval(n_basis_funcs=10, label="position").to_transformer()
# or equivalently:
position_basis = nmo.basis.TransformerBasis(nmo.basis.MSplineEval(n_basis_funcs=10, label="position"))
position_basis

This gives the basis object the transform method, which is equivalent to compute_features.
However, transformers have some limits:

position_basis.transform(position)

Transformers only accept 2d inputs, whereas nemos basis objects can accept inputs of any dimensionality.
In order to use a basis as a transformer, you’ll need to concatenate all your input in a single 2D array.

position_basis.transform(position[:, np.newaxis])

Let’s now create the composite basis for speed and position.

position_basis = nmo.basis.MSplineEval(n_basis_funcs=10, label="position")
speed_basis = nmo.basis.MSplineEval(n_basis_funcs=15, label="speed")
basis = position_basis + speed_basis
basis = basis.to_transformer()
basis

Stack position and speed in a single TsdFrame to hold all our inputs:

transformer_input = nap.TsdFrame(
    t=position.t,
    d=np.stack([position, speed]).T,
    time_support=position.time_support,
    columns=["position", "speed"],
)

Pass this input to our transformed additive basis.

basis.transform(transformer_input)

Pipelines#

If we want to cross-validate over the basis, we need more one more step: combining the basis and the GLM into a single scikit-learn estimator.
Pipelines to the rescue!

# set the reg strength to the optimal
glm = nmo.glm.PopulationGLM(solver_name="LBFGS", solver_kwargs={"tol": 10**-12})
pipe = pipeline.Pipeline([
    ("basis", basis),
    ("glm", glm)
])
pipe

Pipeline runs basis.transform, then passes that output to glm, so we can do everything in a single line:

pipe.fit(transformer_input, count)

Visualize model predictions!

visualize_model_predictions(pipe, transformer_input)

Cross-validating on the basis#

Now that we have our pipeline estimator, we can cross-validate on any of its parameters!

pipe.steps

Let’s cross-validate on:

The number of the basis functions of the position basis
The functional form of the basis for speed
Let’s retrieve the those attributes from the pipeline

# the label of the pipeline step retrieves the basis
print(pipe["basis"])

# the position basis can by retreived by its label
print("\n", pipe["basis"]["position"])

# the n_basis_funcs is an attribute
print("\n", pipe["basis"]["position"].n_basis_funcs)

# with the same syntax we can retreive the speed basis
print("\n", pipe["basis"]["speed"])

Construct param_grid, using __ to stand in for .
In scikit-learn pipelines, we access nested parameters using double underscores:
- pipe["basis"]["position"].n_basis_funcs - normal Python syntax
- "basis__position__n_basis_funcs" - scikit-learn parameter grid syntax

param_grid = {
    "basis__position__n_basis_funcs": [5, 10, 20],
    "basis__speed": [nmo.basis.MSplineEval(15),
                      nmo.basis.BSplineEval(15),
                      nmo.basis.RaisedCosineLinearEval(15)],
}

Cross-validate as before:

cv = model_selection.GridSearchCV(pipe, param_grid, cv=cv_folds)
cv.fit(transformer_input, count)

Investigate results:

pd.DataFrame(cv.cv_results_)

Can easily grab the best estimator, the pipeline that did the best:

best_estim = cv.best_estimator_
best_estim

Visualize model predictions!

visualize_model_predictions(best_estim, transformer_input)

Find the best basis!

As an exercise, spend 10 minutes exploring the possible basis objects and seeing which performs the best.

You should use the pipe object we defined in this section.
You will need to redefine the param_grid dictionary, setting basis__speed and basis__position (or their attributes, e.g., basis__position__n_basis_funcs) to a range of values. Remember that all combinations are tested, so if you e.g., select 5 choices for each, you’ll be testing 25 different combinations!

param_grid = {
    "basis__position": ...,
    "basis__speed": ...,
}

After defining param_grid, reinitialize cv (you can do so with the same arguments).
Then call cv.fit and re-run pd.DataFrame(cv.cv_results_) to summarize the results.
Finally, visualize the best estimator with visualize_model_predictions(cv.best_estimator_, transformer_input)

Who can find the best set of basis objects?

Feature selection#

Let’s move on to feature selection. Our goal is to compare alternative models: position + speed, position only, or speed only.

Problem: scikit-learn’s cross-validation assumes the pipeline input stays constant, but each model needs different features. How do we solve this?

Solution: Use a “null” basis that produces zero features!

We’ll create this null basis using CustomBasis, which defines a basis from custom functions.

# define a function that creates an empty array (n_samples, 0)
def func(x):
    return np.zeros((x.shape[0], 0))

# create a null transformer basis using the custom basis class
null_basis = nmo.basis.CustomBasis([func]).to_transformer()

# verify: this creates an empty feature array
null_basis.compute_features(position).shape

Why is this useful? We can combine null_basis with actual bases to create different models that all accept the same input!

Let’s define the bases for our three models:

Position + speed: combine position and speed bases
Position only: combine position basis with null basis (speed features is empty)
Speed only: combine null basis with speed basis (position features is empty)

# combine them to define each model
basis_all = (position_basis + speed_basis).to_transformer()
basis_position = (position_basis + null_basis).to_transformer()
basis_speed = (null_basis + speed_basis).to_transformer()

# assign labels (optional but helpful for readability)
basis_all.label = "position + speed"
basis_position.label = "position only"
basis_speed.label = "speed only"

These bases can all transform the same transformer_input (a TsdFrame with columns for position and speed), but they generate design matrices with different numbers of features:

# "position + speed" design: 25 features (10 + 15)
print("position + speed design matrix shape:")
print(basis_all.transform(transformer_input).shape)

# "position" design: 10 features (10 + 0)
print("\nposition design matrix shape:")
print(basis_position.transform(transformer_input).shape)

# "speed" design: 15 features (0 + 15)
print("\nspeed design matrix shape:")
print(basis_speed.transform(transformer_input).shape)

To cross-validate over different basis compositions, we need to understand how they’re stored in our pipeline. The additive basis is stored as a basis attribute inside the TransformerBasis object:

# the "basis" step in our pipeline contains a TransformerBasis
# which has a "basis" attribute storing the actual additive basis
pipe["basis"].basis

Now we can create a parameter grid for cross-validation. The key is the string "basis__basis":

First basis: the name of the pipeline step
Second basis: the attribute of the TransformerBasis object
This double-underscore notation is how scikit-learn accesses nested parameters

# create parameter grid with our three basis compositions
param_grid = {
    "basis__basis": [
        basis_all,      # position + speed
        basis_position, # position only
        basis_speed     # speed only
    ],
}

# define and fit GridSearchCV
cv = model_selection.GridSearchCV(pipe, param_grid, cv=cv_folds)
cv.fit(transformer_input, count)

Let’s examine the model comparison results:

cv_df = pd.DataFrame(cv.cv_results_)

# display the key columns: which basis was used, its score, and ranking
cv_df[["param_basis__basis", "mean_test_score", "rank_test_score"]]

Position emerges as the predictor with the greatest explanatory power, while speed adds only marginal benefits.

Find the model!

In this section, we only compared a single choice of regularization strength and basis objects for each feature. As an exercise, spend 10 minutes combining what we learned here with the earlier sections: for each feature combination (position, speed, position + speed), try several different basis objects and, optionally, different regularization strengths.

Don’t forget to visualize your model’s predictions!

Who can find the best model?

Next Steps#

For the next project, you can use all the tools showcased here to find a better encoding model for these hippocampal neurons.

Suggestions:

Extend the model by including theta phase as a predictor
Use the NeMoS MultiplicativeBasis to capture interactions between theta phase and position

References#

The data in this tutorial comes from Grosmark, Andres D., and György Buzsáki. “Diversity in neural firing dynamics supports both rigid and learned hippocampal sequences.” Science 351.6280 (2016): 1440-1443.