Pre-defined datasets#

Pre-defined demo and benchmark datasets.

This package contains actual implementations of datasets. If you want to add a commonly used (and public dataset) to CEBRA, this is the right package to do it. Datasets here can be loaded e.g. for testing, reproducing reference results and benchmarking. When contributing to this package, you should ensure that the data is publicly available under a suitable license.

This module is a registry and currently contains the options [‘demo-discrete’, ‘demo-continuous’, ‘demo-mixed’, ‘demo-discrete-multisession’, ‘demo-continuous-multisession’, ‘demo-continuous-unified’, ‘allen-movie-one-ca-VISp-10-train-10-111’, ‘allen-movie-one-ca-VISp-10-train-10-222’, ‘allen-movie-one-ca-VISp-10-train-10-333’, ‘allen-movie-one-ca-VISp-10-train-10-444’].

To retrieve a list of options, call:

>>> print(cebra.datasets.get_options())
['demo-discrete', 'demo-continuous', 'demo-mixed', ...]

To obtain an initialized instance, call cebra.datasets.init, defined in cebra.registry.add_helper_functions(). The first parameter to provide is the datasets name to use, which is one of the available options presented above. Then the required positional arguments specific to the module are provided, if needed.

You can register additional options by defining and registering classes with a name. To do that, you can add a decorator on top of it: @cebra.datasets.register("my-cebra-datasets").

Later, initialize your class similarly to the pre-defined options, using cebra.datasets.init with the datasets name set to my-cebra-datasets.

Note that these customized options will not be automatically added to this docstring.

Synthetic datasets#

class cebra.datasets.gaussian_mixture.ContinuousGaussianMixtureDataset(noise='poisson')#

Bases: SingleSessionDataset

A dataset of synthetically generated continuous labels and the corresponding 2D latents and 100D noisy observations.

Parameters:: noise (str) – The applied noise distribution applied.

property continuous_index#

The continuous index, if available.

The continuous index along with a similarity metric is used for drawing positive and/or negative samples.

Returns:: Tensor of shape (N,d), representing the index for all N samples in the dataset.

class cebra.datasets.gaussian_mixture.ContinuousGaussianMixtureDataset_continuous_gaussian_mixture_gaussian(*args, **kwargs)#

Bases: ContinuousGaussianMixtureDataset

A dataset of synthetically generated continuous labels and the corresponding 2D latents and 100D noisy observations.

Parameters:: noise – The applied noise distribution applied.

class cebra.datasets.gaussian_mixture.ContinuousGaussianMixtureDataset_continuous_gaussian_mixture_laplace(*args, **kwargs)#

Bases: ContinuousGaussianMixtureDataset

A dataset of synthetically generated continuous labels and the corresponding 2D latents and 100D noisy observations.

Parameters:: noise – The applied noise distribution applied.

class cebra.datasets.gaussian_mixture.ContinuousGaussianMixtureDataset_continuous_gaussian_mixture_poisson(*args, **kwargs)#

Bases: ContinuousGaussianMixtureDataset

A dataset of synthetically generated continuous labels and the corresponding 2D latents and 100D noisy observations.

Parameters:: noise – The applied noise distribution applied.

class cebra.datasets.gaussian_mixture.ContinuousGaussianMixtureDataset_continuous_gaussian_mixture_refractory_poisson(*args, **kwargs)#

Bases: ContinuousGaussianMixtureDataset

A dataset of synthetically generated continuous labels and the corresponding 2D latents and 100D noisy observations.

Parameters:: noise – The applied noise distribution applied.

class cebra.datasets.gaussian_mixture.ContinuousGaussianMixtureDataset_continuous_gaussian_mixture_uniform(*args, **kwargs)#

Bases: ContinuousGaussianMixtureDataset

A dataset of synthetically generated continuous labels and the corresponding 2D latents and 100D noisy observations.

Parameters:: noise – The applied noise distribution applied.

Rat Hippocampus dataset#

Rat hippocampus dataset

References

Grosmark, A.D., and Buzsáki, G. (2016). Diversity in neural firing dynamics supports both rigid and learned
hippocampal sequences. Science 351, 1440–1443.
Chen, Z., Grosmark, A.D., Penagos, H., and Wilson, M.A. (2016). Uncovering representations of sleep-associated
hippocampal ensemble spike activity. Sci. Rep. 6, 32193.
Grosmark, A.D., Long J. and Buzsáki, G. (2016); Recordings from hippocampal area CA1, PRE, during and POST novel spatial learning. CRCNS.org. http://dx.doi.org/10.6080/K0862DC5

class cebra.datasets.hippocampus.SingleRatDataset(name='achilles', root='data', download=True)#

Bases: SingleSessionDataset

A single rat hippocampus tetrode recording while the rat navigates on a linear track.

Neural data is spike counts binned into 25ms time window and the continuous behavior label is position and the running direction (left, right) of a rat. The behavior label is structured as 3D array consists of position, right, and left.

Parameters:: name – The name of the rat to use. Choose among ‘achilles’, ‘buddy’, ‘cicero’ and ‘gatsby’.

property continuous_index#

The continuous index, if available.

The continuous index along with a similarity metric is used for drawing positive and/or negative samples.

Returns:: Tensor of shape (N,d), representing the index for all N samples in the dataset.

decode(x_train, y_train, x_test, y_test)#

kNN decoding function.

Perform a kNN decoding for n_neighbors = 1,4,9,26,25 with the given train set and test set.

Parameters:

x_train – The train set data
y_train – The train set label
x_test – The test set data
y_test – The test set label

class cebra.datasets.hippocampus.SingleRatTrialSplitDataset(name='achilles', split_no=0, split=None, root='data')#

Bases: SingleRatDataset

A single rat hippocampus tetrode recording while the rat navigates on a linear track with 3-fold splits.

Neural data is spike counts binned into 25ms time window and the behavior is position and the running direction (left, right) of a rat. The behavior label is structured as 3D array consists of position, right, and left. The neural and behavior recordings are parsed into trials (a round trip from one end of the track) and the trials are split into a train, valid and test set with k=3 nested cross validation.

Parameters:

name – The name of a rat to use. Choose among ‘achilles’, ‘buddy’, ‘cicero’ and ‘gatsby’.
split_no – The k for k-fold split. Choose among 0, 1, 2.
split – The split to use. Choose among ‘train’, ‘valid’, ‘test’, ‘all’, and ‘wo_test’ (all trials except test split).

class cebra.datasets.hippocampus.SingleRatCorruptDataset(name, seed, root='data')#

Bases: SingleRatDataset

A single rat hippocampus tetrode recording while the rat navigates on a linear track with a shuffled behavior label.

Neural data is spike counts binned into 25ms time window and the behavior is position and the running direction (left, right) of a rat. The behavior label is structured as 3D array consists of position, right, and left and it is shuffled in random orders.

Parameters:

name – The name of the rat to use. Choose among ‘achilles’, ‘buddy’, ‘cicero’ and ‘gatsby’.
seed – The random seed to set the shuffling.

class cebra.datasets.hippocampus.MultipleRatsTrialSplitDataset(split_no=0, split=None)#

Bases: DatasetCollection

4 rats hippocampus tetrode recording while the rat navigates on a linear track with 3-fold splits.

Neural and behavior recordings of 4 rats. For each rat, neural data is spike counts binned into 25ms time window and the behavior is position and the running direction (left, right) of a rat. The behavior label is structured as 3D array consists of position, right, and left. Neural and behavior recordings of each rat are parsed into trials (a round trip from one end of the track) and the trials are split into a train, valid and test set with k=3 nested cross validation.

Parameters:

split_no – The k for k-fold split. Choose among 0, 1, and 2.
split – The split to use. Choose among ‘train’, ‘valid’, ‘test’, ‘all’, and ‘wo_test’ (all trials except test split).

class cebra.datasets.hippocampus.MultipleRatsTrialSplitDataset_rat_hippocampus_multisubjects_3fold_trial_split_0(*args, **kwargs)#

Bases: MultipleRatsTrialSplitDataset

4 rats hippocampus tetrode recording while the rat navigates on a linear track with 3-fold splits.

Neural and behavior recordings of 4 rats. For each rat, neural data is spike counts binned into 25ms time window and the behavior is position and the running direction (left, right) of a rat. The behavior label is structured as 3D array consists of position, right, and left. Neural and behavior recordings of each rat are parsed into trials (a round trip from one end of the track) and the trials are split into a train, valid and test set with k=3 nested cross validation.

Parameters:

split_no – The k for k-fold split. Choose among 0, 1, and 2.
split – The split to use. Choose among ‘train’, ‘valid’, ‘test’, ‘all’, and ‘wo_test’ (all trials except test split).

class cebra.datasets.hippocampus.MultipleRatsTrialSplitDataset_rat_hippocampus_multisubjects_3fold_trial_split_1(*args, **kwargs)#

Bases: MultipleRatsTrialSplitDataset

4 rats hippocampus tetrode recording while the rat navigates on a linear track with 3-fold splits.

Neural and behavior recordings of 4 rats. For each rat, neural data is spike counts binned into 25ms time window and the behavior is position and the running direction (left, right) of a rat. The behavior label is structured as 3D array consists of position, right, and left. Neural and behavior recordings of each rat are parsed into trials (a round trip from one end of the track) and the trials are split into a train, valid and test set with k=3 nested cross validation.

Parameters:

split_no – The k for k-fold split. Choose among 0, 1, and 2.
split – The split to use. Choose among ‘train’, ‘valid’, ‘test’, ‘all’, and ‘wo_test’ (all trials except test split).

class cebra.datasets.hippocampus.MultipleRatsTrialSplitDataset_rat_hippocampus_multisubjects_3fold_trial_split_2(*args, **kwargs)#

Bases: MultipleRatsTrialSplitDataset

4 rats hippocampus tetrode recording while the rat navigates on a linear track with 3-fold splits.

Neural and behavior recordings of 4 rats. For each rat, neural data is spike counts binned into 25ms time window and the behavior is position and the running direction (left, right) of a rat. The behavior label is structured as 3D array consists of position, right, and left. Neural and behavior recordings of each rat are parsed into trials (a round trip from one end of the track) and the trials are split into a train, valid and test set with k=3 nested cross validation.

Parameters:

split_no – The k for k-fold split. Choose among 0, 1, and 2.
split – The split to use. Choose among ‘train’, ‘valid’, ‘test’, ‘all’, and ‘wo_test’ (all trials except test split).

class cebra.datasets.hippocampus.SingleRatCorruptDataset_rat_hippocampus_achilles_corrupt_0(*args, **kwargs)#