Dimensionality_Reduction_for_Neural_Data.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AQViYQWusLlk"
   },
   "source": [
    "*You may navigate the notebook through the table of contents on the left.*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-ovwIXKO_0OE"
   },
   "source": [
    "---\n",
    "# **Dimensionality reduction for neural data**\n",
    "---\n",
    "\n",
    "Leonardo Agueci <sup>1</sup>, Arthur Pellegrino<sup>2</sup> and N. Alex Cayco Gajic<sup>1</sup>\n",
    "\n",
    "<sup>1</sup>Ecole Normale Superieure Paris, <sup>2</sup>University of Edinburgh"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vicaP9ha1RME"
   },
   "source": [
    "# Introduction\n",
    "\n",
    "<p align=\"center\" width=\"100%\">\n",
    "<img src=\"https://drive.google.com/uc?id=1bWAEsI-pGOSUWmMW3TfSmTVk1B32jSkR\" width=\"550px;\">\n",
    "</p>\n",
    "\n",
    "In this tutorial you will extensively learn about principal component analysis (PCA) [1,2] and its application on neural data. \n",
    "- In __Part I__, we introduce you to the theory underlying this dimensionality reduction method, and its relationship with singular value decomposition (SVD). Then, we will extend its usage to tensors, basic step for applying this method to standard neural data; as a ___bonus___, you will find a brief introduction to tensor component analysis (TCA) and to sliceTCA. Finally, you will find a description of the bi-cross validation method. \n",
    "- In __Part II__ you will use this theoretical background to implement PCA.\n",
    "- __Part III__ focuses on how to cross-validate the number of components.\n",
    "- In __Part IV__ you will apply PCA to neural data recorded during hand movement.\n",
    "- Finally, __Part V__ we will finally try to gain some understanding of what could be encoded in our neural recordings using some simple linear decoders\n",
    "\n",
    "___Remarks:___\n",
    "\n",
    "- Code cells are partially filled sometimes, for helping you with coding and visualization. Although you can change any part of the code as you prefer (e.g. for improving visualization), we recommend to add you code only in the parts delimited as follows (for plots, you might find some settings that might be helpful):\n",
    "\n",
    "```(python)\n",
    "# Fill in your code below\n",
    "###########################\n",
    "\n",
    "# Instruction 1 \n",
    "# your code here\n",
    "\n",
    "# Instruction 2\n",
    "plt.plot(# your code here , some_useful_settings)\n",
    "\n",
    "###########################\n",
    "```\n",
    "\n",
    "- __If you want to run this notebook offline, make sure your python environment has installed the following libraries:__\n",
    "    - matplotlib\n",
    "    - numpy\n",
    "    - pickle\n",
    "    - mpl_toolkits\n",
    "    - sklearn\n",
    "    - tqdm\n",
    "    - scipy\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "id": "Jsq_YHti1TzE"
   },
   "source": [
    "# **Part I.** The theory behind (multi)linear dimensionality reduction\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "\n",
    "One of the reasons for the enduring popularity of PCA is its analytical interpretability. To demonstrate this, we will present two mathematically equivalent formulations of PCA: 1) identifying the directions of maximum projected variance, and 2) finding the best low-dimensional approximation to the data matrix. By setting up these two questions as constrained optimization problems (following [2]), it can be shown that the two formulations yield the same solution: the eigenvalues and eigenvectors of the data covariance matrix. We will then cover the relationship between PCA and singular value decomposition (SVD), and the Eckart-Young Theorem which guarantees that the SVD is the closest low-rank approximation to the data matrix. The equivalence of these two formulations (the closest low rank approximation and vs directions of maximum variance) will set up the intuition necessary to understand data tensor extensions such as TCA or sliceTCA."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true,
    "id": "hucfkKCJFmrG"
   },
   "source": [
    "## The equivalent formulations of PCA"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true,
    "id": "Sj3mrj-dFYGx"
   },
   "source": [
    "### PCA as the directions of maximum projected variance\n",
    "\n",
    "When discussing PCA, we follow convention that we are working with a data matrix $\\mathbf{X} \\in \\mathbb{R}^{N \\times K}$ where $N$ is the number of neurons and $K$ the number of samples. We further assume that the data has already been centered, that is the activity of each neuron has been subtracted its mean. The sample covariance matrix $\\mathbf{\\hat\\Sigma}\\in \\mathbb{R}^{N \\times N}$ is then defined as $\\mathbf{\\hat \\Sigma} = \\frac{1}{N_\\text{samples}}\\mathbf{X}\\mathbf{X}^T$.\n",
    "\n",
    "We start by seeking the vector $\\mathbf{u} \\in \\mathbb{R}^N$ that maximizes the variance of the projected data. We further add the constraint that $\\bf u$ must be normalized ($\\mathbf{u}^T\\mathbf{u}=1$, else the problem is unbounded). Then, the projected data is simply $\\mathbf{u}^T\\mathbf{X}\\in \\mathbb{R}^{K}$. The projected variance is then given by:\n",
    "\n",
    "\\begin{align}\n",
    "\\text{var}(\\mathbf{u}^T\\mathbf{X}) =& \\frac{1}{N_\\text{samples}}(\\mathbf{u}^T \\mathbf{X})(\\mathbf{X}^T \\mathbf{u}) \\\\\n",
    "=&\\mathbf{u}^T \\left ( \\frac{1}{N_\\text{samples}}\\mathbf{X}\\mathbf{X}^T \\right ) \\mathbf{u} \\\\\n",
    "=& \\mathbf{u}^T \\mathbf{\\hat \\Sigma} \\mathbf{u}.\n",
    "\\end{align}\n",
    "\n",
    "One way to solve such an optimization problem is to use [Lagrange multipliers](https://en.wikipedia.org/wiki/Lagrange_multiplier). First, we define the Lagrange function:\n",
    "\n",
    "$$ L(\\mathbf{u},\\lambda) = \\mathbf{u}^T \\mathbf{\\hat \\Sigma u} + \\lambda(1-\\mathbf{u}^T\\mathbf{u})$$\n",
    "\n",
    "A necessary condition for $\\mathbf{u}$ to be a maximum under the imposed constraint is that it must be a zero of the gradient of the Lagrange function:\n",
    "\n",
    "$$ \\nabla_{\\mathbf u} L = \\mathbf{\\hat \\Sigma u} - \\lambda \\mathbf{u} = 0 $$\n",
    "\n",
    "That is, $ \\mathbf{\\hat \\Sigma u} = \\lambda \\mathbf{u}$, therefore  $\\mathbf{u}$ must be an eigenvector of the data covariance matrix and $\\lambda$ its corresponding eigenvalue. Combining this with the projected variance as derived above, we get $\\text{var}(\\mathbf{u}^T\\mathbf{X}) = \\mathbf{u}^T \\mathbf{\\hat \\Sigma} \\mathbf{u} = \\lambda\\mathbf{u}^T  \\mathbf{u} = \\lambda$. Therefore the vector that maximizes the projected variance is the eigenvector of $\\bf \\hat \\Sigma$ with the largest eigenvalue.\n",
    "\n",
    "We can show inductively that by repeating this argument (with an additional constraint that each vector is orthogonal to the previously identified vectors), that the directions of maximum projected variance are given by the eigenvectors of $\\bf \\hat \\Sigma$ and the projected variances by their corresponding eigenvalues.\n",
    "\n",
    "The first $R$ principal components are given by projecting the data $\\bf X$ onto the $R$ eigenvectors of the covariance matrix with largest eigenvalues. The eigenvectors are often called the loadings or weights, and the projected data are often called the scores.\n",
    "\n",
    "Finally, we note that since $\\bf \\hat \\Sigma$ is symmetric, its eigenvectors form an orthonormal basis for $\\mathbb{R}^N$. In the context of PCA, this means that the loading vectors are orthogonal and the scores are uncorrelated."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true,
    "id": "IWFzLPGIM4dR"
   },
   "source": [
    "### PCA as the low-dimensional projection that minimizes reconstruction error\n",
    "\n",
    "In this second formulation we will attempt to find the low-dimensional projection of $\\mathbf{X}$ with the smallest *reconstruction error*. We'll denote each sample (i.e. column) of the data matrix as ${\\bf x}_n\\in\\mathbb{R}^N$. We thus seek an $R$-dimensional subspace such that, on average over all samples, the data projected on this subspace $\\mathbf{\\hat x}_n\\in\\mathbb{R}^N$ is closest to ${\\bf x}_n$. We'll call our candidate orthonormal basis $\\mathbf{u}_1,\\dots,\\mathbf{u}_R\\in\\mathbb{R}^N$, which spans this $R$-dimensional subspace within the $N$-dimensional activity space. Then, the estimate of each sample can be written as the linear combination:\n",
    "\n",
    "$${\\bf \\hat x}_n =\\sum_{i=1}^R \\beta_{ni} \\mathbf{u}_i$$\n",
    "\n",
    "for some coefficients $\\beta_{ni}$, which we would like to find as a function of $\\mathbf{u}_i$ and $\\mathbf{x}_n$.\n",
    "\n",
    "Note that we're still missing $N-R$ vectors needed to form a complete orthonormal basis for the full activity space: we'll simply call such vectors $\\mathbf{u}_{R+1},\\dots,\\mathbf{u}_N$. This means that we can rewrite each sample as:\n",
    "\n",
    "$${\\bf x}_n =\\sum_{i=1}^N \\alpha_{ni} \\mathbf{u}_i$$\n",
    "\n",
    "Since the $\\mathbf{u}_i$'s are orthonormal, the coefficients are simply given by the projections of the data onto each basis vector: $\\alpha_{ni} =\\mathbf{x}_n^T \\mathbf{u}_{i} $. You can verify this by taking the dot product of the above equation with $\\mathbf{u}_j$ (remembering that $\\mathbf{u}_i$'s being orthonormal is defined as: $\\mathbf{u}_i\\cdot \\mathbf{u}_j=1$ if $i=j$ else $0$).\n",
    "\n",
    "We can write the mean-squared error as:\n",
    "$$ E =  \\frac{1}{N \\cdot N_\\text{samples}} \\sum_{n=1}^{N_\\text{samples}} \\lVert \\mathbf{x}_n - \\mathbf{\\hat x}_n \\rVert^2 $$\n",
    "First, using our expressions for the samples above, we can write the residual on the right as:\n",
    "\n",
    "\\begin{align}\n",
    " \\mathbf{x}_n - \\mathbf{\\hat x}_n = \\sum_{i=1}^R (\\alpha_{ni}-\\beta_{ni})\\mathbf{u}_i + \\sum_{j=R+1}^N \\alpha_{nj}\\mathbf{u}_j\n",
    "\\end{align}\n",
    "\n",
    "Now, noting that $\\lVert \\mathbf{x}_n - \\mathbf{\\hat x}_n \\rVert^2 = ( \\mathbf{x}_n - \\mathbf{\\hat x}_n )\\cdot  (\\mathbf{x}_n - \\mathbf{\\hat x}_n)$, and using the fact that the $\\mathbf{u}_i$'s are orthonormal, the norm of the residual can be simplified as:\n",
    "\n",
    "\\begin{align}\n",
    "\\lVert \\mathbf{x}_n - \\mathbf{\\hat x}_n \\rVert^2 = \\sum_{i=1}^R (\\alpha_{ni}-\\beta_{ni})^2 + \\sum_{j=R+1}^N \\alpha_{nj}^2\n",
    "\\end{align}\n",
    "\n",
    "Finally, we can minimize the error by differentiating it as:\n",
    "$$\\frac{d E}{d \\beta_{ni}} =  \\frac{2}{N \\cdot N_\\text{samples}} (\\beta_{ni} - \\alpha_{ni})$$\n",
    "\n",
    "which means that the $\\beta_{ni}$ that minimizes the MSE is given simply by $\\beta_{ni} = \\alpha_{ni} = \\mathbf{x}_n^T \\mathbf{u}_i$. In retrospect this makes sense since the error should be minimized by projecting the data onto the low-dimensional subspace. We can now expand the error as:\n",
    "\n",
    "\\begin{align}\n",
    "E =& \\frac{1}{N \\cdot N_\\text{samples}} \\sum_{n=1}^{N_\\text{samples}} \\sum_{j=R+1}^N \\alpha_{nj}^2 \\\\\n",
    "=& \\frac{1}{N \\cdot N_\\text{samples}} \\sum_{n=1}^{N_\\text{samples}} \\sum_{j=R+1}^N (\\mathbf{x}_n^T \\mathbf{u}_j)^2 \\\\\n",
    "=& \\frac{1}{N \\cdot N_\\text{samples}} \\sum_{n=1}^{N_\\text{samples}} \\sum_{j=R+1}^N (\\mathbf{u}_j^T\\mathbf{x}_n )(\\mathbf{x}_n^T \\mathbf{u}_j) \\\\\n",
    "=& \\frac{1}{N} \\sum_{j=R+1}^N \\mathbf{u}_j^T \\mathbf{\\hat \\Sigma u}_j\n",
    "\\end{align}\n",
    "Where last step follows from $\\mathbf{x}_n$ being mean-centered. It now remains to find an orthonormal basis $\\mathbf{u}_j$ that minimizes the mean squared error. In the case that $R=N-1$, we have the same Lagrangian as before (up to a scaling factor):\n",
    "$$ L(\\mathbf{u},\\lambda) = \\frac{1}{N} \\mathbf{u}^T \\mathbf{\\hat \\Sigma u} + \\lambda(1-\\mathbf{u}^T\\mathbf{u})$$\n",
    "Again the solutions are the eigenvectors of the covariance matrix. This time, since we want to *minimize* the reconstruction error, we choose the eigenvector with smallest eigenvalue for the residual, which means that the low-dimensional subspace we project the data on is spanned by the remaining eigenvectors with larger eigenvalues.\n",
    "\n",
    "*In general,* it can be shown that taking the $R$ largest eivenvectors of the covariance matrix as the principal subspace solves the constrained minimization of the MSE, resulting in the same solution as in the previous section. *Therefore,* PCA can be formulated in two mathematically equivalent forms: maximizing the directions of projected variance, or minimizing the reconstruction error of a low-dimensional projection.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true,
    "id": "80f8FubC76HK"
   },
   "source": [
    "### PCA as a low-rank approximation to the data matrix\n",
    "\n",
    "So far we saw PCA as solving two different constrained optimization problems by projecting of the data onto the eigenvectors of the covariance matrix. Further, since the covariance matrix has an orthonormal set of eigenvectors, this projection can be considered as a change of basis from Cartesian coordinates (i.e., each dimension represents a different neuron) to a more \"natural\" basis of the data (i.e., each dimension represents a different covariability pattern across the population). Mathematically this projection is given by the equation:\n",
    "$$ \\mathbf{Z} = \\mathbf{U}^T \\mathbf{X}$$\n",
    "where $\\bf U \\in \\mathbb{R}^{N\\times N}$ is a matrix whose columns are given by the orthonormal eigenvectors of $\\mathbf{\\hat \\Sigma}$, in order of descending eigenvalue magnitude (which is the variance of the projected data, therefore this ranks the components by how much variance of the data they capture).\n",
    "\n",
    "Note that since its columns $\\mathbf{u}_i$ are orthonormal, $\\bf U$ is a unitary matrix. An important property of unitary matrices is that $\\mathbf{U}^{-1} = \\mathbf{U}^T$. To see this, note that the $\\mathbf{U}^T \\mathbf{U} = \\mathbf{I}$ falls out naturally from matrix multiplication, as $(\\mathbf{U}^T \\mathbf{U})_{ij}=\\mathbf{u}_i^T\\mathbf{u}_j = \\delta_{ij}$ by orthonormality. To prove that $ \\mathbf{U} \\mathbf{U}^T = \\mathbf{I}$, we can multiply both sides of the equation above by $\\mathbf{U}$ on the left and $\\mathbf{U}^{-1}$ on the right as follows:\n",
    "\\begin{align}\n",
    "\\mathbf{U}^T \\mathbf{U} &= \\mathbf{I} \\\\\n",
    "\\mathbf{U}(\\mathbf{U}^T \\mathbf{U})\\mathbf{U}^{-1} &= \\mathbf{U}\\mathbf{U}^{-1}\\\\\n",
    "\\mathbf{U}\\mathbf{U}^T &= \\mathbf{I} \\\\\n",
    "\\end{align}\n",
    "\n",
    "This is important in the context of PCA because it means we can easily move $\\mathbf{U}$ to the other side of the equation to project the scores back into the original basis, returning the original data matrix:\n",
    "$$\\mathbf{U} \\mathbf{Z} = \\mathbf{X}$$\n",
    "and now we see that PCA is a *matrix decomposition*, meaning that it decomposes the data matrix into a product of two matrices: the new basis vectors (or loadings) $\\bf U$ and the projected data (or scores) $\\mathbf Z$.\n",
    "\n",
    "In dimensionality reduction, the purpose is to reduce the number of features used to capture the data. In PCA, this is done by taking the first $R$ components, which correspond to the largest eigenvalues (and thus the directions of largest projected variance). To do this we simply truncate the projection matrix after the first $R$ columns to get $\\mathbf{ U }_{R}\\in \\mathbb{R}^{N\\times R}$. Then the $R$-dimensional scores are still given by the same formula:\n",
    "$$ \\mathbf{Z}_R = \\mathbf{U}_{R}^T \\mathbf{X}$$\n",
    "\n",
    "However we cannot simply move $\\mathbf{U}_{R}$ to the other side of the equation. While we still have $\\mathbf{U}_{R}^T  \\mathbf{U}_{R} = \\mathbf{I}$ (due to orthonormality of the columns), we can no longer guarantee that $ \\mathbf{U}_{R} \\mathbf{U}_{R}^T  = \\mathbf{I}$ since our proof above relied on the existence of $\\mathbf{U}_{R}^{-1}$, which does not exist for non-square matrices ($R< N$)! Geometrically, this makes sense since we can project from a high-dimensional space to a low-dimensional space but cannot invert the projection to recover the exact original data.\n",
    "\n",
    "We can try to revert the transformation anyway by projecting the top $R$ components back into the $N$-dimensional space using $\\mathbf{U}_R$. This results in the following *approximation* to the data matrix:\n",
    "$$ \\mathbf{ \\hat X} = \\mathbf{U}_{R} \\mathbf{Z}_R$$\n",
    "Another way to write this approxmination is to expand it as a sum of the outer products of the column.\n",
    "$$ \\mathbf{ \\hat X} = \\sum_{r=1}^R \\mathbf{u}_{r} \\mathbf{ z}^T_r$$\n",
    "where $ \\mathbf{u}_{r}$ are the columns of $ \\mathbf{U}$ and $ \\mathbf{ z}^T_r$ are the rows of $\\mathbf{Z}$. Each column of $\\mathbf{\\hat X}$ is a linear combination of the $R$ basis vectors, so the rank of $ \\mathbf{ \\hat X} $ is (at most) $R$. Equivalently, the approximated data will lie on an $R$-dimensional subspace of $\\mathbb{R}^N$ spanned by the new basis $\\mathbf{u}_1,\\dots,\\mathbf{u}_R$.\n",
    "\n",
    "This brings us back to the low-dimensional projection view that we saw previously, which further suggests that it may be the best low-rank approximation of the data matrix. This is indeed the case, as we will see in the next section.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true,
    "id": "CvhZAkSUHDzY"
   },
   "source": [
    "### Relationship between PCA and SVD\n",
    "\n",
    "PCA is very closely related to a more general matrix decomposition called the singular value decomposition (SVD). It is an amazing fact that any rank-$R$ matrix can be decomposed into three matrices:\n",
    "$$ {\\mathbf X} = {\\mathbf U S} \\mathbf{V}^T$$\n",
    "where ${\\mathbf U} \\in \\mathbb{R}^{N \\times R}$ and $\\mathbf{V} \\in \\mathbb{R}^{K \\times R}$ have orthonormal columns (called the left and right singular vectors, respectively) that define new rotated bases for their respective vector spaces, and $\\mathbf{S}\\in \\mathbb{R}^{R \\times R}$ is a positive diagonal matrix (whose diagonal elements $\\sigma_i$ are called the singular values, and are ranked in decreasing order). Most often, you will equivalently find it defined with, $U\\in\\mathbb{R}^{N\\times \\min\\{N,K\\}}$, $S\\in\\mathbb{R}^{\\min \\{N,K\\}\\times \\min\\{N,K\\}}$, $V\\in\\mathbb{R}^{\\min\\{N,K\\}\\times N}$, in which case the rank of $\\bf X$ is simply the number of non-zero singular values. The SVD can be thought of as decomposing a matrix into how it operates on vectors: by taking a $K$-dimensional vector, rotating it into a new basis, shrinking or stretching the new coordinates according to $\\sigma_i$, and projecting it into $\\mathbb{R}^N$.\n",
    "\n",
    "How does this relate to PCA? To see this, let's again assume that $\\mathbf{X}$ is mean-centered and calculate the covariance matrix using the SVD:\n",
    "\\begin{align}\n",
    "\\mathbf{\\hat \\Sigma} =& \\frac{1}{N_\\text{samples}}\\mathbf{X}\\mathbf{X}^T \\\\\n",
    "=& \\frac{1}{N_\\text{samples}}\\mathbf{USV}^T\\mathbf{VSU}^T \\\\\n",
    "=& \\frac{1}{N_\\text{samples}}\\mathbf{US}^2\\mathbf{U}^T \\\\\n",
    "\\end{align}\n",
    "which is an eigenvalue decomposition of the covariance matrix (up to the scaling). Therefore the transformed basis vectors ${\\mathbf U}$ in SVD and PCA are identical, and the eigenvalues and singular values are related as $\\lambda_i = \\sigma_i^2/N_\\text{samples}$. Consistent with the previous definitions, the scores in PCA are given by $\\mathbf{Z=U}^T\\mathbf{X=U}^T\\mathbf{USV}^T=\\mathbf{SV}^T$. Therefore, in the language of SVD, dimensionality reduction with PCA is equivalent to truncating the decomposition after $R$ singular values. This means taking only the first $R$ columns of ${\\mathbf U}$ and ${\\mathbf V}$ as well as the upper $R\\times R$ corner of $\\mathbf S$ to get the following rank-$R$ approximation:\n",
    "$$ \\mathbf{\\hat X} = \\mathbf {U}_R \\mathbf{S}_{R} \\mathbf{V}_R^T$$\n",
    "\n",
    "Further, the [Eckart-Young Theorem](https://rich-d-wilkinson.github.io/MATH3030/3-5-lowrank.html) states that the truncated SVD approximation above gives the rank-$R$ that is closest to the data matrix in the Frobenius norm:\n",
    "$$\\min_{\\text{rank}(\\mathbf{Y}) \\leq R} \\lVert  \\mathbf{X - Y} \\rVert_F$$\n",
    "with $\\lVert  \\mathbf{X - \\hat X} \\rVert_F = \\left(\\sum_{i=R+1}^N \\sigma_i^2\\right)^\\frac{1}{2}$. Equivalently, this means that taking the top $R$ principal components PCA results in the best low-rank approximation to the data with the following MSE:\n",
    "$$E = \\frac{1}{N \\cdot N_\\text{samples}}\\sum_{i=R+1}^{N}\\sigma_i^2 $$\n",
    "\n",
    "Together with the maths demonstrated above, this shows that there is a deep connection between minimizing the error of a low rank (or low-dimensional) approximation and capturing covariability patterns in neural data. In the following section we will discuss this idea in extension to neural data tensors."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true,
    "id": "_Eb5QTQ-3Esa"
   },
   "source": [
    "## From matrices to tensors\n",
    "<a id='tensors'></a>\n",
    "So far, we have focused on reducing the dimensionality of a matrix $\\mathbf{X}\\in \\mathbb{R}^{N \\times K}$. However, in neuroscience we often work with a data *tensor* $\\mathcal{X}\\in \\mathbb{R}^{N \\times T \\times K}$, where $N$ is the number of neurons, $K$ the number of *trials* and $T$ the number of time points.\n",
    "\n",
    "The first key element to notice, is that, unlike for matrices, we cannot naturally separate the variables at hand into *features* (neurons) and *samples* (time points). Instead there are two relevant notions of sampling (across time points and across trials) which we would like to keep separated to interpret the data [4]. In the following sections, we will look at three complementary ways to account for this, from which stem three dimensionality reduction methods : PCA on unfoldings [3], TCA [4] and sliceTCA [5]."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true,
    "id": "RSR7L-Rx3M4k"
   },
   "source": [
    "### PCA on unfoldings\n",
    "\n",
    "#### Neuron-unfolding\n",
    "\n",
    "The first strategy, is to go back to the matrix case by *unfolding* the data tensor. For example, if we wish to keep neurons as features, we can unfold the tensor $\\mathcal{X}\\in \\mathbb{R}^{N \\times T \\times K}$ into a matrix $\\mathbf{X}\\in \\mathbb{R}^{N\\times (TK)}$. If we now apply the low-rank approximation perspective on PCA to this matrix, the data can be approximated as:\n",
    "\n",
    "$$\n",
    "  \\mathbf{X}\\approx \\mathbf{\\hat X} = \\mathbf{U}\\mathbf{Z} \\in \\mathbb{R}^{N\\times (TK)}\n",
    "$$\n",
    "\n",
    "where $\\mathbf{U}\\in\\mathbb{R}^{N\\times R}$ and $\\mathbf{Z}\\in \\mathbb{R}^{R \\times(TK)}$. As seen previously, this can be equivalently written as the sum of the following outer products:\n",
    "\n",
    "$$\n",
    "\\mathbf{\\hat X} = \\sum_{r=1}^R \\mathbf{u}_r\\mathbf{z}^T_r\n",
    "$$\n",
    "\n",
    "where $\\mathbf{u}_r \\in \\mathbb{R}^N$ are the columns of $\\bf U$ and $\\mathbf{z}_r\\in\\mathbb{R}^{TK}$ are the rows of $\\bf V$.\n",
    "\n",
    "\n",
    "Next, we can *fold* each $\\mathbf{z}_r$ back into a matrix $\\mathbf{W}_r\\in\\mathbb{R}^{T \\times K}$. The outer product $\\mathbf{u}_r\\mathbf{z}^T_r$ is sometimes written in the form $\\mathbf{u}_r\\otimes \\mathbf{z}_r$, which allows us to generalize the notation for the outer product of two vectors $(\\mathbf{u}_r\\otimes \\mathbf{z}_r)_{ij}=\\mathbf{u}_{r,i}\\mathbf{z}_{r,j}$  to the outer product of a vector with a matrix $(\\mathbf{u}_r\\otimes \\mathbf{W}_r)_{ijk}=\\mathbf{u}_{r,i}\\mathbf{W}_{r,jk}$. The figure below shows how a single *component* can be seen as either the outer product of a vector with a matrix or with a (long) vector.\n",
    "\n",
    "<p align=\"center\" width=\"100%\">\n",
    "<img src=\"https://drive.google.com/uc?id=1oxa02j3jDSrY2izD9M_bpoHtWa6Ja6I1\" width=\"500px;\">\n",
    "</p>\n",
    "\n",
    "This way, we can rewrite our (SVD/PCA) approximation directly as a tensor:\n",
    "\n",
    "$$\n",
    "\\mathcal{X} \\approx \\mathcal{\\hat X} = \\sum_{r=1}^R \\mathbf{u}_r \\otimes \\mathbf{W}_r \\in \\mathbb{R}^{N \\times T \\times K}\n",
    "$$\n",
    "\n",
    "And, since the Frobenius norm is defined element-wise, the same objective as for the matrix case can be used to determine the decomposition:\n",
    "\n",
    "$$\n",
    "E(\\mathbf{u}_1,..., \\mathbf{u}_r, \\mathbf{W}_1, ..., \\mathbf{W}_r) = \\lVert \\mathcal{X} - \\mathcal{\\hat X} \\rVert_F\n",
    "$$\n",
    "\n",
    "#### Trial-unfolding\n",
    "\n",
    "We arbitrarily decided that the *neuron* leg would be the features (the $\\mathbf{u}$'s). When working with matrices, this doesn't matter; in the end, the term $\\mathbf{u}_r \\mathbf{z}_r^T$ is symmetric, in the sense that it can be transposed (i.e. swap $\\mathbf{u}_r$ and $\\mathbf{z}_r$) to now consider the time points as features. When moving to tensors, this doesn't work anymore, as $\\mathbf{z}_r$ now represents two variables (e.g. time and trial).\n",
    "\n",
    "In the context of tensors, if we consider *time* as our features, we just have to unfold the tensor along another of its *legs*:\n",
    "\n",
    "<p align=\"center\" width=\"100%\">\n",
    "<img src=\"https://drive.google.com/uc?id=1dHfvXK4ay3gZHx2_Om35MiRgWrOcNIDY\" width=\"500px;\">\n",
    "</p>\n",
    "\n",
    "And similarly for trials [3]:\n",
    "\n",
    "<p align=\"center\" width=\"100%\">\n",
    "<img src=\"https://drive.google.com/uc?id=1AOPTNowx6LejxaGaMJGjqD58_KEBJcvE\" width=\"500px;\">\n",
    "</p>\n",
    "\n",
    "To keep track of the different vectors and matrices at hand, we can indicate the leg(s) of the tensor they describe as a supsercript. For trial-unfolding PCA, we would therefore write:\n",
    "\n",
    "$$\n",
    "\\mathcal{X} \\approx \\mathcal{\\hat X} = \\sum_{r=1}^R \\mathbf{u}_r^{trial} \\otimes \\mathbf{W}^{neuron, time}_r\n",
    "$$\n",
    "\n",
    "Note that, formally, $\\mathbf{u}_r^{trial} \\otimes \\mathbf{W}^{neuron, time}_r \\in \\mathbb{R}^{K \\times N \\times T}$, but we have implicitely permuted the indices to match the data tensor in $\\mathbb{R}^{N \\times T \\times K}$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true,
    "id": "R7Ms_CNW3TyG"
   },
   "source": [
    "### (bonus) SliceTCA\n",
    "\n",
    "We have just seen that applying PCA on unfoldings of a data tensor was equivalent to writing the approximation of the data tensor as a sum of outer product of vectors (the feature) and matrices (the folded scores vectors). We can generalize this concept by allowing different types of unfoldings within the same sum. Preserving the same notation as before,\n",
    "\n",
    "$$\n",
    "\\mathcal{X} \\approx \\mathcal{\\hat X} = \\sum_{r=1}^{R_{neuron}} \\mathbf{u}_r^{neuron} \\otimes \\mathbf{W}^{time, trial}_r  + \\sum_{r=1}^{R_{time}} \\mathbf{u}_r^{time} \\otimes \\mathbf{W}^{neuron, trial}_r  + \\sum_{r=1}^{R_{trial}} \\mathbf{u}_r^{trial} \\otimes \\mathbf{W}^{neuron, time}_r\n",
    "$$\n",
    "\n",
    "This decomposition is called sliceTCA [5].\n",
    "\n",
    "<p align=\"center\" width=\"100%\">\n",
    "<img src=\"https://drive.google.com/uc?id=1H5RkoRhITPLbAGHiLMKaK_WsQqk-WzfM\" width=\"600px;\">\n",
    "</p>\n",
    "\n",
    "Note that by setting (e.g.) $R_{time}=R_{trial} = 0$, we fall back onto (e.g.) the neuron-unfolding PCA decomposition. So that sliceTCA is a generalization (or in a set-theoretic sense, a superclass) of PCA on unfoldings."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true,
    "id": "r0dkEPTK3V2W"
   },
   "source": [
    "### (bonus) TCA\n",
    "\n",
    "Another strategy for going from PCA on unfoldings to a tensor decomposition is to constrain the slices to be of rank $1$. Consider again a PCA on (any) unfolding decomposition:\n",
    "\n",
    "$$\n",
    "\\mathcal{X} \\approx \\mathcal{\\hat X} = \\sum_{r=1}^R \\mathbf{u}_r \\otimes \\mathbf{W}_r \\in \\mathbb{R}^{N \\times T \\times K}\n",
    "$$\n",
    "\n",
    "If we now constrain the matrix to be of rank $1$, that is $\\mathbf{W}_r = \\mathbf{v}_r \\mathbf{w}_r^T = \\mathbf{v}_r \\otimes \\mathbf{w}_r$:\n",
    "\n",
    "<p align=\"center\" width=\"100%\">\n",
    "<img src=\"https://drive.google.com/uc?id=1STP28EUzPHYZNEkh0uC0PiUPmf4W1fx5\" width=\"500px;\">\n",
    "</p>\n",
    "\n",
    "And we thus obtain the decomposition:\n",
    "\n",
    "$$\n",
    "\\mathcal{X} \\approx \\mathcal{\\hat X} = \\sum_{r=1}^R \\mathbf{u}_r \\otimes \\mathbf{v}_r \\otimes \\mathbf{w}_r \\in \\mathbb{R}^{N \\times T \\times K}\n",
    "$$\n",
    "\n",
    "This decomposition is known within neuroscience as tensor component analysis (TCA, [4]) but in other fields (such as physics) is known by other names: CANDECOMP, PARAFAC, or simply the CP decomposition. Therefore, PCA is a generalization (or superclass) of TCA."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true,
    "id": "By4NmZOwYJDH"
   },
   "source": [
    "### Relationship between tensor decomposition methods\n",
    "\n",
    "Overall, the relationsip between the three tensor decomposition methods presented here can be summarized by the following Venn diagram [5]:\n",
    "\n",
    "<p align=\"center\" width=\"100%\">\n",
    "<img src=\"https://drive.google.com/uc?id=1mNdFEP7gnDK72Q-jDyT1jkKs1sG4q698\" width=\"300px;\">\n",
    "</p>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "id": "fv-sEwUwBnb9"
   },
   "source": [
    "# **Part II.** Vizualizing and denoising data with PCA on a toy data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "\n",
    "To familiarize yourselves with PCA/SVD, you will first write your own implementation of PCA (without using `sklearn`). Second, you will try to *guess* the rank of some simulated data using the bi-cross-validation method.\n",
    "\n",
    "**WARNING:** In the maths above we used the convention that $\\mathbf{X} \\in \\mathbb{R}^{N\\times K} $ (i.e., the columns are the samples), as this is commonly used in neuroscience and in the PCA literature. However, other communities (notably in machine learning) use the rows as samples instead: $\\mathbf{X} \\in \\mathbb{R}^{K\\times N}$. Be mindful when you write your code that the packages you use may be using the latter notation (for example, in that case, the $\\bf V$ in the SVD will represent the basis vectors, rather than the $\\bf U$.).\n",
    "\n",
    "First, execute the cell below to get started and generate your data matrix $M$ of 10 neurons and 101 time points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-25T16:34:43.597797Z",
     "start_time": "2024-03-25T16:34:43.577338Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "hidden": true,
    "id": "YJYsEJRa8DqI",
    "outputId": "061d070a-7704-4451-cec0-1f7d02286db2"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "np.random.seed(3)\n",
    "\n",
    "def exp_quad(x, l=0.1, s=1): \n",
    "    return np.exp(-x**2/(2*l**2))*s**2\n",
    "\n",
    "n = 10\n",
    "k = np.random.randint(1, n//2+1)\n",
    "time_steps = 101\n",
    "ts = np.linspace(0, 1, time_steps)\n",
    "noise = 0.01\n",
    "kernel = exp_quad\n",
    "\n",
    "latent_gaussian = np.random.randn(time_steps, k)\n",
    "covariance_matrix = kernel(np.subtract.outer(ts, ts))\n",
    "L = np.linalg.cholesky(covariance_matrix+10**-6*np.eye(time_steps))\n",
    "\n",
    "neuron_basis = np.random.randn(n, k)\n",
    "time_basis = L @ latent_gaussian\n",
    "\n",
    "M = time_basis @ neuron_basis.T + np.random.randn(n) + np.random.randn(time_steps, n)/3\n",
    "\n",
    "print('Data matrix M (sample_size, neuron):', M.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true,
    "id": "Ifljaff4D1fU"
   },
   "source": [
    "**Step 1.** The first step of neural dimensionality reduction is *always* to look at the raw data first. Plot the single neuron activity generated by the previous code cell over time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true,
    "id": "hAmcHlYdk4ph"
   },
   "outputs": [],
   "source": [
    "# Fill in your code below\n",
    "fig = plt.figure(figsize=(6,3), constrained_layout=True)\n",
    "ax = fig.add_subplot()\n",
    "\n",
    "ax.plot( # your code here)\n",
    "ax.set_xlabel('Time')\n",
    "ax.set_ylabel('Neural activity')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true,
    "id": "65kcjvPYWTha"
   },
   "source": [
    "**Step 2.** Center the data matrix (subtract from each neuron its mean activity). Make a scatter plot of two arbitrarily chosen neurons, using each as one of the two dimensions. What do you notice?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true,
    "id": "99eiybV7mKkr"
   },
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(4,4), constrained_layout=True)\n",
    "ax = fig.add_subplot()\n",
    "\n",
    "# Fill in your code below\n",
    "###########################\n",
    "\n",
    "# Center each neuron activity in M\n",
    "M_centered = # your code here\n",
    "\n",
    "# Plot the centered activity for neurons i and j of your choise\n",
    "ax.scatter()\n",
    "\n",
    "###########################\n",
    "\n",
    "ax.set_aspect('equal')\n",
    "ax.spines[['left', 'bottom']].set_position((\"data\", 0))\n",
    "ax.spines[['top', 'right']].set_visible(False)\n",
    "ax.set_xlabel('Neuron '+str(i), loc='right')\n",
    "ax.set_ylabel('Neuron '+str(j), loc='top')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true,
    "id": "zbjO7fENTit3"
   },
   "source": [
    "**Step 3.** Perform PCA using ``` np.linalg.svd ```.  Be careful with rows and columns, and verify whether the basis vectors should be $\\bf U$ or $\\bf V$. Make the following two plots:\n",
    "* a scatter plot of projections of the data onto the first two PCs using the $S$, $V$, and $U$.\n",
    "* a plot showing how the projected variance decays as a function of the number of components."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-25T16:50:25.414714Z",
     "start_time": "2024-03-25T16:50:25.406068Z"
    },
    "hidden": true,
    "id": "O3Z4ZUCxmvxV"
   },
   "outputs": [],
   "source": [
    "# Fill in your code below\n",
    "###############################\n",
    "\n",
    "# Generate U, S, V using SVD\n",
    "U, S, V = # your code here\n",
    "\n",
    "# Generate the projected data\n",
    "projection = # your code here\n",
    "\n",
    "# Compute the variance along each dimension (careful with matrix axes!)\n",
    "projected_variance = # your code here\n",
    "\n",
    "###############################\n",
    "\n",
    "fig = plt.figure(figsize=(8,3), constrained_layout=True)\n",
    "ax = fig.add_subplot(1,2,1)\n",
    "\n",
    "ax.scatter(projection[:,0], projection[:,1], alpha=0.5)\n",
    "ax.set_xlabel('PC1')\n",
    "ax.set_ylabel('PC2')\n",
    "\n",
    "ax = fig.add_subplot(1,2,2)\n",
    "ax.plot(projected_variance,'o-')\n",
    "ax.set_xlabel('Number of components')\n",
    "ax.set_ylabel('Projected variance')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true,
    "id": "05UiKsNblGkW"
   },
   "source": [
    "**Step 4.** Now make the same plot as above using the eigenvectors of the covariance matrix.\n",
    "* Calculate the sample covariance matrix.\n",
    "* Use ``` np.linalg.eig ```  to get the eigenvectors and eigenvalues of the covariance matrix. Be careful with the rows and columns.\n",
    "* Make a scatter plot of the data matrix on the first two eigenvectors of the covariance matrix, just as you did in Step 3. Are the two plots identical? Should they be?\n",
    "* Finally plot the projected variance using the eigenvalues. Are the two plots identical?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true,
    "id": "Thc-337HnBy-"
   },
   "outputs": [],
   "source": [
    "# Fill in your code below\n",
    "###############################\n",
    "\n",
    "# Generate the covariance matrix\n",
    "Sigma = # your code here\n",
    "\n",
    "# Compute the eigenvalues using np.linalg.eig()\n",
    "evals, evecs = # your code here\n",
    "\n",
    "# Generate the projected data\n",
    "projection = # your code here\n",
    "\n",
    "# Compute the variance using eigenvalues\n",
    "projected_variance = # your code here\n",
    "\n",
    "###############################\n",
    "\n",
    "# Let's plot what we got\n",
    "fig = plt.figure(figsize=(8,3), constrained_layout=True)\n",
    "ax = fig.add_subplot(1,2,1)\n",
    "\n",
    "ax.scatter(projection[:,0], projection[:,1], alpha=0.5)\n",
    "ax.set_xlabel('PC1')\n",
    "ax.set_ylabel('PC2')\n",
    "\n",
    "ax = fig.add_subplot(1,2,2)\n",
    "projected_variance = np.sort(projected_variance)[::-1]\n",
    "ax.plot(projected_variance,'o-')\n",
    "ax.set_xlabel('Number of components')\n",
    "ax.set_ylabel('Projected variance')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true,
    "id": "Pj0cCNXnmLyQ"
   },
   "source": [
    "**Step 5.** PCA is often used as a preprocessing step for other dimensionality reduction methods due to its denoising properties. Let's try to recompose the neural firing rates using the first two principal components. Does it look denoised to you? Try to change the number of components to see the difference.<br>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "n_components = 2\n",
    "\n",
    "# Fill in your code below\n",
    "##############################\n",
    "\n",
    "# let's reconstruct the data by using only few components (always remember to put back the mean)\n",
    "M_reconstructed = # your code here\n",
    "\n",
    "##############################\n",
    "\n",
    "fig = plt.figure(figsize=(10,3), constrained_layout=True)\n",
    "ax = fig.add_subplot(1,2,1)\n",
    "\n",
    "ax.set_title('Initial data')\n",
    "ax.plot(ts, M)\n",
    "ax.set_xlabel('Time')\n",
    "ax.set_ylabel('Neural activity')\n",
    "\n",
    "ax = fig.add_subplot(1,2,2)\n",
    "\n",
    "ax.set_title('Reconstructed data')\n",
    "ax.plot(ts, M_reconstructed)\n",
    "ax.set_xlabel('Time')\n",
    "ax.set_ylabel('Neural activity')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true,
    "id": "xiEPnSYgqSTt"
   },
   "source": [
    "**Step 6.** A frequently used method to calculate the dimensionality is to plot the cumulative fraction of variance explained in the data and see when the curve levels off. Since the projected variance is given by the eigenvalues, this can be evaluated by a cumulative sum of eigenvalues normalized by the sum of all eigenvalues. Plot this as a function of the number of components. What do you think the dimensionality is?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true,
    "id": "ko9OTnGQqSpW"
   },
   "outputs": [],
   "source": [
    "# Fill in your code below\n",
    "##############################\n",
    "\n",
    "# Compute the comulative (remember to normalize)\n",
    "cumulative = # your code here\n",
    "\n",
    "##############################\n",
    "\n",
    "fig = plt.figure(figsize=(4,3), constrained_layout=True)\n",
    "ax = fig.add_subplot()\n",
    "ax.plot(cumulative, 'o-')\n",
    "ax.set_xlabel('Number of components')\n",
    "ax.set_ylabel('Variance explained')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rmCBrK9ZZG2U"
   },
   "source": [
    "# **Part III.**  Cross-validating dataset dimensionality"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "A key step in dimensionality reduction methods is to determine the dimensionality of the data, that is to decide on the number of components to use. There exist two methods for doing it: (1) bi-cross-validation [6], which can be used for PCA, NMF, and similar matrix factorizations, and (2) masking entries of the data matrix or tensor during optimization, which is more general. Here, we will focus on ___bi-cross-validation___.\n",
    "\n",
    "You can read more about cross-validation methods for PCA on this nice [blog post](https://alexhwilliams.info/itsneuronalblog/2018/02/26/crossval/)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-25T17:30:07.324843Z",
     "start_time": "2024-03-25T17:30:07.319937Z"
    },
    "heading_collapsed": true
   },
   "source": [
    "## The method"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true,
    "id": "hxVUyYJ23l00"
   },
   "source": [
    "### Bi-cross-validation for PCA\n",
    "\n",
    "The SVD formulation of PCA allows for a very nice way to cross-validate the number of components [6]. The basic idea is that we will split the data into *two* test/training sets: one in neurons, and one in time. That is, we split the data matrix into block structure:\n",
    "$\\mathbf{X} = \\begin{bmatrix} \\mathbf{X}^{tr}_{tr} & \\mathbf{X}_{tr}^{te}\\\\ \\mathbf{X}_{te}^{tr} & \\mathbf{X}^{te}_{te} \\end{bmatrix}$.\n",
    "The subscripts corresponds to neuron indices, and the superscripts to sample indices.\n",
    "\n",
    "For each $R$, we use the truncated SVD on the data from *all* neurons during the *training* samples to identify the $R$ basis vectors:\n",
    "$$\\mathbf{X}^{tr} = \\mathbf{U}_R\\mathbf{Z}_{R}^{tr}$$\n",
    "or, in block form:\n",
    "$$\\begin{bmatrix} \\mathbf{X}^{tr}_{tr} \\\\ \\mathbf{X}_{te}^{tr} \\end{bmatrix} = \\begin{bmatrix} \\mathbf{U}_{R,tr} \\\\ \\mathbf{U}_{R,te} \\end{bmatrix} \\mathbf{Z}_{R}^{tr}$$\n",
    "\n",
    "We now want to cross-validate how well $\\mathbf{U}_R$ describes the data in the test samples. This is in general ill-defined since we are moving from a low-dimensional subspace to a high-dimensional space, but the key idea is to use the *training* neurons in the *test* samples, in addition to the low-dimensional basis we just identified from the *training* samples, to predict the *test* neuron's activity in the *test* samples. Starting from the equivalent block formula for the test samples, we have:\n",
    "$$\\begin{bmatrix} \\mathbf{X}^{te}_{tr} \\\\ \\mathbf{X}_{te}^{te} \\end{bmatrix} = \\begin{bmatrix} \\mathbf{U}_{R,tr} \\\\ \\mathbf{U}_{R,te} \\end{bmatrix} \\mathbf{Z}_{R}^{te}$$\n",
    "\n",
    "We first take advantage of the top block of the equation, using the basis we just identified, to obtain an estimate for the scores during the test samples $\\mathbf{\\hat Z}_{R}^{te}$, i.e. we solve the following linear regression problem:\n",
    "\n",
    "$$\\mathbf{X}^{te}_{tr} = \\mathbf{U}_{R,tr}\\mathbf{Z}_{R}^{te}$$\n",
    "Now we plus this into the bottom block of the equation to obtain our estimate for the test neurons during test samples:\n",
    "$$\\mathbf{\\hat X}^{te}_{te} = \\mathbf{U}_{R,te}\\mathbf{\\hat Z}_{R}^{te}$$\n",
    "We can now calculate the fraction of explained variance in the estimate as we [normally would](https://en.wikipedia.org/wiki/Fraction_of_variance_unexplained#:~:text=In%20statistics%2C%20the%20fraction%20of,by%20the%20explanatory%20variables%20X.):\n",
    "$$ \\text{frac var exp} = 1 - \\frac{\\text{MSE}}{\\text{var}}$$\n",
    "We repeat for different $R$ and as usual look for the value that maximizes the cross-validated explained variance.\n",
    "\n",
    "There are a couple practicalities to note from this method that stem from the linear regression problem. First, note that the number of neurons in the training set has to be greater than $R$, or else the linear regression problem is underdetermined. This puts a limit to what dimensionalities we can test depending on how many neurons we can afford to use for training. Second, since we are using a subset of neurons to predict activity in other neurons, this method does not work well if the test neurons are independent of the training neurons. In other words it assumes some amount of shared variability between the neurons.\n",
    "\n",
    "However, this formulation of bi-cross-validation is very particular to matrix decompositions and does not extend easily to tensors or to nonlinear methods."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true,
    "id": "VqwfH7fNZD_p"
   },
   "source": [
    "### Cross-validation for tensors\n",
    "\n",
    "To cross-validate other methods, we can simply hold out or mask entries of the data matrix/tensor while fitting the parameters of the model (e.g. the $\\mathbf{u}$'s) on the remaining entries [4]. However, we must be careful to take into account temporal correlations since the tensor entries are not independently generated. That is, if we hold out single entries randomly throughout the matrix, the training and testing dataset will be strongly correlated. In fact, assuming the data is continuous in time (as would be the case for imaging data, or with firing rate estimates), with a sufficiently high sampling rate (i.e. sufficiently small time bins), the entry $\\mathcal{X}_{i,j,k}$ of the data tensor (supposed to be in the test set) can be arbitrarily well approximated by averaging the $\\mathcal{X}_{i,j-1,k}$ and $\\mathcal{X}_{i,j+1,k}$ (which may both be in the train set). Instead, we advocate for holding out blocks and trimming the ends for the most fair separation between test and training data [5]. What timescale should we use? Some good bets are the timescale of the calcium indicator (imaging data) the smoothing kernel (for ephys). Or we can simply use the autocorrelation timescale for any arbitrary time series.\n",
    "\n",
    "<p align=\"center\" width=\"100%\">\n",
    "<img src=\"https://drive.google.com/uc?id=13dBa2-ccAWAeXf9oasoNJTO3bNvfPJOW\" width=\"300px;\">\n",
    "</p>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HUS6Snb8a3bc"
   },
   "source": [
    "## Application of PCA bi-cross validation on a toy dataset\n",
    "\n",
    "Now we will move on to a more realistic  model, where the true rank of the data is unclear due to noise. We will use ___PCA bi-cross-validation___ to try to guess the cross-validated dimensionality. For this we will generate higher-dimensional data. Execute the following code to generate your new data matrix $M$ of 100 neurons and 500 time points to get started.\n",
    "\n",
    "*(Note that for time series data we would need to be careful about temporal correlations between adjacent entries. In this particular tutorial we draw the samples from a multivariate Gaussian with no notion of time so we don't have to worry about it!)*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-25T17:12:17.401801Z",
     "start_time": "2024-03-25T17:12:17.384282Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "TFpJqSwIPWtu",
    "outputId": "2b605f87-1714-47df-bbe3-233156173633"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "np.random.seed(1)\n",
    "\n",
    "n = 100\n",
    "k = np.random.randint(1, n//3+1) # Mystery rank\n",
    "sample_size = 500\n",
    "alpha = 0.0 # Noise term\n",
    "\n",
    "u = np.random.randn(n, k)/np.sqrt(n*k)\n",
    "v = np.random.randn(n, n)/np.sqrt(n*k)\n",
    "C = u @ u.T + alpha*(v @ v.T) # Positive semi-definite symmetric matrix <= covariance matrix\n",
    "\n",
    "M = np.random.multivariate_normal(np.random.randn(n)/np.sqrt(n), C, sample_size) + np.random.randn(sample_size, n)/10\n",
    "\n",
    "print('Shapes:')\n",
    "print('  Data matrix (sample_size, neuron):', M.shape)\n",
    "print('  Covariance matrix (neuron, neuron):', C.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Step 0.** Before starting, it is useful to define some helper functions, computing the mean squared error (MSE) and the explained variance, since we will use them a lot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_mse(prediction, target):\n",
    "    \"\"\"\n",
    "    Function computing the mean squared error between predicted and target data\n",
    "    Inputs: \n",
    "        prediction (numpy array): predicted values\n",
    "        target (numpy array): target values\n",
    "    returns:\n",
    "        mse (float): mean squared error\n",
    "    \"\"\"\n",
    "    # Fill in your code below\n",
    "    ##############################\n",
    "    return # your code\n",
    "    ##############################\n",
    " \n",
    "def compute_var_exp(prediction, target):\n",
    "    \"\"\"\n",
    "    Function computing the variance explained between predicted and target data.\n",
    "    It is usually defined as 1 - MSE/mean(target^2)\n",
    "    Inputs: \n",
    "        prediction (numpy array): predicted values\n",
    "        target (numpy array): target values\n",
    "    returns:\n",
    "        var_exp (float): variance explained\n",
    "    \"\"\"\n",
    "    # Fill in your code below\n",
    "    ##############################\n",
    "    return # your code\n",
    "    ##############################"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wU7kreYFZucm"
   },
   "source": [
    "**Step 1.** Plot the cumulative fraction of variance explained in $M$ as a function of the number of principal components. You can use either the eigendecomposition or the SVD, as you prefer. Can you guess what the dimensionality is?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ztFHs2dmvHZS"
   },
   "outputs": [],
   "source": [
    "# Fill in your code below\n",
    "##############################\n",
    "\n",
    "# Center the data\n",
    "M_centered = # your code\n",
    "\n",
    "# Compute the explained variance\n",
    "# your code here\n",
    "#\n",
    "\n",
    "# (if you use eigendecomposition, don't forget to sort the eigenvalues)\n",
    "# your code here\n",
    "\n",
    "# compute the cumulative fraction of variance explained (normalized)\n",
    "cumulative = # your code here\n",
    "\n",
    "##############################\n",
    "\n",
    "fig = plt.figure(figsize=(6,3), constrained_layout=True)\n",
    "ax = fig.add_subplot()\n",
    "ax.plot(cumulative,'o-')\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5lI_-eJBaHGl"
   },
   "source": [
    "**Step 2.** Now let's start the bi-cross-validation. Do the following steps.\n",
    "* First (we already did it for you), split the data into training and test sets in both time and trials. We have written an example code using ```train_test_split``` from ```sklearn``` but feel free to modify / implement it your own way. Use 80% for training both in neurons and in time. Make sure it's clear to keep track of which indices are neurons and which are time.\n",
    "* Second, choose a specific rank $r$ to first test your code. (But remember that it can't be larger than the number of neurons used for training)\n",
    "* Third, apply PCA on the data for training samples (all neurons) using SVD.\n",
    "* Next, use linear regression to solve for the scores in the test samples. You might want to use ```np.linalg.lstsq```.\n",
    "* Finally, make a scatter plot of the actual data for the training neurons, test samples against the values you predicted from PCA and linear regression (also for training neurons, test samples - so it's not yet cross validated). Calculate and print the fraction of variance explained. How accurate is it?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wLEPd9S3zev3"
   },
   "outputs": [],
   "source": [
    "# Creating the indeces for selecting train/test sets over samples/neurons\n",
    "neur_indices_tr,neur_indices_te = train_test_split(np.arange(M.shape[1]), train_size = .8, shuffle=True)\n",
    "time_indices_tr,time_indices_te = train_test_split(np.arange(M.shape[0]), train_size = .8, shuffle=True)\n",
    "\n",
    "# Selecting the training and test data\n",
    "M_tr_all = M_centered[time_indices_tr,:] # Note that data must be centered before!\n",
    "M_te_all = M_centered[time_indices_te,:]\n",
    "M_te_tr = M_te_all[:,neur_indices_tr]\n",
    "M_te_te = M_te_all[:,neur_indices_te]\n",
    "\n",
    "\n",
    "# Fill in your code below\n",
    "############################\n",
    "\n",
    "# Choose the rank r\n",
    "# your code here\n",
    "\n",
    "# Apply SVD on the training data\n",
    "# your code here\n",
    "\n",
    "# Now, we solve the scores with linear regression. Let's use 'u' (small) for the train basis, \n",
    "# 'z' for the scores and 'x' forthe data\n",
    "u = # your code here\n",
    "x = # your code here\n",
    "\n",
    "# Use linear regression ( u z = x) to solve the scores on test data\n",
    "z = np.linalg.lstsq(u,x,rcond=None)[0]\n",
    "\n",
    "# Compute the prediction\n",
    "M_te_tr_predict = # your code here\n",
    "\n",
    "############################\n",
    "\n",
    "# Compute variance explained\n",
    "var_exp_tr = compute_var_exp(M_te_tr_predict,M_te_tr)\n",
    "\n",
    "fig = plt.figure(figsize=(4,3), constrained_layout=True)\n",
    "ax = fig.add_subplot()\n",
    "ax.scatter(M_te_tr_predict.flatten(), M_te_tr.flatten(),alpha=0.5)\n",
    "ax.axline((0,0), (0.7,0.7), color='black', linestyle='--')\n",
    "print(var_exp_tr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2bUWo2dwbXkP"
   },
   "source": [
    "**Step 3.** Now we will try to do the cross validation to the test neurons, test samples data.\n",
    "* Predict the activity of the test neurons in the test dataset. Make the same plot as in the last cell, and calculate the cross-validated fraction of variance explained. How does it compare to the previous step?\n",
    "* Repeat this and the last cell for a few different choices of rank. What happens when the rank is very small ? Very large ?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "T-NDQKws1dl3"
   },
   "outputs": [],
   "source": [
    "# Fill in your code below\n",
    "##########################\n",
    "\n",
    "# Compute the prediction on test neurons \n",
    "M_te_te_predict = # your code here\n",
    "\n",
    "# Compute variance explained\n",
    "var_exp_te = # your code here\n",
    "\n",
    "##########################\n",
    "\n",
    "fig = plt.figure(figsize=(4,4), constrained_layout=True)\n",
    "ax = fig.add_subplot()\n",
    "ax.scatter(M_te_te_predict.flatten(), M_te_te.flatten(),alpha=0.5)\n",
    "ax.axline((-0.5,-0.5), (0.5,0.5), color='black', linestyle='--')\n",
    "plt.show()\n",
    "\n",
    "print(\"Variance explained (test): \",var_exp_te)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Ft9VINuQbjn7"
   },
   "source": [
    "**Step 4.** Here we will do the bi-cross-validation systematically to try to guess the rank.\n",
    "* Copy/paste the code above into a loop to calculate the fraction of variance explained (both in training and test neurons) for different numbers of components. Test all possible dimensionalities.\n",
    "* Plot the variance explained in both the training and test neurons over the number of components. What's the dimensionality?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qaFomKtE1b0H"
   },
   "outputs": [],
   "source": [
    "# enter here your guess for the rank\n",
    "your_guess = 3\n",
    "\n",
    "\n",
    "r_max = M_te_tr.shape[1]\n",
    "dims = np.arange(1,r_max+1)\n",
    "var_exp_tr = np.empty(dims.shape)\n",
    "var_exp_te = np.empty(dims.shape)\n",
    "\n",
    "for r in dims:\n",
    "    # Reuse the last two cells of code for computing train/test explained variance\n",
    "    # Fill in your code below\n",
    "    ###########################\n",
    "\n",
    "    # Apply SVD on the training data\n",
    "    U, S, V = np.linalg.svd(M_tr_all, full_matrices=False)\n",
    "\n",
    "    # Now, we solve the scores with linear regression. Let's use 'u' for the train basis, \n",
    "    # 'z' for the scores and 'x' forthe data\n",
    "    u = # your code here\n",
    "    x = # your code here\n",
    "    # Use linear regression ( u z = x) to solve the scores on test data\n",
    "    z = np.linalg.lstsq(u,x,rcond=None)[0] # scores\n",
    "\n",
    "    # Compute the prediction for training neurons\n",
    "    M_te_tr_predict = # your code here\n",
    "\n",
    "    # Compute explained variance\n",
    "    var_exp_tr[r-1] = compute_var_exp(M_te_tr_predict,M_te_tr)\n",
    "\n",
    "    # Compute the prediction on test neurons\n",
    "    M_te_te_predict = # your code here\n",
    "\n",
    "    # Compute variance explained\n",
    "    var_exp_te[r-1] = compute_var_exp(M_te_te_predict,M_te_te)\n",
    "    ###########################\n",
    "\n",
    "fig = plt.figure(figsize=(4,3), constrained_layout=True)\n",
    "ax = fig.add_subplot()\n",
    "ax.plot(dims,var_exp_tr,'-o', label = \"train data\")\n",
    "ax.plot(dims,var_exp_te,'-o', label = \"test data\")\n",
    "ax.legend()\n",
    "ax.set_xlabel(\"Rank\")\n",
    "ax.set_ylabel(\"Variance explained\")\n",
    "ax.set_ylim([0,1.05])\n",
    "plt.show()\n",
    "\n",
    "k_guess = dims[np.argmax(var_exp_te)]\n",
    "\n",
    "print(\"Computed rank: \", k_guess)\n",
    "print(\"Real rank (the mistery rank k):\", k)\n",
    "print('Your guess was '+('Correct !' if your_guess == k_guess else 'Incorrect !'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Vu9CqeQFb2zP"
   },
   "source": [
    "# **Part IV.** Application to neural data analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "The data consists of microelectrode array recordings from M1 and PMd during a delayed center-out reach task. It is one of the datasets from the [Neural Latent Benchmark](https://neurallatents.github.io/) project, whose aim is to provide standardized freely accessible neural datasets for the sake of developing data analysis and modeling methods.\n",
    "\n",
    "\n",
    "\n",
    "---\n",
    "\n",
    "\n",
    "\n",
    "The task is a delayed center-out reach task with barriers, resulting in a variety of straight and curved trajectories. Neural activity was recorded from the dorsal premotor and primary motor cortices, and cursor, monkey gaze position, and monkey hand position and velocity are also provided. You can find more information on the dataset in the articles [7-8].\n",
    "\n",
    "\n",
    "<p align=\"center\" width=\"100%\">\n",
    "<img src=\"https://neurallatents.github.io/assets/maze_fig1.png\" width=\"550px;\">\n",
    "</p>\n",
    "\n",
    "---\n",
    "It has been found that M1/PMd activity play a key role in motor control during this type of tasks.\n",
    "___Today, we will explore this hypothesis by analyzing how the activity of the population of neurons is able to decode behavioral information.___"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZhA_2V8ZkTav"
   },
   "source": [
    "## Data imports\n",
    "\n",
    "The neural data is presented in the Neurodata Without Border (NWB) format, which is a commonly used format for storing and sharing neural data. Here, all the code was written for you, you may just run the following cells, they will download the part of the dataset from this experiment which is relevant to this tutorial: one recording session. <br>\n",
    "___Note that this can take a while (<5min) to run. So you should avoid restarting your runtime environment or rerunning those cells.___"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-khOJ_4KHdhn"
   },
   "source": [
    "### Run this (offline users)\n",
    "\n",
    "__We install some packages using pip. You will need to adapt the first cell code if you use conda!__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Import/install necessary libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "00nl4ELuKj1t"
   },
   "outputs": [],
   "source": [
    "# You should not evaluate this cell multiple times\n",
    "\n",
    "from IPython.display import clear_output\n",
    "\n",
    "# Download the data\n",
    "!pip install pandas\n",
    "!pip install --upgrade numpy\n",
    "!pip install git+https://github.com/neurallatents/nlb_tools.git\n",
    "!pip install dandi\n",
    "!pip install pydantic[email]\n",
    "\n",
    "clear_output() #clear cell output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-25T18:25:41.060456Z",
     "start_time": "2024-03-25T18:25:41.031801Z"
    },
    "id": "KcHq3W1qFt-8"
   },
   "outputs": [],
   "source": [
    "from nlb_tools.nwb_interface import NWBDataset\n",
    "from IPython.display import clear_output\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pickle\n",
    "\n",
    "np.random.seed(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Download/load the raw dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The continuous data provided with the MC_Maze datasets includes:\n",
    "* `cursor_pos` - x and y position of the cursor controlled by the monkey\n",
    "* `eye_pos` - x and y position of the monkey's point of gaze on the screen, in mm\n",
    "* `hand_pos` - x and y position of the monkey's hand, in mm\n",
    "* `hand_vel` - x and y velocities of the monkey's hand, in mm/s, computed offline using `np.gradient`\n",
    "* `spikes` - spike times binned at 1 ms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "_JClxUMAKoCr",
    "outputId": "b0559c1f-9eef-4dfc-9064-736af93cfffe"
   },
   "outputs": [],
   "source": [
    "#The part of the dataset of interest\n",
    "\n",
    "!dandi download https://gui.dandiarchive.org/#/dandiset/000128\n",
    "dataset = NWBDataset(\"000128/sub-Jenkins/\", \"*train\", split_heldout=False)\n",
    "clear_output()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Dataset visualization (table)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tmxrx0frKvP1"
   },
   "outputs": [],
   "source": [
    "dataset.data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ygvPEzOnKxmr"
   },
   "outputs": [],
   "source": [
    "dataset.trial_info"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Dataset preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kSYTVBhT4u0T"
   },
   "outputs": [],
   "source": [
    "#The part of the dataset of interest\n",
    "bin_size = 5 # ms\n",
    "time_window = (-1000, 800)\n",
    "timesteps = int((time_window[1]-time_window[0])//bin_size)\n",
    "n_neurons = len(dataset.data.spikes.columns)\n",
    "\n",
    "# resample the dataset\n",
    "dataset.resample(bin_size)\n",
    "# align trials to movement onset\n",
    "alldata = dataset.make_trial_data(align_field='move_onset_time', align_range=time_window)\n",
    "\n",
    "trials = alldata.trial_id.unique()\n",
    "clear_output()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trials = alldata.trial_id.unique()\n",
    "\n",
    "nns = np.zeros([len(trials),n_neurons, timesteps]) # here we will save the neural activity\n",
    "hhs = np.zeros([len(trials), 4, timesteps]) # here we will save the hand position\n",
    "info_df = [] # this will contain the trial info\n",
    "\n",
    "names = ['trial', 'ttype', 'tversion', 'target_x', 'target_y', 'success', 'num_targets', 'active']\n",
    "\n",
    "for t, df in alldata.groupby('trial_id'):\n",
    "    nns[t] = df.spikes.values.T\n",
    "    hhs[t] = np.concatenate([df.hand_pos.values.T, df.hand_vel.values.T],axis=0)\n",
    "\n",
    "    info_df.append([\n",
    "        dataset.trial_info.trial_id[t],\n",
    "        dataset.trial_info.trial_type[t],\n",
    "        dataset.trial_info.trial_version[t],\n",
    "        dataset.trial_info.target_pos[t][0][0], \n",
    "        dataset.trial_info.target_pos[t][0][1],\n",
    "        dataset.trial_info.success[t],\n",
    "        dataset.trial_info.num_targets[t],\n",
    "        dataset.trial_info.active_target[t]\n",
    "\t])\n",
    "\n",
    "info_df = pd.DataFrame(info_df, columns=names)\n",
    "\n",
    "# here we extract a part of the dataset for faster computations\n",
    "## Note that we selected sparser targets, for visualization purposes\n",
    "cond = [26, 29, 1, 13, 32, 28, 33, 23, 10, 6 ,20, 38]\n",
    "\n",
    "info_df = info_df.loc[info_df.num_targets==1]\n",
    "info_df = info_df.loc[np.isin(info_df.ttype, cond)]\n",
    "\n",
    "info_straight = info_df.loc[info_df.tversion==0].reset_index(drop=True)\n",
    "info_curved = info_df.loc[info_df.tversion==1].reset_index(drop=True)\n",
    "\n",
    "hand_straight = hhs[info_straight.trial.values]\n",
    "hand_curved = hhs[info_curved.trial.values]\n",
    "\n",
    "neurons_straight = nns[info_straight.trial.values]\n",
    "neurons_curved = nns[info_curved.trial.values]\n",
    "\n",
    "info_straight.to_pickle('info_straight.p')\n",
    "info_curved.to_pickle('info_curved.p')\n",
    "\n",
    "np.save('hand_straight.npy', hand_straight)\n",
    "np.save('hand_curved.npy', hand_curved)\n",
    "\n",
    "np.save('neurons_straight.npy', neurons_straight)\n",
    "np.save('neurons_curved.npy', neurons_curved)\n",
    "\n",
    "del alldata, dataset, nns, hhs, info_df, info_straight, info_curved, hand_straight, hand_curved, neurons_straight, neurons_curved"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "### Run this (colab users)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "from IPython.display import clear_output\n",
    "\n",
    "!pip install pandas\n",
    "!pip install git+https://github.com/neurallatents/nlb_tools.git\n",
    "!pip install gdown\n",
    "\n",
    "# download the preprocessed data\n",
    "!gdown https://drive.google.com/uc?id=1fhmEvzfelRuYn0ZWpoxTrk8eFLIqh3Jh -O adv-dim-red.zip\n",
    "!unzip adv-dim-red.zip\n",
    "!rm adv-dim-red.zip\n",
    "!mv adv-dim-red-data/* .\n",
    "!rm -r adv-dim-red-data\n",
    "\n",
    "clear_output()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loading the dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we load the neural data, behavioral data and extra trial data information. <br><br>\n",
    "__We will mainly focus our analysis on the straight part of the dataset.__ <br><br>\n",
    "\n",
    "If you are curious (and if you have some spare time once finished), you could also try to have a look at the other part of the dataset, where the monkeys need to reach the target while avoiding some obstacles. This force the trajectories to be longer, more curved and complex. You can find these data in the files `neurons_curved.py`, `hand_curved.py` and `info_curved.p`.\n",
    "\n",
    "But let's focus on the simple case. Here, we load and print some info on the dataset. Just run the cell, no need to write anything here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-25T18:24:47.692748Z",
     "start_time": "2024-03-25T18:24:47.528536Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "from sklearn.linear_model import Ridge\n",
    "from sklearn.model_selection import GridSearchCV\n",
    "from sklearn.preprocessing import StandardScaler, MinMaxScaler\n",
    "from sklearn.model_selection import train_test_split\n",
    "import pickle\n",
    "from tqdm import tqdm\n",
    "\n",
    "import scipy.ndimage as spnd\n",
    "\n",
    "np.random.seed(10)\n",
    "\n",
    "\n",
    "# loading neural and behavioral datasets\n",
    "neural_data = np.load('neurons_straight.npy')\n",
    "behavioral_data = np.load('hand_straight.npy')\n",
    "\n",
    "# task variables\n",
    "with open('info_straight.p', 'rb') as f:\n",
    "    trial_data = pickle.load(f).reset_index()\n",
    "\n",
    "# defining some useful constants\n",
    "time_window = (-1000, 800) # ms\n",
    "bin_size = 5 # ms\n",
    "onset_time = -time_window[0]//bin_size\n",
    "n_trials, n_neurons, timesteps = neural_data.shape\n",
    "time_axis = np.arange(time_window[0], time_window[1], bin_size)\n",
    "\n",
    "\n",
    "print('\\033[1;4m# Data format\\033[0m\\n')\n",
    "print('Time window:', time_window, 'ms\\n')\n",
    "print('Bin size:', bin_size, 'ms\\n')\n",
    "print('Movement onset: 0 ms\\n')\n",
    "print('Neural data shape (trials, neurons, time bins):\\n', neural_data.shape)\n",
    "print('\\nBehavioral data shape (trials, x/y position + x/y velocity, time bins):\\n', behavioral_data.shape)\n",
    "print('\\ntrial_data (Dataframe) columns\\n',list(trial_data.columns))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Some preprocessing\n",
    "\n",
    "Before starting our analysis, we will need some preprocessing on the data. <br>\n",
    "First, since we have several targets, we need a simple variable for labeling each trial. We will use the x/y target positions to extract the target angle, our label. For helping us plotting, we will assign a specific color for each target angle. This part is already coded.\n",
    "\n",
    "Then, we need to convert the PSTH dataset into firing rates estimates. A simple way to do it is to ___apply a gaussian filter on it, applied on the time axis for each specific neuron___. \n",
    "- Since our dataset might not be binned at 1ms (info in the cell above), first compute the gaussian std in bins that you want to use for the smoothing. We will use a std of 50 ms.\n",
    "- Apply the smoothing on the neural data, saving it in a variable called `rate_neural_data`. If you don't know how to do the smoothing, you can try having a look at `scipy.ndimage.spnd.gaussian_filter1d()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate target vangle and get index for sorting trials by target angle\n",
    "trial_data['angle'] = np.arctan2(trial_data.target_y.values, trial_data.target_x.values)\n",
    "\n",
    "# Here, we define a color map for the different targets. \n",
    "# You can use 't_map[angle]' for getting the color\n",
    "targets = np.sort(trial_data['angle'].unique())\n",
    "colors = matplotlib.cm.rainbow(np.linspace(0,1,len(targets)))\n",
    "t_map = {angle: colors[i] for i, angle in enumerate(targets)}\n",
    "\n",
    "# Fill in your code below\n",
    "##############################\n",
    "# Gaussian filter time window\n",
    "gaussian_filter_std = 50 # ms\n",
    "\n",
    "# convert the std in bins size\n",
    "gaussian_filter_std_bins = # your code here\n",
    "\n",
    "# smooth along time axis, using a Gaussian filter (careful with the axis)\n",
    "rate_neural_data = # your code here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "id": "kXjoeNrXlWR5"
   },
   "source": [
    "## Simple statistics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true,
    "id": "A_hVAJmmmT_K"
   },
   "source": [
    "### Analysis of behavior\n",
    "\n",
    "__Step 1 (Hand Trajectories).__ First, let's plot the hand trajectories. Try to make two subplots where trajectories are color coded by the target angle (present in the trail_data).\n",
    "- In the first plot, plot all the color coded trajectories\n",
    "- in the second plot, just plot the mean trajectory for each target angle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1,2,figsize=(12,6), dpi=60,constrained_layout=True)\n",
    "\n",
    "# let's extract the angles from trial_data\n",
    "angles = trial_data['angle'].values\n",
    "\n",
    "# Here, we define a color map for the different targets. You can use 't_map[angle]' for getting the color\n",
    "targets = np.sort(trial_data['angle'].unique())\n",
    "colors = matplotlib.cm.rainbow(np.linspace(0,1,len(targets)))\n",
    "t_map = {angle: colors[i] for i, angle in enumerate(targets)}\n",
    "\n",
    "for trial_id,trial_traj in enumerate(behavioral_data):\n",
    "    # Fill in your code below\n",
    "    ###############################\n",
    "    # remember that trials and angles are in the same order\n",
    "    angle = # your code\n",
    "\n",
    "    ax[0].plot() # your code here\n",
    "    ###############################\n",
    "\n",
    "ax[0].set_xlabel('x (mm)')\n",
    "ax[0].set_ylabel('y (mm)')\n",
    "ax[0].set_title('Hand movement, different trials')\n",
    "\n",
    "for angle in targets:\n",
    "    # Fill in your code below\n",
    "    ###############################\n",
    "    # create the index vector, for selecting the trajectories\n",
    "    indeces = # your code here\n",
    "\n",
    "    # extract the trajectories and compute the mean over trials\n",
    "    mean_trajectory = # your code here\n",
    "\n",
    "    ax[1].plot() #your code here\n",
    "    ###############################\n",
    "\n",
    "ax[1].set_xlabel('x (mm)')\n",
    "ax[1].set_ylabel('y (mm)')\n",
    "ax[1].set_title('Hand movement, average trajectory')\n",
    "\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "As you can see, the animal perform stereotyped, slightly curved trajectories.\n",
    "\n",
    "__Step 2 (Hand Velocity).__ Let's now plot the hand velocity. You can reuse the same code, just selecting the indeces in the behavioral_data matrix corresponding to velocity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1,2,figsize=(12,6), dpi=60,constrained_layout=True)\n",
    "\n",
    "# let's extract the angles from trial_data\n",
    "angles = trial_data['angle'].values\n",
    "\n",
    "# Here, we define a color map for the different targets.\n",
    "# You can use 't_map[angle]' for getting the color\n",
    "targets = np.sort(trial_data['angle'].unique())\n",
    "colors = matplotlib.cm.rainbow(np.linspace(0,1,len(targets)))\n",
    "t_map = {angle: colors[i] for i, angle in enumerate(targets)}\n",
    "\n",
    "for trial_id,trial_traj in enumerate(behavioral_data):\n",
    "    # Fill in your code below\n",
    "    ###############################\n",
    "    # remember that trials and angles are in the same order\n",
    "    angle = # your code\n",
    "\n",
    "    ax[0].plot() # your code here\n",
    "    ###############################\n",
    "\n",
    "ax[0].set_xlabel('x (mm/s)')\n",
    "ax[0].set_ylabel('y (mm/s)')\n",
    "ax[0].set_title('Hand velocity, different trials')\n",
    "\n",
    "for angle in targets:\n",
    "    # Fill in your code below\n",
    "    ###############################\n",
    "    # create the index vector, for selecting the trajectories\n",
    "    indeces = # your code here\n",
    "\n",
    "    # extract the trajectories and compute the mean over trials\n",
    "    mean_trajectory = # your code here\n",
    "    ###############################\n",
    "    ax[1].plot()\n",
    "\n",
    "ax[1].set_xlabel('x (mm/s)')\n",
    "ax[1].set_ylabel('y (mm/s)')\n",
    "ax[1].set_title('Hand velocity, average')\n",
    "\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Much more messy right? Before moving on, just try to spot the main differences between trajectories and velocities! What do you see?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true,
    "id": "lWiB0VoTKgtQ"
   },
   "source": [
    "__Step 3.__ Here, we plot the single x or y velocities, for gaining some information on movement timing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(3,1,figsize=(5,7), dpi=100, constrained_layout=True)\n",
    "\n",
    "chosen_angles = targets[[2,5,10]]\n",
    "\n",
    "for a in chosen_angles:\n",
    "\n",
    "    # Fill in your code below\n",
    "    ############################\n",
    "\n",
    "    # create the index vector for selected angle and extract the trajectories\n",
    "    indeces = # your code here\n",
    "    velocities = # your code here\n",
    "\n",
    "    ### Plotting hand mean velocity\n",
    "\n",
    "    # compute the velocity module, mean/std over trials\n",
    "    velocity_module = # your code here\n",
    "    mean_module = # your code here\n",
    "    std_module = # your code here\n",
    "\n",
    "    ############################\n",
    "\n",
    "    # let's plot mean and std\n",
    "    ax[0].plot(time_axis, mean_module, \n",
    "               c=t_map[a], linewidth=1, label=f'${a/np.pi*180:.2f}^\\circ$')\n",
    "    ax[0].fill_between(time_axis, \n",
    "                       mean_module-std_module, \n",
    "                       mean_module+std_module, \n",
    "                       color=t_map[a], alpha=0.2)\n",
    "\n",
    "    ### Plotting hand x-velocity\n",
    "\n",
    "    # Fill in your code below\n",
    "    ############################\n",
    "\n",
    "    # compute the x_velocity module, mean/std over trials\n",
    "    x_velocity = # your code here\n",
    "    x_velocity_mean = # your code here\n",
    "    x_velocity_std = # your code here\n",
    "\n",
    "    ############################\n",
    "\n",
    "    # \n",
    "    ax[1].plot(time_axis, x_velocity_mean, \n",
    "               c=t_map[a], linewidth=1, label=f'${a/np.pi*180:.2f}^\\circ$')\n",
    "    ax[1].fill_between(time_axis, \n",
    "                       x_velocity_mean-x_velocity_std, \n",
    "                       x_velocity_mean+x_velocity_std, \n",
    "                       color=t_map[a], alpha=0.2)\n",
    "\n",
    "    ### Plotting hand y-velocity\n",
    "\n",
    "    # Fill in your code below\n",
    "    ############################\n",
    "\n",
    "    # compute the x_velocity module, mean/std over trials\n",
    "    y_velocity = # your code here\n",
    "    y_velocity_mean = # your code here\n",
    "    y_velocity_std = # your code here\n",
    "\n",
    "    ############################\n",
    "\n",
    "    ax[2].plot(time_axis, y_velocity_mean, \n",
    "               c=t_map[a], linewidth=1, label=f'${a/np.pi*180:.2f}^\\circ$')\n",
    "    ax[2].fill_between(time_axis, \n",
    "                       y_velocity_mean-y_velocity_std, \n",
    "                       y_velocity_mean+y_velocity_std, \n",
    "                       color=t_map[a], alpha=0.2)\n",
    "\n",
    "# Some settings\n",
    "ax[0].set_title('Hand velocity module')\n",
    "ax[0].set_ylabel('|Velocity| (mm/s)')\n",
    "ax[1].set_title('Hand x-velocity')\n",
    "ax[1].set_ylabel('Velocity (mm/s)')\n",
    "ax[2].set_title('Hand y-velocity')\n",
    "ax[2].set_ylabel('Velocity (mm/s)')\n",
    "for i in range(3):\n",
    "    ax[i].vlines(0, -700,1400, color='k', linestyle='--', linewidth=1, label = 'mov. onset',zorder=3)\n",
    "    ax[i].set_ylim(-700,1400)\n",
    "    ax[i].legend(loc='upper left')\n",
    "    ax[i].set_xlabel('Time (ms)')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "What do you observe? Is there some parts of the data data might not be useful to include for our analysis?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true,
    "id": "1yA8Y7MLUWH3"
   },
   "source": [
    "### Analysis of the neural data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true,
    "id": "lESeuP4FAwI5"
   },
   "source": [
    "#### Neuron-wise average\n",
    "\n",
    "Plot a histogram of the average number of spikes per trial for all neurons."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# Fill in your code below"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true,
    "id": "PUcFmXkg5xth"
   },
   "source": [
    "*  What do you observe? (You can use more (~50) bins for improving visualization)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true,
    "id": "TXxNUD5p-Tup"
   },
   "source": [
    "#### Time-wise average\n",
    "\n",
    "We plot now the average (over neurons and trials) number of spikes as a function of time, and the same with its standard deviation to compare them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1,2,figsize=(12,4),constrained_layout=True)\n",
    "\n",
    "# Fill in your code below\n",
    "##############################\n",
    "\n",
    "# Compute the mean and standard deviation of the activity\n",
    "mean_activity = # your code here\n",
    "std_activity = # your code here\n",
    "\n",
    "# plot the mean here\n",
    "ax[0].plot() # your code here\n",
    "\n",
    "# same as above\n",
    "ax[1].plot() # your code here\n",
    "\n",
    "# plot the std\n",
    "ax[1].fill_between() # your code here\n",
    "\n",
    "##############################\n",
    "\n",
    "\n",
    "# other settings\n",
    "ax[0].axvline(0, color='red', label='mov. onset', linestyle='--')\n",
    "ax[0].set_xlabel('Time (ms)')\n",
    "ax[0].set_ylabel('Average number of spikes')\n",
    "ax[0].legend()\n",
    "ax[1].axvline(0, color='red', label='mov. onset', linestyle='--')\n",
    "ax[1].set_xlabel('Time (ms)')\n",
    "ax[1].set_ylabel('Average number of spikes')\n",
    "ax[1].legend()\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true,
    "id": "Xm0yW1vT_rGC"
   },
   "source": [
    "*  What do you observe? Could we guess something from this plot?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "### Trial to trial variability\n",
    "\n",
    "As we just saw in the last plot, there is a lot of variability in the data, at basically any time. A good way to improve our understanding on it is to look at the neural activity more closely, plotting the PSTH for some neurons, as well as the per-target average firing rate over time.\n",
    "\n",
    "\n",
    "___Sort the spikes according to the target angle, in order to put together activity for the same target. This will help reading the data.___\n",
    "\n",
    "(You might want to check out other neurons once you have the code correctly running)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "neurons = [0,30,76]\n",
    "\n",
    "# Fill in your code here\n",
    "###############################\n",
    "\n",
    "# first we want to sort the spikes according to the target angle\n",
    "sorted_trial_data = # your code here\n",
    "\n",
    "# we extract the new index of the sorted trials\n",
    "trial_idx = # your code here\n",
    "\n",
    "# then, we extract the ordered angles label\n",
    "angles = # your code here\n",
    "\n",
    "# apply the same sorting to the neural data and the rates\n",
    "sorted_neural_data =  # your code here\n",
    "sorted_rate_neural_data = # your code here\n",
    "\n",
    "###############################\n",
    "\n",
    "# Code for plotting\n",
    "\n",
    "fig = plt.figure(figsize=(12,9), dpi=100)\n",
    "# heat maps sorted by target angle for three example neurons\n",
    "for ni,n in enumerate(neurons):\n",
    "    plt.subplot(3,3,ni+1)\n",
    "    plt.imshow(sorted_neural_data[:,n], aspect='auto', vmin=0, vmax=1,extent=[time_window[0],time_window[1], 0, len(rate_neural_data)], cmap='magma', origin='lower')\n",
    "    plt.plot([0,0], [0,len(sorted_neural_data)], 'w--')\n",
    "    if ni==0: \n",
    "        plt.ylabel('trials')\n",
    "    plt.title(f'neuron {n}, PSTHs')\n",
    "\n",
    "# heat maps sorted by target angle for three example neurons\n",
    "for ni,n in enumerate(neurons):\n",
    "    plt.subplot(3,3,ni+4)\n",
    "    plt.imshow(sorted_rate_neural_data[:,n], aspect='auto', extent=[time_window[0],time_window[1], 0, len(rate_neural_data)], cmap='magma', origin='lower')\n",
    "    plt.plot([0,0], [0,len(sorted_rate_neural_data)], 'w--')\n",
    "    if ni==0: \n",
    "        plt.ylabel('trials')\n",
    "    plt.title(f'neuron {n}, firing rates')\n",
    "\n",
    "# condition-averaged, scaled & smoothed PSTHs\n",
    "for ni,n in enumerate(neurons):\n",
    "    plt.subplot(3,3,ni+7)\n",
    "    for ai,a in enumerate(targets):\n",
    "        indeces = (angles == a)\n",
    "        plt.plot(time_axis, 0.015*ai+np.mean(sorted_rate_neural_data[indeces,n], axis=0), '-', color=t_map[a], zorder=-ai)\n",
    "        plt.plot([0,0], [-0.05,.4], 'k--')\n",
    "        if ni==0: \n",
    "            plt.ylabel('average activity per target')\n",
    "        plt.xlim([time_axis[0],time_axis[-1]])\n",
    "        plt.xlabel('time from movement onset (s)')\n",
    "        plt.tight_layout()\n",
    "    plt.ylim(-0.02,0.28)\n",
    "    plt.title(f'neuron {n}, target-averaged activity')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true,
    "id": "lsSFZeUXUexa"
   },
   "source": [
    "*  Qualitatively comment on the tuning of these three neurons (you may want to look at other neurons)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8TrnajlbKmsl"
   },
   "source": [
    "## Dimensionality reduction: PCA\n",
    "\n",
    "The above approaches have a key pitfall: information about the data is lost when averaging, and we fail to find properties of the whole neural population when looking at single neurons. We therefore turn to dimensionality reduction.\n",
    "\n",
    "As anticipated in **Part I**, PCA is a method that can be applied on matrices. Our dataset is a 3-dimensional array, a tensor, with shape `(n_trials, n_neurons, timesteps)`. For this reason we will need to make it 2-dimensional. This opens to the two types of scenarios we anticipated in section [From Matrices to Tensors](#tensors): neuron-unfolded and trial-unfolded PCA.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Getting ready"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "* Before applying PCA, we will first need to rescale our dataset. Although the standard way is to z-score it, we will use here a [min-max normalization](https://en.wikipedia.org/wiki/Feature_scaling), i.e. moving the dataset values to the range [0, 1].\n",
    "\n",
    "___Apply the min-max normalization for each specific neuron___\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fill in your code\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* Since we will keep using PCA, we will define our function for computing pca on unfolded data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_pca_unfolded(data,n_pcs=None):\n",
    "    \"\"\"\n",
    "    Function computing pca using SVD (np.linalg.svd).\n",
    "    The last dimension (data.shape[2]) will be the unfolded one.\n",
    "    If n_pcs is None, n_pcs = all components\n",
    "       Returns:\n",
    "       - the first n_pcs principal components of the data \n",
    "       - the eigenvector matrix\n",
    "    \"\"\"\n",
    "    # checking the inputs\n",
    "    assert data.ndim==3, \"Data must be 3 dimensional\"\n",
    "    assert isinstance(n_pcs,int) or n_pcs is None, \"If provided, n_pcs must be an integer\"\n",
    "\n",
    "    if n_pcs is None:\n",
    "        n_pcs = data.shape[2]\n",
    "    \n",
    "    # Fill in your code\n",
    "    ################################\n",
    "        \n",
    "    # unfold the data. The final shape will be (dim1*dim2, dim3)\n",
    "    unfolded_data = # your code here\n",
    "\n",
    "    ################################\n",
    "\n",
    "    # SVD\n",
    "    _, _, V = np.linalg.svd(unfolded_data, full_matrices=False)\n",
    "\n",
    "    # select the first n_pcs components\n",
    "    if n_pcs is not None:\n",
    "        V = V[:n_pcs]\n",
    "\n",
    "    # project the data on the \n",
    "    proj_data = data @ V.T\n",
    "\n",
    "    return proj_data, V, "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mdZ_rwNY4cSZ"
   },
   "source": [
    "### Neural covariability with PCA\n",
    "\n",
    "**Step 1.** Apply PCA to the neuron-unfolded neural data, using our new function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ODJbPbPqmtCn"
   },
   "outputs": [],
   "source": [
    "# Fill in your code below\n",
    "\n",
    "# Let's first check that the dataset has the correct shape\n",
    "correct_shape = # your code here\n",
    "flag = norm_rate_neural_data.shape == correct_shape\n",
    "assert flag, \"make sure that norm_rate_neural_data has correct shape.\"\n",
    "\n",
    "n_pcs = 20\n",
    "\n",
    "# apply PCA using the function compute_pca_unfolded\n",
    "rate_neural_data_scores, V_neuron = # your code here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2ob8mEEwTBcD"
   },
   "source": [
    "**Step 2.** Compute and plot cumulative variance explained as a function of the number of principal components. What do you see?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2qxuLOK5oNx0"
   },
   "outputs": [],
   "source": [
    "var_explained = []\n",
    "for number_of_pcs in range(1,k+1):\n",
    "    # Fill in your code here\n",
    "    ##########################\n",
    "    # let's select the first number_of_pcs principal components from the scores\n",
    "    sub_scores = # your code here\n",
    "\n",
    "    # use V_neuron to reconstruct the data\n",
    "    reconstructed_data = # your code here\n",
    "\n",
    "    ##########################\n",
    "    \n",
    "    # compute and save the variance explained\n",
    "    mse = ((reconstructed_data - norm_rate_neural_data)**2).mean()\n",
    "    var_explained.append(1-mse/norm_rate_neural_data.var())\n",
    "\n",
    "fig = plt.figure(figsize=(4,4), constrained_layout=True)\n",
    "ax = fig.add_subplot()\n",
    "\n",
    "ax.plot(np.arange(1, k+1), var_explained, '-o')\n",
    "ax.set_xlabel('Number of PCs')\n",
    "ax.set_ylabel('Variance Explained')\n",
    "ax.set_title('Neuron-unfolded PCA')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bmKo8PnsTq3w"
   },
   "source": [
    "**Step 3.**\n",
    "\n",
    "* Plot the neural activity along the first three PCs in a 3D space. The code was prepared for plotting with three different views, for better visualization.\n",
    "* Color each trajectory by the target angle.\n",
    "* You may start by plotting a single trajectory, for only two target angles.\n",
    "* Finally, increase the plotted trials to 5 and the angles to 3."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "MkvlMTAuomMP"
   },
   "outputs": [],
   "source": [
    "### Some settings\n",
    "\n",
    "# we extract the angles label\n",
    "angles = trial_data.angle.values\n",
    "# targets and color settings\n",
    "targets = np.sort(trial_data['angle'].unique())\n",
    "colors = matplotlib.cm.rainbow(np.linspace(0,1,len(targets)))\n",
    "t_map = {angle: colors[i] for i, angle in enumerate(targets)}\n",
    "\n",
    "# view settings. You can change them for rotating the plot. (elevation, azimut)\n",
    "views = [(30, -60),(30,120),(-50,-60)]\n",
    "\n",
    "# Some parameters you can tune\n",
    "selected_angles = targets[[0,2,5]]\n",
    "plotted_trials = 1\n",
    "principal_components = [0,1,2]\n",
    "\n",
    "\n",
    "### Plot\n",
    "fig = plt.figure(figsize=(10,3), constrained_layout=True)\n",
    "gs = fig.add_gridspec(1, 3)\n",
    "\n",
    "\n",
    "fig.suptitle(f'Neuron unfolded PCA {np.array(principal_components)+1}')\n",
    "# looping over views\n",
    "for k, view in enumerate(views):\n",
    "    ax = fig.add_subplot(gs[k],projection='3d')\n",
    "    elev, azim = view\n",
    "    for a in selected_angles:\n",
    "        # Fill in your code\n",
    "        ############################\n",
    "        \n",
    "        # create a trial mask by using 'angles' (already done before)\n",
    "        trial_indeces = # your code here\n",
    "\n",
    "        # extract the trajectories using trial \n",
    "        trajectories = # your code here\n",
    "\n",
    "        # select only the first plotted_trials-th ones\n",
    "        trajectories  = # your code here\n",
    "\n",
    "        # extract the principal components you want to plot\n",
    "        trajectories = # your code here\n",
    "        \n",
    "        ############################\n",
    "        for trajectory in trajectories:\n",
    "            ax.plot(trajectory[:, 0], trajectory[:, 1], trajectory[:, 2],\n",
    "                color = t_map[a], linewidth = 1.5, alpha=0.8)\n",
    "    \n",
    "    ax.view_init(elev=elev, azim=azim)\n",
    "    ax.set_xlabel(f'PC{principal_components[0]+1}')\n",
    "    ax.set_ylabel(f'PC{principal_components[1]+1}')\n",
    "    ax.set_zlabel(f'PC{principal_components[2]+1}')\n",
    "    \n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* What do you observe? not that much, right?\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "**Step 4** Before giving up, let's try to plot the target-averaged trajectories.\n",
    "* We will use the same code, but plotting the mean trajectory for each target angle.\n",
    "* For making the plot more readable, we will also highlight start/end points and the movement onset.\n",
    "* (Bonus) use `linestyle='o-'` to make the datapoint visible. This allows you to encode the time information. It might be useful to plot every 4/5 timesteps to properly visualize the points\n",
    "* (Bonus) try to explore different sets of principal components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Some settings\n",
    "\n",
    "# we extract the angles label\n",
    "angles = trial_data.angle.values\n",
    "# targets and color settings\n",
    "targets = np.sort(trial_data['angle'].unique())\n",
    "colors = matplotlib.cm.rainbow(np.linspace(0,1,len(targets)))\n",
    "t_map = {angle: colors[i] for i, angle in enumerate(targets)}\n",
    "\n",
    "# view settings. You can change them for rotating the plot. (elevation, azimut)\n",
    "views = [(30, -60),(30,120),(-50,-60)]\n",
    "\n",
    "# Some parameters you can tune\n",
    "selected_angles = targets\n",
    "principal_components = [0,1,3]\n",
    "mov_onset_time = onset_time\n",
    "\n",
    "\n",
    "### Plot\n",
    "fig = plt.figure(figsize=(15,5), constrained_layout=True)\n",
    "gs = fig.add_gridspec(1, 3)\n",
    "\n",
    "\n",
    "fig.suptitle(f'Neuron unfolded PCA {np.array(principal_components)+1}')\n",
    "# looping over views\n",
    "for k, view in enumerate(views):\n",
    "    ax = fig.add_subplot(gs[k],projection='3d')\n",
    "    elev, azim = view\n",
    "    for a in selected_angles:\n",
    "\n",
    "        # Fill in your code\n",
    "        ############################\n",
    "\n",
    "        # create a trial mask by using 'angles' (already done before)\n",
    "        trial_indeces = # your code here\n",
    "\n",
    "        # extract the trajectories\n",
    "        trajectories = # your code here\n",
    "        \n",
    "        # compute the mean over trials\n",
    "        mean_trajectory  = # your code here\n",
    "        \n",
    "        # extract the principal components you want to plot\n",
    "        mean_trajectory = # your code here\n",
    "        \n",
    "        # plot mean_trajectory\n",
    "        ax.plot( # your code here, use the next line settings for better visualization \n",
    "            color = t_map[a], markersize=3,alpha=0.8, linewidth=1)\n",
    "        # plot start points\n",
    "        ax.scatter( # your code here, use the next line setting for better visualization\n",
    "                        s=70, edgecolor='k',color = t_map[a],label='start' if a==targets[0] else None)\n",
    "        # plot mov. onset points\n",
    "        ax.scatter( # your code here, use the next line setting for better visualization\n",
    "                        marker = '*', s=70, edgecolor='k',color = t_map[a],label='mov. onset' if a==targets[0] else None)\n",
    "        # plot stop points\n",
    "        ax.scatter( # your code here, use the next line setting for better visualization\n",
    "                        marker = 'v', s=70, edgecolor='k',color = t_map[a],label='stop' if a==targets[0] else None)\n",
    "\n",
    "    ax.view_init(elev=elev, azim=azim)\n",
    "    ax.set_xlabel(f'PC{principal_components[0]+1}')\n",
    "    ax.set_ylabel(f'PC{principal_components[1]+1}')\n",
    "    ax.set_zlabel(f'PC{principal_components[2]+1}')\n",
    "    if k == 1:\n",
    "        ax.legend()\n",
    "    \n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8l71v-QDWD2L"
   },
   "source": [
    "* OK, we start seeing something! What do you observe? try to connect these plots with the behavioral data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Restricting the window"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you already know, PCA is an unsupervised dimensionality reduction method that just try to find axis (aka components) that maximize variance. Although very useful, sometimes it might not be what we want.\n",
    "\n",
    "Indeed, this might not lead to anything useful if we try to catch some specific information on the task, like target position/direction.\n",
    "\n",
    "In order to help PCA finding some better representation, let's try to restrict the fitting interval to a time window in which it is more likely to catch the info.\n",
    "\n",
    "Below, you will find some code that does exactly that. Since it can be confusing, we already prepared the code for you. ___The only thing you need to do is to find the window!___"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def projection_on_time_window(data, start_time, stop_time, n_pcs=4):\n",
    "    \"\"\"\n",
    "    Please enter start/stop times in milliseconds\n",
    "    \"\"\"\n",
    "    assert data.shape[2] == n_neurons, \"The data should be in the form ( n_trials, timesteps, n_neurons)\"\n",
    "    start = np.where(time_axis==start_time)[0][0]\n",
    "    stop = np.where(time_axis==stop_time)[0][0]\n",
    "\n",
    "    onset_data = norm_rate_neural_data[:,start:stop]\n",
    "    proj_onset_data, V_onset = compute_pca_unfolded(onset_data, n_pcs)\n",
    "    proj_full_data = norm_rate_neural_data @ V_onset.T\n",
    "\n",
    "    angles = trial_data.angle.values\n",
    "    times = (0,20)\n",
    "\n",
    "    fig, ax = plt.subplots(1,3,figsize=(13,4), constrained_layout=True)\n",
    "    for j,pcs in enumerate([(0,1),(1,2),(0,2)]):\n",
    "        for current_angle in targets:\n",
    "            trajectory = proj_full_data[angles==current_angle].mean(axis=0)\n",
    "            ax[j].plot(trajectory[start:stop, pcs[0]], trajectory[start:stop, pcs[1]],'-o',\n",
    "                    color = t_map[current_angle], linewidth = 1,markersize=2, zorder=1)\n",
    "            ax[j].scatter(trajectory[stop, pcs[0]], trajectory[stop, pcs[1]],marker='o',\n",
    "                    color = t_map[current_angle], edgecolor='k',s=120, label='stop' if current_angle==targets[0] else None, zorder=2)\n",
    "            ax[j].scatter(trajectory[start, pcs[0]], trajectory[start, pcs[1]],marker='v',\n",
    "                        color = t_map[current_angle], edgecolor='k',s=120,label='start' if current_angle==targets[0] else None, zorder=2)\n",
    "        ax[j].set_xlabel(f'PC{pcs[0]+1}')\n",
    "        ax[j].set_ylabel(f'PC{pcs[1]+1}')\n",
    "        ax[j].set_title(f'Neuron-unfolded PCA {np.array(pcs)+1}')\n",
    "        if j == 1:\n",
    "            ax[j].legend()\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "    pcs = [0,1,2]\n",
    "    times = (0,380)\n",
    "    fig = plt.figure(figsize=(13,4), constrained_layout=True)\n",
    "    gs = fig.add_gridspec(1, 3)\n",
    "    fig.suptitle(f'Neuron unfolded PCA {np.array(pcs)+1}')\n",
    "    for k, view in enumerate([(30, -60),(30,45),(-50,-60)]):\n",
    "        ax = fig.add_subplot(gs[k],projection='3d')\n",
    "        elev, azim = view\n",
    "        for current_angle in targets: \n",
    "            trajectory = proj_full_data[angles==current_angle].mean(axis=0)\n",
    "            ax.plot(trajectory[times[0]:times[1], pcs[0]], trajectory[times[0]:times[1], pcs[1]], trajectory[times[0]:times[1], pcs[2]],'-o',\n",
    "                    color = t_map[current_angle], linewidth = 1,markersize=2, zorder=1)\n",
    "            ax.scatter(trajectory[-1, pcs[0]], trajectory[-1, pcs[1]], trajectory[-1, pcs[2]],marker='^',edgecolor='k',\n",
    "                        color = 'k', linewidth = 1,s=20,label='end' if current_angle==targets[0] else None, zorder=2)\n",
    "            ax.scatter(trajectory[0, pcs[0]], trajectory[0, pcs[1]], trajectory[0, pcs[2]],marker='v',edgecolor='k',\n",
    "                        color = 'k', linewidth = 1,s=20,label='start' if current_angle==targets[0] else None, zorder=2)\n",
    "            ax.scatter(trajectory[stop, pcs[0]], trajectory[stop, pcs[1]], trajectory[stop, pcs[2]],marker='o',edgecolor='k',\n",
    "                    color = t_map[current_angle],s=80, label='window stop' if current_angle==targets[0] else None, zorder=3)\n",
    "            ax.scatter(trajectory[start, pcs[0]], trajectory[start, pcs[1]], trajectory[stop, pcs[2]],marker='v',edgecolor='k',\n",
    "                        color = t_map[current_angle],s=80,label='window start' if current_angle==targets[0] else None, zorder=3)\n",
    "        ax.view_init(elev=elev, azim=azim)\n",
    "        ax.set_xlabel(f'PC{pcs[0]+1}')\n",
    "        ax.set_ylabel(f'PC{pcs[1]+1}')\n",
    "        ax.set_zlabel(f'PC{pcs[2]+1}')\n",
    "        if k == 1:\n",
    "            ax.legend()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Play with the values to see how the projections change\n",
    "\n",
    "start_time = -900\n",
    "stop_time = 700\n",
    "\n",
    "projection_on_time_window(norm_rate_neural_data, start_time, stop_time)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UXnkqWF24kcB"
   },
   "source": [
    "### Trial covariability with PCA\n",
    "\n",
    "Let's now have a look at trial covariability. \n",
    "* Apply PCA to the **trial-unfolded** neural data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fill in your code below\n",
    "\n",
    "# Let's first check that the dataset has the correct shape. Use the variables timesteps, n_trials, and n_neurons\n",
    "correct_shape = ( # your code here ) # tuple of three elements\n",
    "flag = norm_rate_neural_data.shape == correct_shape\n",
    "assert flag, \"make sure that norm_rate_neural_data has the correct shape.\"\n",
    "\n",
    "n_pcs = 20\n",
    "\n",
    "# apply PCA using the function compute_pca_unfolded\n",
    "rate_neural_data_scores, V_neuron = # your code here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* Plot cumulative variance explained as a function of the number of principal components.\n",
    "* Compare this with the equivalent neuron-unfolded plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "trial_var_explained = []\n",
    "for number_of_pcs in range(1,k+1):\n",
    "    # Fill in your code below\n",
    "    ###########################\n",
    "\n",
    "    # let's select the first number_of_pcs principal components from the scores\n",
    "    sub_scores = # your code here\n",
    "\n",
    "    # use V_neuron to reconstruct the data\n",
    "    reconstructed_data = # your code here\n",
    "\n",
    "    ###########################\n",
    "    \n",
    "    # compute and save the variance explained\n",
    "    mse = ((reconstructed_data - norm_rate_neural_data)**2).mean()\n",
    "    trial_var_explained.append(1-mse/norm_rate_neural_data.var())\n",
    "\n",
    "fig = plt.figure(figsize=(4,4), constrained_layout=True)\n",
    "ax = fig.add_subplot()\n",
    "\n",
    "ax.plot(np.arange(1, k+1), trial_var_explained, '-o', label='trial-unfolded')\n",
    "ax.plot(np.arange(1, k+1), neuron_var_explained, '-o', label='neuron-unfolded')\n",
    "ax.set_xlabel('Number of PCs')\n",
    "ax.set_ylabel('Variance Explained')\n",
    "ax.set_title('Neuron-unfolded PCA')\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RdJN4Yr5Xq-R"
   },
   "source": [
    "* What is missing to really assess how much of the data is captured by each unfolding?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "T7WzpgptZs_4"
   },
   "source": [
    "**Step 5.** Plot the trial PCs which, as a reminder, are of shape #trials. Additionally sort and color code them by angle."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pcs_to_plot = 4\n",
    "\n",
    "fig, axes = plt.subplots(1, pcs_to_plot, figsize=(3*pcs_to_plot,3), constrained_layout=True)\n",
    "\n",
    "angles = trial_data['angle']\n",
    "sorted_id = np.argsort(angles)\n",
    "sorted_angles = angles[sorted_id].values\n",
    "\n",
    "# We sort by trials to make the plot easier to read\n",
    "V_trial_sorted = V_trial[:, sorted_id]\n",
    "\n",
    "for i in range(pcs_to_plot):\n",
    "    axes[i].scatter(# your code here (leave the rest of the line as it is) , color = [t_map[a] for a in sorted_angles])\n",
    "    axes[i].set_xlabel('Trial (sorted)')\n",
    "    axes[i].set_title('PC'+str(i+1))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Q1wBImWubX8T"
   },
   "source": [
    "* What do you observe?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bi-cross validation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now try to compute the dimensionality of our data using bi-cross validation. Do you remember how to do it? For make things faster, we did the job for you ;)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# let's create again our neuron-unfolded data\n",
    "correct_shape = (n_trials, timesteps, n_neurons)\n",
    "#norm_rate_neural_data = norm_rate_neural_data.transpose(2,1,0)\n",
    "flag = norm_rate_neural_data.shape == correct_shape\n",
    "assert flag, \"make sure that norm_rate_neural_data has correct shape.\"\n",
    "\n",
    "neuron_unfolded_data = np.concatenate([d for d in norm_rate_neural_data])\n",
    "\n",
    "neur_indices_tr,neur_indices_te = train_test_split(np.arange(neuron_unfolded_data.shape[1]), train_size = .8, shuffle=True)\n",
    "samples_indices_tr,samples_indices_te = train_test_split(np.arange(neuron_unfolded_data.shape[0]), train_size = .8, shuffle=True)\n",
    "\n",
    "# Selecting the training and test data\n",
    "M_tr_all = neuron_unfolded_data[samples_indices_tr,:]\n",
    "M_te_all = neuron_unfolded_data[samples_indices_te,:]\n",
    "M_te_tr = M_te_all[:,neur_indices_tr]\n",
    "M_te_te = M_te_all[:,neur_indices_te]\n",
    "# enter here your guess for the rank\n",
    "your_guess = 8\n",
    "\n",
    "r_max = 30\n",
    "dims = np.arange(1,r_max+1)\n",
    "var_exp_tr = np.zeros((r_max+1))\n",
    "var_exp_te = np.zeros((r_max+1))\n",
    "\n",
    "# looping over all different ranks\n",
    "for r in tqdm(dims):\n",
    "\n",
    "    # Apply SVD on the training data\n",
    "    _, S, V = np.linalg.svd(M_tr_all, full_matrices=False)\n",
    "\n",
    "    # Now, we solve the scores with linear regression. Let's use 'u' for the train basis, \n",
    "    # 'z' for the scores and 'x' forthe data\n",
    "    u = V[:r,neur_indices_tr].T\n",
    "    x = M_te_tr.T\n",
    "    # Use linear regression ( u z = x) to solve the scores on test data\n",
    "    z = np.linalg.lstsq(u,x,rcond=None)[0] # scores\n",
    "\n",
    "    # Compute the prediction for training neurons\n",
    "    M_te_tr_predict = (u @ z).T\n",
    "\n",
    "    # Compute explained variance for te_tr\n",
    "    mse = ((M_te_tr_predict - M_te_tr)**2).mean()\n",
    "    var_exp_tr[r] = 1-mse/M_te_tr.var()\n",
    "\n",
    "    # Compute the prediction on test neurons\n",
    "    M_te_te_predict = (V[:r,neur_indices_te].T @ z).T\n",
    "\n",
    "    # Compute variance explained for te_te\n",
    "    mse = ((M_te_te_predict - M_te_te)**2).mean()\n",
    "    var_exp_te[r] = 1-mse/M_te_te.var()\n",
    "\n",
    "fig = plt.figure(figsize=(4,3), constrained_layout=True)\n",
    "ax = fig.add_subplot()\n",
    "ax.plot(np.arange(r_max+1),var_exp_tr,'-o', label = \"train data\")\n",
    "ax.plot(np.arange(r_max+1),var_exp_te,'-o', label = \"test data\")\n",
    "ax.legend()\n",
    "ax.set_xlabel(\"Rank\")\n",
    "ax.set_ylabel(\"Variance explained\")\n",
    "ax.set_ylim([0,1.05])\n",
    "plt.show()\n",
    "\n",
    "k_guess = dims[np.argmax(var_exp_te)]\n",
    "\n",
    "print(\"Computed rank: \", k_guess)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* What do you observe? \n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# **Part V.** Supervised Learning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Until know, we have been focusing on a pure unsupervised methods, i.e. methods that allow to reduce the dimensionality of the dataset without any labeling telling us how good the obtained representation is.<br>\n",
    "\n",
    "Apart from pure visualization (and fun), our initial goal was to try connecting neural activity with behavioral observation. So, let's now try to use some (simple) supervised learning method for improving our understanding on how PMd/M1 could drive arm trajectories.\n",
    "\n",
    " We will fit our model on the neural data and on a PCA projection of the data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Some useful functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "from sklearn.linear_model import Ridge\n",
    "from sklearn.model_selection import GridSearchCV\n",
    "from sklearn.preprocessing import StandardScaler, MinMaxScaler\n",
    "from sklearn.model_selection import train_test_split\n",
    "import pickle\n",
    "from tqdm import tqdm\n",
    "\n",
    "import scipy.ndimage as spnd\n",
    "\n",
    "def plot_predictions(model, X, y, time_axis, variable = 'position'):\n",
    "    assert variable in ['position','velocity'], \"Please specify 'position' or 'velocity'\"\n",
    "    predictions = np.stack([model.predict(X[i]) for i in range(X.shape[0])])\n",
    "    fig, axs = plt.subplots(2, 3, figsize=(14, 5))\n",
    "    trials = X.shape[0]\n",
    "    colors = matplotlib.cm.rainbow(np.linspace(0,1,trials))\n",
    "    color_map = {trial_id: colors[i] for i, trial_id in enumerate(range(trials))}\n",
    "    for i in range(trials):\n",
    "        # X component\n",
    "        axs[0][0].plot(time_axis, y[i,:,0], linewidth=0.7, color=color_map[i])\n",
    "        axs[1][0].plot(time_axis, predictions[i,:,0], linewidth=0.7, color=color_map[i])\n",
    "        # Y component\n",
    "        axs[0][1].plot(time_axis, y[i,:,1], linewidth=0.7, color=color_map[i])\n",
    "        axs[1][1].plot(time_axis, predictions[i,:,1], linewidth=0.7, color=color_map[i])\n",
    "\n",
    "        # True and predicted trajectories\n",
    "        if variable == 'velocity':\n",
    "            true_traj = np.cumsum(y[i], axis=0) * bin_size / 1000\n",
    "            pred_traj = np.cumsum(predictions[i], axis=0) * bin_size / 1000\n",
    "        else:\n",
    "            true_traj = y[i]\n",
    "            pred_traj = predictions[i]\n",
    "        axs[0][2].plot(true_traj[:, 0], true_traj[:, 1], linewidth=0.7, color=color_map[i])\n",
    "        axs[1][2].plot(pred_traj[:, 0], pred_traj[:, 1], linewidth=0.7, color=color_map[i])\n",
    "    \n",
    "    # Set up shared axes\n",
    "    for i in range(2):\n",
    "        axs[i][0].set_xlim(time_axis[0], time_axis[-1])\n",
    "        axs[i][1].set_xlim(time_axis[0], time_axis[-1])\n",
    "        axs[i][2].set_xlim(-180, 180)\n",
    "        axs[i][2].set_ylim(-140, 140)\n",
    "\n",
    "    # Add labels\n",
    "    axs[0][0].set_title(f'X {variable} (mm/s)')\n",
    "    axs[0][1].set_title(f'Y {variable} (mm/s)')\n",
    "    axs[0][2].set_title('Reach trajectory')\n",
    "    axs[0][0].set_ylabel('Target')\n",
    "    axs[1][0].set_ylabel('Prediction')\n",
    "    plt.show()\n",
    "\n",
    "def compute_pca_unfolded(data,n_pcs=None):\n",
    "    \"\"\"\n",
    "    Function computing pca using SVD (np.linalg.svd).\n",
    "    The last dimension (data.shape[2]) will be the unfolded one.\n",
    "    If n_pcs is None, n_pcs = all components\n",
    "       Returns:\n",
    "       - the first n_pcs principal components of the data \n",
    "       - the eigenvector matrix\n",
    "    \"\"\"\n",
    "    assert data.ndim==3, \"Data must be 3 dimensional\"\n",
    "    assert isinstance(n_pcs,int) or n_pcs is None, \"If provided, n_pcs must be an integer\"\n",
    "\n",
    "    if n_pcs is None:\n",
    "        n_pcs = data.shape[2]\n",
    "    \n",
    "    # unfold the data. The final shape will be (dim1*dim2, dim3)\n",
    "    unfolded_data = np.concatenate([d for d in data])\n",
    "\n",
    "    # SVD\n",
    "    _, _, V = np.linalg.svd(unfolded_data, full_matrices=False)\n",
    "\n",
    "    # select the first n_pcs components\n",
    "    if n_pcs is not None:\n",
    "        V = V[:n_pcs]\n",
    "\n",
    "    # project the data on the \n",
    "    proj_data = data @ V.T\n",
    "\n",
    "    return proj_data, V\n",
    "\n",
    "def train_model(x_train, y_train, x_test, y_test, name):\n",
    "\n",
    "    print(\"Training '%s'\" %name)\n",
    "    # Fit decoder\n",
    "    gscv = GridSearchCV(Ridge(), {'alpha': np.logspace(-4, 0, 5)})\n",
    "    gscv.fit(x_train, y_train)\n",
    "\n",
    "    # evaluating decoder\n",
    "    train_score = gscv.score(x_train, y_train)\n",
    "    test_score = gscv.score(x_test, y_test)\n",
    "\n",
    "    \n",
    "    print(f\" Decoding R2 (train): {train_score}\")\n",
    "    print(f\" Decoding R2 (test): {test_score}\")\n",
    "\n",
    "    return gscv.best_estimator_\n",
    "\n",
    "def run_regression(input_dataset, target_dataset, n_pca_components, time_lag, time_window, selected_target, y_variable = 'position'):\n",
    "\n",
    "    bin_lag = time_lag//bin_size\n",
    "    new_time_axis = np.arange(time_window[0],time_window[1],bin_size)\n",
    "    bin_window = [np.where(time_axis==val)[0][0] for val in time_window]\n",
    "\n",
    "\n",
    "    targets = np.sort(trial_data['angle'].unique())\n",
    "    selected_target = targets[3]\n",
    "    angles = trial_data['angle'].values\n",
    "    \n",
    "\n",
    "    # z-scoring the dataset (using standard scaler this time)\n",
    "\n",
    "    assert input_dataset.shape == (n_trials, timesteps, n_neurons), f\"Please check the shape of the dataset. Correct = {(n_trials, timesteps, n_neurons)}\"\n",
    "    input_dataset = np.stack([StandardScaler().fit_transform(input_dataset[:,:,n]) for n in range(n_neurons)],axis=-1)\n",
    "\n",
    "    # Input datasets (neural data)\n",
    "\n",
    "    ## Apply time window\n",
    "    input_dataset = input_dataset[:,bin_window[0]:bin_window[1]]\n",
    "\n",
    "    ## PCA-transformed dataset\n",
    "    assert input_dataset.shape == (n_trials, bin_window[1]-bin_window[0], n_neurons), \"Please check the shape of the dataset\"\n",
    "    rate_dataset_pca, _ = compute_pca_unfolded(input_dataset, n_pca_components)\n",
    "\n",
    "\n",
    "    rate_datasets = [input_dataset,rate_dataset_pca]\n",
    "    rate_datasets_names = [\"rate dataset\",\"pca-reduced dataset\"]\n",
    "\n",
    "    # Target datasets\n",
    "\n",
    "    ## define the lagged window\n",
    "    window = (bin_window[0] + bin_lag, bin_window[1] + bin_lag)\n",
    "    ## select correct time window and variables\n",
    "    targets_dataset = target_dataset[:,:, window[0]:window[1]].transpose(0,2,1)\n",
    "\n",
    "\n",
    "    # Creating train/test datasets\n",
    "\n",
    "    train_indeces, test_indeces = train_test_split(np.arange(targets_dataset.shape[0]), train_size = .7, shuffle=True)\n",
    "\n",
    "    train_rate_datasets = [dataset[train_indeces] for dataset in rate_datasets]\n",
    "    test_rate_datasets = [dataset[test_indeces] for dataset in rate_datasets]\n",
    "\n",
    "    y_train = targets_dataset[train_indeces]\n",
    "    y_test = targets_dataset[test_indeces]\n",
    "\n",
    "    # making the datasets 2D (samples x n. features)\n",
    "    train_datasets = [np.concatenate([data for data in dataset]) for dataset in train_rate_datasets]\n",
    "    test_datasets = [np.concatenate([data for data in dataset]) for dataset in test_rate_datasets]\n",
    "\n",
    "\n",
    "    y_train = np.concatenate([data for data in y_train])\n",
    "    y_test = np.concatenate([data for data in y_test])\n",
    "\n",
    "\n",
    "    # Preparing plot dataset\n",
    "\n",
    "    # angles_train = trial_data_straight.loc[train_indeces, 'angles'].values\n",
    "    angles_test = angles[test_indeces]\n",
    "\n",
    "    # train_index = (angles_train == selected_target)\n",
    "    test_selection_index = (angles_test == selected_target)\n",
    "\n",
    "    # train_selection = [dataset[train_index] for dataset in train_rate_datasets]\n",
    "    test_selection_X = [dataset[test_selection_index] for dataset in test_rate_datasets]\n",
    "    test_selection_y = targets_dataset[test_indeces][test_selection_index]\n",
    "\n",
    "    for i, name in enumerate(rate_datasets_names):\n",
    "        model = train_model(train_datasets[i], y_train, test_datasets[i], y_test, name)\n",
    "\n",
    "        plot_predictions(model, test_selection_X[i], test_selection_y, new_time_axis, y_variable)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hand position decoder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we fit a linear decoder (more precisely, we will perform a Ridge regression) on the neural activity, using it as a decoder for ___predicting hand position___. \n",
    "\n",
    "Since coding might take a while, we prepared a function called `run regression`, that implement the regression. What you are asked to do is to set the data and the parameters that are used in the regression:\n",
    "- `input_dataset`: the neural activity\n",
    "- `target_dataset`: the behavior dataset (positions)\n",
    "- `n_pca_components`\n",
    "- `time_lag`: allows the model to take into account brain-to-movement delays\n",
    "- `time_window`: selects a particular part of the data for the regression\n",
    "\n",
    "Try to play with the parameters in order to make the model working better."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Tunable parameters\n",
    "X_dataset = #\n",
    "y_dataset = #\n",
    "selected_target = targets[2]\n",
    "n_pca_components = 3\n",
    "time_lag = 0 # (ms) Taking into account the delay between cortical activation and the movement\n",
    "time_window = (-500, 500) # (ms) selecting a time window\n",
    "\n",
    "\n",
    "# Run the regression\n",
    "run_regression(X_dataset, y_dataset, n_pca_components, time_lag, time_window, selected_target, 'position')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hand velocity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, try to do the same for the hand velocity instead, and see if there is any improvement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Tunable parameters\n",
    "X_dataset = #\n",
    "y_dataset = #\n",
    "selected_target = targets[2]\n",
    "n_pca_components = 3\n",
    "time_lag = 0 # (ms) Taking into account the delay between cortical activation and the movement\n",
    "time_window = (-500, 500) # (ms) selecting a time window\n",
    "\n",
    "\n",
    "# Run the regression\n",
    "run_regression(X_dataset, y_dataset, n_pca_components, time_lag, time_window, selected_target, 'velocity')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* A little better (maybe). What do you think is happening? Can we do better?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CP7VEBVqy-GN"
   },
   "source": [
    "# References\n",
    "\n",
    "[1] Cunningham, J. P., & Yu, B. M. (2014). Dimensionality reduction for large-scale neural recordings. Nature neuroscience, 17(11), 1500-1509.\n",
    "\n",
    "[2] Bishop, C. M. (2009). Pattern recognition and machine learning. Springer, New York.\n",
    "\n",
    "[3] Seely, J. S., Kaufman, M. T., Ryu, S. I., Shenoy, K. V., Cunningham, J. P., & Churchland, M. M. (2016). Tensor analysis reveals distinct population structure that parallels the different computational roles of areas M1 and V1. PLoS computational biology, 12(11), e1005164.\n",
    "\n",
    "[4] Williams, A. H., Kim, T. H., Wang, F., Vyas, S., Ryu, S. I., Shenoy, K. V., ... & Ganguli, S. (2018). Unsupervised discovery of demixed, low-dimensional neural dynamics across multiple timescales through tensor component analysis. Neuron, 98(6), 1099-1115.\n",
    "\n",
    "[5] Pellegrino*, A., Stein*, H., & Cayco-Gajic, N. A. (2023). Disentangling mixed classes of covariability in large-scale neural data. bioRxiv.\n",
    "\n",
    "[6] Owen, A. B., & Perry, P. O. (2009). Bi-cross-validation of the SVD and the nonnegative matrix factorization. Annals of Applied Statistics, 3(2), 564-594.\n",
    "\n",
    "[7] Mark M. Churchland, John P. Cunningham, Matthew T. Kaufman, Stephen I. Ryu, Krishna V. Shenoy (2010).\n",
    "Cortical Preparatory Activity: Representation of Movement or First Cog in a Dynamical Machine?. Neuron, Volume 68, Issue 3\n",
    "\n",
    "[8] Churchland, M., Cunningham, J., Kaufman, M. et al. (2012). Neural population dynamics during reaching. Nature 487, 51–56.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [
    "-khOJ_4KHdhn"
   ],
   "gpuType": "T4",
   "machine_shape": "hm",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": true,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}