def get_test_metric(hparams,
model_version,
metric='r2',
dtype='test',
multioutput='variance_weighted',
sess_idx=0):
"""Calculate a single R\ :sup:`2` value across all test batches for a decoder.
Parameters
----------
hparams : :obj:`dict`
needs to contain enough information to specify an autoencoder
model_version : :obj:`int` or :obj:`str`
version from test tube experiment defined in :obj:`hparams` or the string 'best'
metric : :obj:`str`, optional
'r2' | 'fc' | 'mse'
dtype : :obj:`str`
type of trials to use for computing metric
'train' | 'val' | 'test'
multioutput : :obj:`str`
defines how to aggregate multiple r2 scores; see r2_score documentation in sklearn
'raw_values' | 'uniform_average' | 'variance_weighted'
sess_idx : :obj:`int`, optional
session index into data generator
Returns
-------
:obj:`tuple`
- hparams (:obj:`dict`): hparams of model used to calculate metrics
- metric (:obj:`int`)
"""
from sklearn.metrics import r2_score, accuracy_score
from behavenet.fitting.utils import get_best_model_and_data
from behavenet.models import Decoder
model, data_generator = get_best_model_and_data(hparams,
Decoder,
load_data=True,
version=model_version)
n_test_batches = len(data_generator.datasets[sess_idx].batch_idxs[dtype])
max_lags = hparams['n_max_lags']
true = []
pred = []
data_generator.reset_iterators(dtype)
for i in range(n_test_batches):
batch, _ = data_generator.next_batch(dtype)
# get true latents/states
if metric == 'r2' or metric == 'mse':
if 'ae_latents' in batch:
curr_true = batch['ae_latents'][0].cpu().detach().numpy()
elif 'labels' in batch:
curr_true = batch['labels'][0].cpu().detach().numpy()
else:
raise ValueError('no valid key in {}'.format(batch.keys()))
elif metric == 'fc':
curr_true = batch['arhmm_states'][0].cpu().detach().numpy()
else:
raise ValueError('"%s" is an invalid metric type' % metric)
# get predicted latents
curr_pred = model(batch['neural'][0])[0].cpu().detach().numpy()
true.append(curr_true[max_lags:-max_lags])
pred.append(curr_pred[max_lags:-max_lags])
if metric == 'r2':
metric = r2_score(np.concatenate(true, axis=0),
np.concatenate(pred, axis=0),
multioutput=multioutput)
elif metric == 'mse':
metric = np.mean(
np.square(
np.concatenate(true, axis=0) - np.concatenate(pred, axis=0)))
elif metric == 'fc':
metric = accuracy_score(
np.concatenate(true, axis=0),
np.argmax(np.concatenate(pred, axis=0), axis=1))
return model.hparams, metric, true, pred
def make_ae_reconstruction_movie_wrapper(hparams,
save_file,
trial=None,
sess_idx=0,
version='best',
include_linear=False,
max_frames=400,
frame_rate=15):
"""Produce movie with original video, reconstructed video, and residual.
This is a high-level function that loads the model described in the hparams dictionary and
produces the necessary predicted video frames.
Parameters
----------
hparams : :obj:`dict`
needs to contain enough information to specify an autoencoder
save_file : :obj:`str`
full save file (path and filename)
trial : :obj:`int`, optional
if :obj:`NoneType`, use first test trial
sess_idx : :obj:`int`, optional
session index into data generator
version : :obj:`str` or :obj:`int`, optional
test tube model version
include_linear : :obj:`bool`, optional
include reconstruction from corresponding linear ae (i.e. ame number of latents)
max_frames : :obj:`int`, optional
maximum number of frames to animate from a trial
frame_rate : :obj:`float`, optional
frame rate of saved movie
"""
# build model(s)
if hparams['model_class'] == 'ae':
from behavenet.models import AE as Model
elif hparams['model_class'] == 'cond-ae':
from behavenet.models import ConditionalAE as Model
else:
raise NotImplementedError('"%s" is an invalid model class' %
hparams['model_class'])
model_ae, data_generator = get_best_model_and_data(hparams,
Model,
version=version)
if include_linear:
import copy
hparams_lin = copy.copy(hparams)
hparams_lin['model_type'] = 'linear'
if 'lin_experiment_name' in hparams:
hparams_lin['experiment_name'] = hparams['lin_experiment_name']
model_lin, _ = get_best_model_and_data(hparams_lin,
Model,
load_data=False,
version=version)
else:
model_lin = None
# push images through decoder
if trial is None:
# choose first test trial
trial = data_generator.batch_idxs[sess_idx]['test'][0]
batch = data_generator.datasets[sess_idx][trial]
ims_orig_pt = batch['images'][:max_frames]
if hparams['model_class'] == 'cond-ae':
labels_pt = batch['labels'][:max_frames]
else:
labels_pt = None
ims_recon_ae = get_reconstruction(model_ae, ims_orig_pt, labels=labels_pt)
if include_linear:
ims_recon_lin = get_reconstruction(model_lin,
ims_orig_pt,
labels=labels_pt)
else:
ims_recon_lin = None
# mask images for plotting
if hparams.get('use_output_mask', False):
ims_orig_pt *= batch['masks'][:max_frames]
ims_orig = ims_orig_pt.cpu().detach().numpy()
ims = [ims_orig, ims_recon_ae, 0.5 + (ims_orig - ims_recon_ae)]
titles = ['Original', 'Conv AE reconstructed', 'Conv AE residual']
if include_linear:
ims.append([])
ims.append(ims_recon_lin)
ims.append(0.5 + (ims_orig - ims_recon_lin))
titles.append('')
titles.append('Linear AE reconstructed')
titles.append('Linear AE residual')
n_rows = 2
n_cols = 3
else:
n_rows = 1
n_cols = 3
make_reconstruction_movie(ims=ims,
titles=titles,
n_rows=n_rows,
n_cols=n_cols,
save_file=save_file,
frame_rate=frame_rate)
def make_neural_reconstruction_movie_wrapper(hparams,
save_file,
trial=None,
sess_idx=0,
max_frames=400,
max_latents=8,
frame_rate=15):
"""Produce movie with original video, ae reconstructed video, and neural reconstructed video.
This is a high-level function that loads the model described in the hparams dictionary and
produces the necessary predicted video frames. Latent traces are additionally plotted, as well
as the residual between the ae reconstruction and the neural reconstruction. Currently produces
ae latents and decoder predictions from scratch (rather than saved pickle files).
Parameters
----------
hparams : :obj:`dict`
needs to contain enough information to specify an autoencoder
save_file : :obj:`str`
full save file (path and filename)
trial : :obj:`int`, optional
if :obj:`NoneType`, use first test trial
sess_idx : :obj:`int`, optional
session index into data generator
max_frames : :obj:`int`, optional
maximum number of frames to animate from a trial
max_latents : :obj:`int`, optional
maximum number of ae latents to plot
frame_rate : :obj:`float`, optional
frame rate of saved movie
"""
from behavenet.models import Decoder
from behavenet.models import AE
###############################
# build ae model/data generator
###############################
hparams_ae = copy.copy(hparams)
hparams_ae['experiment_name'] = hparams['ae_experiment_name']
hparams_ae['model_class'] = hparams['ae_model_class']
hparams_ae['model_type'] = hparams['ae_model_type']
model_ae, data_generator_ae = get_best_model_and_data(
hparams_ae, AE, version=hparams['ae_version'])
# move model to cpu
model_ae.to('cpu')
if trial is None:
# choose first test trial
trial = data_generator_ae.batch_idxs[sess_idx]['test'][0]
# get images from data generator (move to cpu)
batch = data_generator_ae.datasets[sess_idx][trial]
ims_orig_pt = batch['images'][:max_frames].cpu() # 400
# push images through ae to get reconstruction
ims_recon_ae = get_reconstruction(model_ae, ims_orig_pt)
# push images through ae to get latents
latents_ae_pt, _, _ = model_ae.encoding(ims_orig_pt)
# mask images for plotting
if hparams.get('use_output_mask', False):
ims_orig_pt *= batch['masks'][:max_frames]
#######################################
# build decoder model/no data generator
#######################################
hparams_dec = copy.copy(hparams)
hparams_dec['experiment_name'] = hparams['decoder_experiment_name']
hparams_dec['model_class'] = hparams['decoder_model_class']
hparams_dec['model_type'] = hparams['decoder_model_type']
model_dec, data_generator_dec = get_best_model_and_data(
hparams_dec, Decoder, version=hparams['decoder_version'])
# move model to cpu
model_dec.to('cpu')
# get neural activity from data generator (move to cpu)
batch = data_generator_dec.datasets[0][
trial] # 0 not sess_idx since decoders only have 1 sess
neural_activity_pt = batch['neural'][:max_frames].cpu()
# push neural activity through decoder to get prediction
latents_dec_pt, _ = model_dec(neural_activity_pt)
# push prediction through ae to get reconstruction
ims_recon_dec = get_reconstruction(model_ae, latents_dec_pt)
# away
make_neural_reconstruction_movie(
ims_orig=ims_orig_pt.cpu().detach().numpy(),
ims_recon_ae=ims_recon_ae,
ims_recon_neural=ims_recon_dec,
latents_ae=latents_ae_pt.cpu().detach().numpy()[:, :max_latents],
latents_neural=latents_dec_pt.cpu().detach().numpy()[:, :max_latents],
save_file=save_file,
frame_rate=frame_rate)
def make_neural_reconstruction_movie_wrapper(hparams,
save_file,
trials=None,
sess_idx=0,
max_frames=400,
max_latents=8,
zscore_by_dim=False,
colored_predictions=False,
xtick_locs=None,
frame_rate=15):
"""Produce movie with original video, ae reconstructed video, and neural reconstructed video.
This is a high-level function that loads the model described in the hparams dictionary and
produces the necessary predicted video frames. Latent traces are additionally plotted, as well
as the residual between the ae reconstruction and the neural reconstruction. Currently produces
ae latents and decoder predictions from scratch (rather than saved pickle files).
Parameters
----------
hparams : :obj:`dict`
needs to contain enough information to specify an autoencoder
save_file : :obj:`str`
full save file (path and filename)
trials : :obj:`int` or :obj:`list`, optional
if :obj:`NoneType`, use first test trial
sess_idx : :obj:`int`, optional
session index into data generator
max_frames : :obj:`int`, optional
maximum number of frames to animate from a trial
max_latents : :obj:`int`, optional
maximum number of ae latents to plot
zscore_by_dim : :obj:`bool`, optional
True to z-score each dim, False to leave relative scales
colored_predictions : :obj:`bool`, optional
False to plot reconstructions in black, True to plot in different colors
xtick_locs : :obj:`array-like`, optional
tick locations in units of bins
frame_rate : :obj:`float`, optional
frame rate of saved movie
"""
from behavenet.models import Decoder
# define number of frames that separate trials
n_buffer = 5
###############################
# build ae model/data generator
###############################
hparams_ae = copy.copy(hparams)
hparams_ae['experiment_name'] = hparams['ae_experiment_name']
hparams_ae['model_class'] = hparams['ae_model_class']
hparams_ae['model_type'] = hparams['ae_model_type']
model_ae, data_generator_ae = get_best_model_and_data(
hparams_ae, Model=None, version=hparams['ae_version'])
# move model to cpu
model_ae.to('cpu')
#######################################
# build decoder model/no data generator
#######################################
hparams_dec = copy.copy(hparams)
hparams_dec['experiment_name'] = hparams['decoder_experiment_name']
hparams_dec['model_class'] = hparams['decoder_model_class']
hparams_dec['model_type'] = hparams['decoder_model_type']
model_dec, data_generator_dec = get_best_model_and_data(
hparams_dec, Decoder, version=hparams['decoder_version'])
# move model to cpu
model_dec.to('cpu')
if trials is None:
# choose first test trial, put in list
trials = data_generator_ae.batch_idxs[sess_idx]['test'][0]
if isinstance(trials, int):
trials = [trials]
# loop over trials, putting black frames/nans in between
ims_orig = []
ims_recon_ae = []
ims_recon_neural = []
latents_ae = []
latents_neural = []
for i, trial in enumerate(trials):
# get images from data generator (move to cpu)
batch = data_generator_ae.datasets[sess_idx][trial]
ims_orig_pt = batch['images'][:max_frames].cpu() # 400
if hparams_ae['model_class'] == 'cond-ae':
labels_pt = batch['labels'][:max_frames]
else:
labels_pt = None
# push images through ae to get reconstruction
ims_recon_ae_curr, latents_ae_curr = get_reconstruction(
model_ae, ims_orig_pt, labels=labels_pt, return_latents=True)
# mask images for plotting
if hparams_ae.get('use_output_mask', False):
ims_orig_pt *= batch['masks'][:max_frames]
# get neural activity from data generator (move to cpu)
# 0, not sess_idx, since decoders only have 1 sess
batch = data_generator_dec.datasets[0][trial]
neural_activity_pt = batch['neural'][:max_frames].cpu()
# push neural activity through decoder to get prediction
latents_dec_pt, _ = model_dec(neural_activity_pt)
# push prediction through ae to get reconstruction
ims_recon_dec_curr = get_reconstruction(model_ae,
latents_dec_pt,
labels=labels_pt)
# store all relevant quantities
ims_orig.append(ims_orig_pt.cpu().detach().numpy())
ims_recon_ae.append(ims_recon_ae_curr)
ims_recon_neural.append(ims_recon_dec_curr)
latents_ae.append(latents_ae_curr[:, :max_latents])
latents_neural.append(
latents_dec_pt.cpu().detach().numpy()[:, :max_latents])
# add blank frames
if i < len(trials) - 1:
n_channels, y_pix, x_pix = ims_orig[-1].shape[1:]
n = latents_ae[-1].shape[1]
ims_orig.append(np.zeros((n_buffer, n_channels, y_pix, x_pix)))
ims_recon_ae.append(np.zeros((n_buffer, n_channels, y_pix, x_pix)))
ims_recon_neural.append(
np.zeros((n_buffer, n_channels, y_pix, x_pix)))
latents_ae.append(np.nan * np.zeros((n_buffer, n)))
latents_neural.append(np.nan * np.zeros((n_buffer, n)))
latents_ae = np.vstack(latents_ae)
latents_neural = np.vstack(latents_neural)
if zscore_by_dim:
means = np.nanmean(latents_ae, axis=0)
std = np.nanstd(latents_ae, axis=0)
latents_ae = (latents_ae - means) / std
latents_neural = (latents_neural - means) / std
# away
make_neural_reconstruction_movie(
ims_orig=np.vstack(ims_orig),
ims_recon_ae=np.vstack(ims_recon_ae),
ims_recon_neural=np.vstack(ims_recon_neural),
latents_ae=latents_ae,
latents_neural=latents_neural,
ae_model_class=hparams_ae['model_class'].upper(),
colored_predictions=colored_predictions,
xtick_locs=xtick_locs,
frame_rate_beh=hparams['frame_rate'],
save_file=save_file,
frame_rate=frame_rate)
def make_ae_reconstruction_movie_wrapper(hparams,
save_file,
trial=None,
sess_idx=0,
version='best',
include_linear=False,
max_frames=400,
frame_rate=15):
"""Produce movie with original video, reconstructed video, and residual.
This is a high-level function that loads the model described in the hparams dictionary and
produces the necessary predicted video frames.
Parameters
----------
hparams : :obj:`dict`
needs to contain enough information to specify an autoencoder
save_file : :obj:`str`
full save file (path and filename)
trial : :obj:`int`, optional
if :obj:`NoneType`, use first test trial
sess_idx : :obj:`int`, optional
session index into data generator
version : :obj:`str` or :obj:`int`, optional
test tube model version
include_linear : :obj:`bool`, optional
include reconstruction from corresponding linear ae (i.e. ame number of latents)
max_frames : :obj:`int`, optional
maximum number of frames to animate from a trial
frame_rate : :obj:`float`, optional
frame rate of saved movie
"""
from behavenet.models import AE
# build model(s)
model_ae, data_generator = get_best_model_and_data(hparams,
AE,
version=version)
if include_linear:
import copy
hparams_lin = copy.copy(hparams)
hparams_lin['model_type'] = 'linear'
if 'lin_experiment_name' in hparams:
hparams_lin['experiment_name'] = hparams['lin_experiment_name']
model_lin, _ = get_best_model_and_data(hparams_lin,
AE,
load_data=False)
else:
model_lin = None
# push images through decoder
if trial is None:
# choose first test trial
trial = data_generator.batch_idxs[sess_idx]['test'][0]
batch = data_generator.datasets[sess_idx][trial]
ims_orig_pt = batch['images'][:max_frames]
ims_recon_ae = get_reconstruction(model_ae, ims_orig_pt)
if include_linear:
ims_recon_lin = get_reconstruction(model_lin, ims_orig_pt)
else:
ims_recon_lin = None
# mask images for plotting
if hparams.get('use_output_mask', False):
ims_orig_pt *= batch['masks'][:max_frames]
make_ae_reconstruction_movie(ims_orig=ims_orig_pt.cpu().detach().numpy(),
ims_recon_ae=ims_recon_ae,
ims_recon_lin=ims_recon_lin,
save_file=save_file,
frame_rate=frame_rate)