braintools.analysis.population.decode

Contents

braintools.analysis.population.decode#

braintools.analysis.population.decode(x, y, n_classes=None, n_folds=5, model='lda', reg=0.001, key=None)#

Cross-validated linear decoding of integer labels from population activity.

The T samples are split into n_folds contiguous folds of equal size m = T // n_folds; the trailing T mod n_folds samples are dropped so that every fold has the same static shape. For each fold a linear classifier is fit in closed form on the other n_folds - 1 folds and its accuracy is measured on the held-out fold. The folds are processed in parallel with jax.vmap.

Two classifiers are available:

  • model='lda': shrinkage linear discriminant analysis [1] [2]. The pooled within-class covariance \(S\) (class covariances with ddof=0, weighted by the class priors) is shrunk towards a scaled identity, \(\Sigma = (1 - \text{reg})\,S + \text{reg}\, \frac{\operatorname{tr} S}{D} I\), and the class scores are \(x^\top \Sigma^{-1}\mu_k - \tfrac12 \mu_k^\top\Sigma^{-1}\mu_k + \log \pi_k\) with empirical priors \(\pi_k\).

  • model='ridge': one-vs-rest ridge regression [3] on +1/-1 targets with an unpenalised intercept, \(W = (X_c^\top X_c + \text{reg}\,I)^{-1} X_c^\top Y_c\); the predicted class has the largest score.

Parameters:

  • x (array_like or brainunit.Quantity) – Activity of shape (T, D) (samples by features); a 1-D array is treated as (T, 1). Quantities are decoded on their mantissa (the accuracy is dimensionless; for 'ridge' the penalty reg is expressed in the squared units of that mantissa).

  • y (array_like of int) – Labels of shape (T,) with values in [0, n_classes).

  • n_classes (int | None) – Number of classes (static). None means int(y.max()) + 1, which needs a concrete y: under jax.jit/jax.vmap it must be given explicitly.

  • n_folds (int) – Number of folds (static, >= 2).

  • model (str) – Classifier.

  • reg (float, default=1e-3) – Regularisation. For 'lda' it is the shrinkage intensity in [0, 1]; for 'ridge' it is the ridge penalty alpha >= 0.

  • key (Array | None) – PRNG key. If given, the samples are randomly permuted before folding (and before the trailing samples are dropped). If None the folds are contiguous blocks in the given order.

Returns:

Per-fold test accuracy, shape (n_folds,).

Return type:

Array

Raises:

  • ValueError – If the shapes are inconsistent, T < n_folds, model is unknown, a concrete reg is out of range, or n_classes is None while y is traced.

  • TypeError – If y is not an integer array or a static argument is not an int.

See also

mutual_information

Model-free dependence between labels and responses.

pca

Reduce dimensionality before decoding.

Notes

Correspondence with scikit-learn. On identical training/test splits the per-fold predictions equal those of

  • sklearn.discriminant_analysis.LinearDiscriminantAnalysis( solver='lsqr', shrinkage=reg) for model='lda': sklearn’s fixed shrinkage applies shrunk_covariance(empirical_covariance(X_k), reg) to every class and sums them weighted by the priors, which by linearity is exactly \(\Sigma\) above;

  • sklearn.linear_model.RidgeClassifier(alpha=reg) for model='ridge' (for two classes sklearn fits a single +1/-1 column, whose sign decision is identical to the arg-max over the two complementary one-vs-rest columns used here).

Absent classes. A class that does not occur in a training fold cannot be predicted for that fold (its score is -inf), as in scikit-learn, whose fitted classes_ only contain the training labels; no nan is produced. Labels outside [0, n_classes) are never predicted and always count as errors.

The function is jittable and vmappable (with explicit n_classes) but not differentiable (the prediction is an arg-max).

References

Examples

>>> import jax
>>> import jax.numpy as jnp
>>> from braintools.analysis import population
>>> k1, k2 = jax.random.split(jax.random.PRNGKey(0))
>>> y = jnp.arange(103) % 3                         # 3 classes
>>> x = jax.random.normal(k1, (103, 8)) + 3.0 * jax.nn.one_hot(y, 8)
>>> acc = population.decode(x, y, n_folds=5)        # 3 trailing samples dropped
>>> acc.shape
(5,)
>>> bool(acc.mean() > 0.9)
True
>>>
>>> # Under jit, pass the number of classes explicitly.
>>> f = jax.jit(lambda x, y, k: population.decode(x, y, n_classes=3,
...                                                 model='ridge', key=k))
>>> bool(f(x, y, k2).mean() > 0.9)
True
Contents