The Taske Aligned Manifold Estimation (TAME-GP), is an extended Poisson-Gaussian CCA with GP-priors for the latent factors.

The code implements a targeted dimensionality reduction of neural spiking data fitting a latent variable probabilistic model.
The observed variables are spike counts from m different simoultaneusly recorded neural populations x₁,..,x_m, and some task relevant variables s.
We introduce a latent variable z₀ capturing correlation between task varaibles and spiking (as in a probabilistic CCA), and other z_j, j=1,..,m latents capturing whithin brain area correlations (as in a factor analysis).

For each factor we include a GP prior with RBF kernels,
,
and we model the likelihood task variables given the factor z₀ as a Gaussian with,
,
while the spike counts of the units of the j-th neural population are assumed independent given the latent variables and Poisson distributed, for the i-th unit we set

.

Given the assumptions, the model factorizes according to,

The resulting graphical model is depicted below,

The code implements parameter learning and inference in the mdoel presented above. Parameters are learned with an approximate EM algorithm where the latent posterior distribution is approximated via the Laplace method. Inference is implemented as numerical optimization of the MAP for the latents. Latent covariacnce is approximated via the Laplace method.

[1] Balzani, Edoardo , et al., "A probabilistic framework for task-aligned intra- and inter-area neural manifold estimation.", Arxiv, 2022