2A - Discrete Shape Analysis.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d35c3df7",
   "metadata": {},
   "source": [
    "## Part 2A: Discrete shape analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27cdcf3b",
   "metadata": {},
   "source": [
    "In this notebook, we will learn the basics of shape analysis on landmarked data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8669f5ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.cm as cm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba34c4a0",
   "metadata": {},
   "source": [
    "### 0. Data loading"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc55ccc8",
   "metadata": {},
   "source": [
    "**0.1** Run the lines below to load point-based object models of objects in the C. elegans dataset.\n",
    "\n",
    "*The C. elegans dataset is presented and prepared in notebook 1 - Data preparation.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0ffe14f",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset='data/C. elegans/point_models.npy'\n",
    "point_models=np.load(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "deecd98a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The number of points in our object model is an important parameter\n",
    "N=point_models.shape[1]\n",
    "print(\"N=\"+str(N))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c9d2cdb",
   "metadata": {},
   "source": [
    "**0.2** Run the lines below to visualize a set of a few randomly selected models from the collection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e91ff6f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "number=10\n",
    "inds=np.random.choice(len(point_models)-1, size=number, replace=False)\n",
    "\n",
    "fig, ax = plt.subplots(1,number, figsize=(number,1))\n",
    "for i in range(number):\n",
    "    point_model=np.array(point_models[inds[i]])\n",
    "    ax[i].scatter(point_model[:,0],point_model[:,1],s=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ed49c8c",
   "metadata": {},
   "source": [
    "### 1. Procrustes alignment"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b314d500",
   "metadata": {},
   "source": [
    "**1.1** Center and scale the data such that they are all located around (0,0) and are of unit norm. Object models normalized in this way are referred to as preshapes.\n",
    "\n",
    "*Hint: use ``np.linalg.norm`` to normalize*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8b33f6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Desired output: preshapes array of the same dimensions as point_model\n",
    "# Add your code here!\n",
    "preshapes="
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc4b1c29",
   "metadata": {},
   "source": [
    "**1.2** Write an objective function to align one object model onto another one, assuming that the two models are in correspondence†. This problem can be formulated as finding the rotation angle that minimizes the root mean square error (RMSE) between the two point sets.\n",
    "\n",
    "*Hint 1: assuming the distance between two models is ``L``, the RMSE is given as ``np.sqrt(np.mean(L**2))``.*\n",
    "\n",
    "*Hint 2: rotating a point set can be done by multiplying with a 2D rotation matrix (https://en.wikipedia.org/wiki/Rotation_matrix).*\n",
    "\n",
    "† Bonus question: do you understand what this assumption means? Which step must be performed when building the object models to ensure that it holds?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce6f4b17",
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.optimize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "301304f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Picks two random models - reload if you are not happy with the selection :)\n",
    "random=np.random.choice(len(point_models)-1, size=2, replace=False)\n",
    "\n",
    "sample_preshape_1=preshapes[random[0]]\n",
    "sample_preshape_2=preshapes[random[1]]\n",
    "\n",
    "plt.scatter(sample_preshape_1[:,0],sample_preshape_1[:,1])\n",
    "plt.scatter(sample_preshape_2[:,0],sample_preshape_2[:,1])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a09bb0fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Desired output: objectiveFunctionAlignment function returning the RMSE between a reference model m1 and \n",
    "# another model m2 rotated by an angle theta\n",
    "# Add your code here!\n",
    "def objectiveFunctionAlignment(theta, m1, m2):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e7c5ea4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optimizes the objective function\n",
    "res=scipy.optimize.minimize(objectiveFunctionAlignment,0,args=(sample_preshape_1,sample_preshape_2),method=\"L-BFGS-B\")\n",
    "theta0=res.x[0]\n",
    "print(theta0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88b54167",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualizes the alignment\n",
    "plt.scatter(sample_preshape_1[:,0],sample_preshape_1[:,1])\n",
    "\n",
    "R=np.array([[np.cos(theta0),-np.sin(theta0)],[np.sin(theta0),np.cos(theta0)]])\n",
    "aligned=(R @ sample_preshape_2.T).T\n",
    "\n",
    "plt.scatter(aligned[:,0],aligned[:,1])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "299cca28",
   "metadata": {},
   "source": [
    "**1.3** Object models in 2D can be seen as vectors of complex numbers associated with an appropriate Hermitian inner product. The lines below transforms your two objects into complex vectors. Run the lines below to convince yourself that the \"vector of complex numbers\" representation is equivalent to the \"vector of points\" one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "236a23c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "complex_sample_preshape_1=sample_preshape_1[:,0]+1j*sample_preshape_1[:,1]\n",
    "complex_sample_preshape_2=sample_preshape_2[:,0]+1j*sample_preshape_2[:,1]\n",
    "\n",
    "plt.scatter(complex_sample_preshape_1.real,complex_sample_preshape_1.imag)\n",
    "plt.scatter(complex_sample_preshape_2.real,complex_sample_preshape_2.imag)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f15a078",
   "metadata": {},
   "source": [
    "**1.4** The solution of the optimal alignment problem 1.2 can actually be expressed in closed-form when relying on complex representations, bypassing the need for an optimization. Run the lines below to verify that you obtain the same solution as in 1.2 (and be amazed by the beauty of maths :))."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e063e94",
   "metadata": {},
   "outputs": [],
   "source": [
    "theta0_complex=np.angle(complex_sample_preshape_2.conj().T @ complex_sample_preshape_1)\n",
    "print(theta0_complex)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6fe5724",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualizes the alignment \n",
    "plt.scatter(complex_sample_preshape_1.real,complex_sample_preshape_1.imag)\n",
    "\n",
    "aligned_complex=np.exp(1j*theta0_complex)*complex_sample_preshape_2\n",
    "\n",
    "plt.scatter(aligned_complex.real,aligned_complex.imag)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41d1bbfc",
   "metadata": {},
   "source": [
    "*In case you are wondering what the hell happened in the lines above, here are more details (if you don't care, you can safely skip to 2 below).*\n",
    "\n",
    "To align the two objects, we must retreive the angle between them. If what we were dealing with was 2D vectors, we would get that angle immediately through the scalar product (https://en.wikipedia.org/wiki/Dot_product). An equivalent of the scalar product can be defined for complex numbers (usually referred to as Hermitian inner product, https://en.wikipedia.org/wiki/Inner_product_space), so we could equivalently retreive that angle this way as demonstrated in the code below.\n",
    "\n",
    "Now it happens that we are not working with 2D vectors (or complex numbers), but with Nx2 matrices (or complex vectors)! Fortunately, the Hermitian inner product adapt straighforwardly to complex vectors, leading to the expression you see above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c778667e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define two arbitrary 2D vectors v1 and v2 and normalize them\n",
    "v1=[10,7]\n",
    "v1/=np.linalg.norm(v1)\n",
    "\n",
    "v2=[5,5]\n",
    "v2/=np.linalg.norm(v2)\n",
    "\n",
    "# Compute the angle between v1 and v2 using the definition of the scalar product\n",
    "angle_real=np.arccos(np.dot(v1,v2))\n",
    "\n",
    "# Turn v1 and v2 into z1 and z2, their complex representation\n",
    "z1=v1[0]+1j*v1[1]\n",
    "z2=v2[0]+1j*v2[1]\n",
    "\n",
    "# Compute the angle between z1 and z2 using the definition of the Hermitian (complex) inner product\n",
    "angle_complex=np.angle(z1.conj()*z2)\n",
    "\n",
    "# Display the results\n",
    "print(\"Angle computed with the real inner product: \"+str(angle_real))\n",
    "print(\"Angle computed with the complex inner product: \"+str(angle_complex))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba4150db",
   "metadata": {},
   "source": [
    "### 2. Kendall shape space"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9807ad97",
   "metadata": {},
   "source": [
    "**2.1** To build the shape space, we will rely on the complex representation introduced in 1.3. Convert your whole collection of preshapes (obtained in 1.1) into complex numbers. \n",
    "\n",
    "*Important note: Make sure that complex_preshapes is a numpy array!*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9530bc3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Desired output: complex_preshapes array of the same length as preshapes\n",
    "# Add your code here!\n",
    "complex_preshapes="
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a493cdb8",
   "metadata": {},
   "source": [
    "**2.2** The function below computes the Fréchet mean of a dataset relying on complex number representations. Use it to extract the Fréchet mean of your dataset and visualize it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf73ba6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def meanFrechet(input_complex_preshapes): # preshapes\n",
    "    '''Input: dataset of complex preshapes.\n",
    "    Output: Fréchet mean (w.r.t. the Procrustes distance).'''\n",
    "    SQ = input_complex_preshapes.T @ input_complex_preshapes.conj()\n",
    "    D,V = np.linalg.eig(SQ)\n",
    "    ds = np.real(D)\n",
    "    ind_max = np.argmax(ds)\n",
    "    ds_max = np.max(ds)\n",
    "    m = V[:,ind_max]\n",
    "    \n",
    "    centered_m=m-np.mean(m)\n",
    "    m_norm=np.sqrt(centered_m.conj().T @ centered_m).real\n",
    "    return m/m_norm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "902be047",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Desired output: mean_shape_Frechet array\n",
    "# Add your code here!\n",
    "mean_shape_Frechet="
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2026175",
   "metadata": {},
   "source": [
    "*In case you are wondering what the hell happened in the lines above, here are more details (if you don't care, you can safely skip to 2.3 below).*\n",
    "\n",
    "It is not straightforward to see how the minimization problem one needs to solve to obtain the Fréchet mean transforms into an eigenvalue decomposition as implemented in the function meanFrechet above. Those interested in more details and a proof are invited to look at Result 8.2 (Section 8.3) of the Dryden & Mardia textbook (ISBN 978-0471958161)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c58462d8",
   "metadata": {},
   "source": [
    "**2.3** Using the alignment procedure introduced in 1.4, align each element of the dataset onto the Fréchet mean (computed in 2.2)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f64a0820",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Desired output: complex_aligned_shapes array of the same dimensions as complex_preshapes\n",
    "# Add your code here!\n",
    "complex_aligned_shapes="
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abd372bb",
   "metadata": {},
   "source": [
    "**2.4** Now that all shapes are aligned, compute the (point-by-point) mean shape of the dataset and visualize it. Does it differ from what you got in 2.2?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc954ac5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Desired output: mean_shape array and its scatter plot\n",
    "# Add your code here!\n",
    "mean_shape="
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35971655",
   "metadata": {},
   "source": [
    "**2.5** The code below and retreives the mean shape of the dataset relying on general Procrustes alignment. Run it a few times: what do you observe? Does the mean shape obtained in this way differ from what you got in 2.4? \n",
    "\n",
    "*As a reminder, general Procrustes alignment  the following iterative procedure:*\n",
    "1. *Randomly select one object from the collection*\n",
    "2. *Align every other objects to it (ordinary Procrustes alignment)*\n",
    "3. *Compute the average of all aligned data points and scale it to unit norm*\n",
    "4. *Align all objects to that new average*\n",
    "5. *Compute the average of all aligned data (not considering the average obtained at step 3 as this is not part of the data!)*\n",
    "6. *Repeat steps 4 and 5 until convergence*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0aaeb31",
   "metadata": {},
   "outputs": [],
   "source": [
    "from copy import deepcopy\n",
    "\n",
    "# 1. Randomly select one object from the collection\n",
    "random=np.random.randint(len(complex_preshapes)-1)\n",
    "complex_reference=complex_preshapes[random]\n",
    "\n",
    "# 2. Align every other objects to it (ordinary Procrustes alignment)\n",
    "aligned_complex=np.zeros(complex_preshapes.shape, dtype=complex)\n",
    "for i in range(len(complex_preshapes)):\n",
    "    aligned_complex[i]=(complex_preshapes[i].conj().T @ complex_reference)*complex_preshapes[i]\n",
    "    \n",
    "# 3. Compute the average of all aligned data points and scale it to unit norm\n",
    "mean_shape_2=np.mean(aligned_complex,0)\n",
    "mean_shape_2/=np.linalg.norm(mean_shape_2)\n",
    "\n",
    "epsilon = 1e-3\n",
    "delta = np.Inf\n",
    "while np.abs(delta) > epsilon:\n",
    "    # 4. Align all objects to that new average\n",
    "    aligned_complex=np.zeros(complex_preshapes.shape, dtype=complex)\n",
    "    for i in range(len(complex_preshapes)):\n",
    "        aligned_complex[i]=(complex_preshapes[i].conj().T @ mean_shape_2)*complex_preshapes[i]\n",
    "\n",
    "    # 5. Compute the average of all aligned data (not considering the average obtained at step 3 as this is not part of the data!)\n",
    "    new_mean_shape_2=np.mean(aligned_complex,0)\n",
    "    new_mean_shape_2/=np.linalg.norm(new_mean_shape_2)\n",
    "\n",
    "    # 6. Repeat steps 4 and 5 until convergence\n",
    "    delta=np.linalg.norm(new_mean_shape_2 - mean_shape_2)\n",
    "    mean_shape_2=deepcopy(new_mean_shape_2)\n",
    "\n",
    "# Visualize the result\n",
    "plt.scatter(mean_shape_2.real,mean_shape_2.imag)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "404b3eeb",
   "metadata": {},
   "source": [
    "**2.6** Relying on the function ``geodesicPath`` below, compute the geodesic path between a randomly picked shape in the dataset and the Fréchet mean, and visualize elements along the path."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d6dae38",
   "metadata": {},
   "outputs": [],
   "source": [
    "def geodesicDistance(z,w): \n",
    "    '''Geodesic distance between [z] and [w], but computed on preshapes.'''\n",
    "    aux=np.abs(z.conj().T @ w)\n",
    "    if aux>1.0: \n",
    "        aux=1.0\n",
    "    return np.arccos(aux)\n",
    "\n",
    "def geodesicPath(z,w,numSteps): \n",
    "    '''Returns elements regularly spaced along the geodesic curve joining z to w (preshapes).'''\n",
    "    ro = geodesicDistance(z,w)\n",
    "    steps = np.arange(numSteps+1)/numSteps\n",
    "\n",
    "    ta = np.angle(z.conj().T @ w)\n",
    "    path = 1/np.sin(ro)*(np.sin((1-steps[:,None])*ro)*np.exp(1j*ta)*z + np.sin(steps[:,None]*ro)*w)\n",
    "    return path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ac81cde",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Picks one random complex preshape and visualize it with the Fréchet mean - reload if you are not happy with the selection :)\n",
    "random=np.random.choice(len(complex_preshapes)-1, size=1, replace=False)[0]\n",
    "sample_preshape=complex_preshapes[random]\n",
    "\n",
    "plt.scatter(mean_shape_Frechet.real,mean_shape_Frechet.imag)\n",
    "plt.scatter(sample_preshape.real,sample_preshape].imag)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cb0ace1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Desired output: path array containing steps number of object models corresponding to shapes on the geodesic path\n",
    "# between mean_shape_Frechet and sample_preshape, and scatter plots visualizing them\n",
    "# Add your code here!\n",
    "steps=\n",
    "path="
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca215afd",
   "metadata": {},
   "source": [
    "**2.7** Compare the length of the above geodesics with the Procrustes distance (the distance you used for the alignment in 1.2). What do you observe? Can you relate that to what you know about the nature of the shape space?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a39d14b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Desired output: geodesic distance between sample_preshape and mean_shape_Frechet, and Euclidean distance between\n",
    "# complex_aligned_shapes[random] and mean_shape_Frechet.\n",
    "# Add your code here!\n",
    "geo_distance=\n",
    "print(\"Geodesic distance: \"+str(geo_distance))\n",
    "\n",
    "procrustes_dist=\n",
    "print(\"Procrustes distance: \"+str(procrustes_dist))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26231aae",
   "metadata": {},
   "source": [
    "**2.8** Compare the Procrustes distance (the distance you used for the alignment in 1.2) with the chord distance on the shape space defined below. What do you observe?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc7e4048",
   "metadata": {},
   "outputs": [],
   "source": [
    "def chordDistance(z,w): # defined on preshapes\n",
    "    '''Chord distance between [z] and [w].'''\n",
    "    return np.sqrt(2.0 - 2.0*np.abs(z.conj().T @ w))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a4a4e61",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Desired output: chord distance between sample_preshape and mean_shape_Frechet, and Euclidean distance between\n",
    "# complex_aligned_shapes[random] and mean_shape_Frechet.\n",
    "# Add your code here!\n",
    "chord_distance=\n",
    "print(\"Chord distance: \"+str(chord_distance))\n",
    "print(\"Procrustes distance: \"+str(procrustes_dist))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98a558b1",
   "metadata": {},
   "source": [
    "### 3. Shape space PCA on the whole dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6d96c0d",
   "metadata": {},
   "source": [
    "A classical trick to perform local linear operations on a Riemannian manifold is to rely on local tangent spaces. The way into and out of the tangent space are the logarithm and exponential maps, respectively. Their expression for the shape space of point-based models we are currently working with are provided below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96de15ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "def logarithmMap(z,w): \n",
    "    '''Computes a preshape pertaining to the equivalence class log_[z]([w]),\n",
    "    where log is relative the shape space Sigma'''\n",
    "    ta = np.angle(z.conj().T @ w)\n",
    "    w_r = np.exp(-1j*ta)*w\n",
    "    ro = geodesicDistance(z,w_r)\n",
    "    return ro/np.sin(ro)*(w_r - np.cos(ro)*z)\n",
    "\n",
    "def exponentialMap(z,v): \n",
    "    '''Computes the exponential of the tangent vector v in C^n at a preshape z'''\n",
    "    t = np.sqrt(v.conj().T @ v).real\n",
    "    if t < 1e-16 :\n",
    "        return z\n",
    "    return np.cos(t)*z + v*np.sin(t)/t"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b13deb34",
   "metadata": {},
   "source": [
    "**3.1** The lines below perform PCA in the tangent plane around the Fréchet mean using the logarithmic and exponential maps appropriately and displays the explained variance. What do you notice?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f8384f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.decomposition import PCA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0abaa7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "tangent_vectors=np.zeros((len(complex_preshapes),2*N))\n",
    "for i in range(len(complex_preshapes)):\n",
    "    tangent_vector=logarithmMap(mean_shape_Frechet, complex_preshapes[i])\n",
    "    tangent_vectors[i,:N]=tangent_vector.real\n",
    "    tangent_vectors[i,N:]=tangent_vector.imag\n",
    "\n",
    "tangent_plane_pca = PCA()\n",
    "transformed_vectors = tangent_plane_pca.fit_transform(tangent_vectors)\n",
    "\n",
    "plt.plot(np.linspace(1,tangent_plane_pca.n_components_,tangent_plane_pca.n_components_),100*tangent_plane_pca.explained_variance_/np.sum(tangent_plane_pca.explained_variance_))\n",
    "plt.xlabel(\"Principal component\")\n",
    "plt.ylabel(\"Variance explained\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57ed4dc7",
   "metadata": {},
   "source": [
    "**3.2** The lines below retreive the four first modes of shape variation and visualize them around the mean. What does it tell you about the shapes present in the data?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0233d2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "l=np.sqrt(tangent_plane_pca.explained_variance_)\n",
    "K=4\n",
    "\n",
    "modes_shape=np.zeros((K,2,N),dtype=np.complex)\n",
    "for i in range(K) :\n",
    "    vector=np.zeros(tangent_plane_pca.n_components_)\n",
    "    vector[i]=1\n",
    "    mode=l[i]*tangent_plane_pca.inverse_transform(vector)\n",
    "    complex_mode=mode[:N]+1j*mode[N:]\n",
    "    modes_shape[i][0]=exponentialMap(mean_shape_Frechet, complex_mode)\n",
    "    modes_shape[i][1]=exponentialMap(mean_shape_Frechet, -complex_mode)\n",
    "    \n",
    "fig, ax = plt.subplots(1, K, figsize=(5*K,5))\n",
    "for i in range(K):\n",
    "    ax[i].scatter(mean_shape_Frechet.real,mean_shape_Frechet.imag)\n",
    "    ax[i].scatter(modes_shape[i][0].real, modes_shape[i][0].imag)\n",
    "    ax[i].scatter(modes_shape[i][1].real, modes_shape[i][1].imag)\n",
    "    \n",
    "    percent_variance=np.round(100*l[i]**2/np.sum(tangent_plane_pca.explained_variance_))\n",
    "    ax[i].set_title(\"Mode \"+str(i+1)+\", \"+str(percent_variance)+\"% of variance.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "331b0980",
   "metadata": {},
   "source": [
    "**3.3** The lines below plot the PCA-transformed data around the Fréchet mean. Each point corresponds to an object. If z_m is the Fréchet mean and l_1, l_2 are the two first eigenvalues of the PCA, the red point corresponds to z_m, the blue points to z_m +/- l_1, and the cyan point to z_m +/- l_2.\n",
    "\n",
    "What does this plot tell you about the dataset?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ef34440",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.scatter(transformed_vectors[:,0], transformed_vectors[:,1], c=\"g\")\n",
    "plt.scatter([0],[0], c=\"r\")\n",
    "plt.scatter([l[0],-l[0]], [0.,0.],c=\"blue\")\n",
    "plt.scatter([0.,0.],[l[1],-l[1]],c=\"cyan\")\n",
    "plt.axis(\"equal\")\n",
    "plt.title(\"PCA plot around the Fréchet mean\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd239157",
   "metadata": {},
   "source": [
    "**[BONUS] 3.4** The lines below perform PCA as above but without using the logarithm and exponential maps. What do you see? Can you relate that to what you observed in 2.7-2.8?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f208a35",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Performs PCA assuming the shape space is linear\n",
    "flattened_aligned_shapes=np.zeros((len(complex_aligned_shapes),2*N))\n",
    "for i in range(len(complex_aligned_shapes)):\n",
    "    flattened_aligned_shapes[i,:N]=complex_aligned_shapes[i].real\n",
    "    flattened_aligned_shapes[i,N:]=complex_aligned_shapes[i].imag\n",
    "    \n",
    "flat_pca = PCA()\n",
    "flat_transformed_data = flat_pca.fit_transform(flattened_aligned_shapes)\n",
    "\n",
    "plt.plot(np.linspace(1,flat_pca.n_components_,flat_pca.n_components_),100*flat_pca.explained_variance_/np.sum(flat_pca.explained_variance_))\n",
    "plt.xlabel(\"Number of principal components\")\n",
    "plt.ylabel(\"Variance explained\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fe906c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualizes the first K modes of variation\n",
    "flat_l=np.sqrt(flat_pca.explained_variance_)\n",
    "K=4\n",
    "\n",
    "flat_modes_shape=np.zeros((K,N,2))\n",
    "for i in range(K) :\n",
    "    vector=np.zeros(flat_pca.n_components_)\n",
    "    vector[i]=1\n",
    "    flat_mode=0.25*flat_l[i]*flat_pca.inverse_transform(vector)\n",
    "    flat_modes_shape[i][:,0]=flat_mode[:N]\n",
    "    flat_modes_shape[i][:,1]=flat_mode[N:]\n",
    "    \n",
    "fig, ax = plt.subplots(1, K, figsize=(5*K,5))\n",
    "for i in range(K):\n",
    "    ax[i].scatter(mean_shape_Frechet.real,mean_shape_Frechet.imag)\n",
    "    ax[i].scatter(flat_modes_shape[i][:,0], flat_modes_shape[i][:,1])\n",
    "    ax[i].scatter(-flat_modes_shape[i][:,0], -flat_modes_shape[i][:,1])\n",
    "    \n",
    "    percent_variance=np.round(100*flat_l[i]**2/np.sum(flat_pca.explained_variance_))\n",
    "    ax[i].set_title(\"Mode \"+str(i+1)+\", \"+str(percent_variance)+\"% of variance.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b156c7b",
   "metadata": {},
   "source": [
    "### 4. Shape space PCA on classes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50101878",
   "metadata": {},
   "source": [
    "We actually know that the shape distribution is bimodal, because there are dead and alive worms. To have a clearer picture of shape variability, we should thus carry out shape PCA on each class individually."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6deea651",
   "metadata": {},
   "source": [
    "**4.1** Run the lines below to split the dataset into a collection of alive and a collection of dead C.elegans samples.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f796d125",
   "metadata": {},
   "outputs": [],
   "source": [
    "label_data='data/C. elegans/labels.npy'\n",
    "labels=np.load(label_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "effccca7",
   "metadata": {},
   "outputs": [],
   "source": [
    "complex_preshapes_dead=complex_preshapes[np.where(labels==0)]\n",
    "complex_preshapes_live=complex_preshapes[np.where(labels==1)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b88fc4fe",
   "metadata": {},
   "source": [
    "**4.2** Run the lines below to visualize a set of a few randomly selected models from each collection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71dae1a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dead C. elegans\n",
    "number=5\n",
    "inds=np.random.choice(len(complex_preshapes_dead)-1, size=number, replace=False)\n",
    "\n",
    "fig, ax = plt.subplots(1,number, figsize=(number,1))\n",
    "for i in range(number):\n",
    "    sample=complex_preshapes_dead[inds[i]]\n",
    "    ax[i].scatter(sample.real,sample.imag,s=1)\n",
    "    ax[i].set_title(inds[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b112ce2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Live C. elegans\n",
    "number=5\n",
    "inds=np.random.choice(len(complex_preshapes_live)-1, size=number, replace=False)\n",
    "\n",
    "fig, ax = plt.subplots(1,number, figsize=(number,1))\n",
    "for i in range(number):\n",
    "    sample=complex_preshapes_live[inds[i]]\n",
    "    ax[i].scatter(sample.real,sample.imag,s=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e23a3059",
   "metadata": {},
   "source": [
    "**4.3** Compute and visualize the Fréchet mean for each individual class, reusing the `meanFrechet` function from 2.2. How do they differ from what you got in 2.2?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80cde0fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Desired output: mean_shape_Frechet_dead and mean_shape_Frechet_live arrays, and scatter plots to visualize them\n",
    "# Add your code here!\n",
    "mean_shape_Frechet_dead=\n",
    "mean_shape_Frechet_live="
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef4dcd93",
   "metadata": {},
   "source": [
    "**4.4** For each class, perform PCA in the tangent plane around the Fréchet mean by adapting the code from 3.1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ccb9bd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Desired output: tangent plane PCA as in 3.1 for complex_preshapes_dead around mean_shape_Frechet_dead (tangent_plane_pca_dead) \n",
    "# and complex_preshapes_live around mean_shape_Frechet_live (tangent_plane_pca_live)\n",
    "# Add your code here!\n",
    "tangent_plane_pca_dead =\n",
    "tangent_plane_pca_live ="
   ]
  },
  {
   "cell_type": "markdown",
   "id": "765555aa",
   "metadata": {},
   "source": [
    "**4.5** For each class, retreive the two first modes of shape variation and visualize them around the mean by adapting the code from 3.2. How do they differ from what you got in 3.2?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7819d4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Desired output: mode visualisation as in 3.2 for tangent_plane_pca_dead (modes_shape_dead) and \n",
    "# tangent_plane_pca_alive (modes_shape_alive)\n",
    "# Add your code here!\n",
    "modes_shape_dead=\n",
    "modes_shape_live="
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3380832a",
   "metadata": {},
   "source": [
    "**4.6** The lines below plot the PCA-transformed data around the Fréchet mean for each class (see 3.3 for a detailed explanations of what the different coloured points correspond to). How does it compare to 3.3?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60b8f566",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1,2)\n",
    "\n",
    "ax[0].scatter(transformed_vectors_dead[:,0], transformed_vectors_dead[:,1], c=\"g\")\n",
    "ax[0].scatter([0],[0], c=\"r\")\n",
    "ax[0].scatter([l_dead[0],-l_dead[0]], [0.,0.],c=\"blue\")\n",
    "ax[0].scatter([0.,0.],[l_dead[1],-l_dead[1]],c=\"cyan\")\n",
    "ax[0].axis(\"equal\")\n",
    "ax[0].set_title(\"Dead\")\n",
    "\n",
    "ax[1].scatter(transformed_vectors_live[:,0], transformed_vectors_live[:,1], c=\"g\")\n",
    "ax[1].scatter([0],[0], c=\"r\")\n",
    "ax[1].scatter([l_live[0],-l_live[0]], [0.,0.],c=\"blue\")\n",
    "ax[1].scatter([0.,0.],[l_live[1],-l_live[1]],c=\"cyan\")\n",
    "ax[1].axis(\"equal\")\n",
    "ax[1].set_title(\"Live\")\n",
    "\n",
    "plt.suptitle(\"PCA plot around the Fréchet mean\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9c026e1",
   "metadata": {},
   "source": [
    "### 5. Statistical shape modelling"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f1f4fd9",
   "metadata": {},
   "source": [
    "**5.1** The following lines generate synthetic live C. elegans shapes relying on a simple statistical model (multivariate Gaussian) built from the covariance matrix of the dataset. Do you understand how the model is constructed? How good do you think this model is and why?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46ff2e7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0117bb1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "count=3\n",
    "num_synthetic_data=5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7747efe5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Live C. elegans\n",
    "synthetic_data_live=np.zeros((num_synthetic_data,N),dtype=np.complex)\n",
    "for i in range(num_synthetic_data):\n",
    "    rand=[]\n",
    "    for k in range(tangent_plane_pca_live.n_components_):\n",
    "        var=tangent_plane_pca_live.explained_variance_[k]*(len(tangent_vectors_live)-1)\n",
    "        if var<1e-6:\n",
    "            sigma=0.0\n",
    "        else:\n",
    "            sigma=np.sqrt(var)\n",
    "        rand.append(scipy.stats.norm.rvs(loc=0, scale=sigma, size=1))\n",
    "    rand=np.array(rand)\n",
    "\n",
    "    v=np.zeros((tangent_plane_pca_live.n_components_))\n",
    "    for k in range(count):\n",
    "        zz=rand.T @ tangent_plane_pca_live.components_.T[:,k]\n",
    "        v+=(zz*tangent_plane_pca_live.components_.T[:,k])\n",
    "    v+=tangent_plane_pca_live.mean_\n",
    "    \n",
    "    complex_v=v[:N]+1j*v[N:]\n",
    "    synthetic_data_live[i]=exponentialMap(mean_shape_Frechet_live, complex_v) \n",
    "    \n",
    "fig, ax = plt.subplots(1, num_synthetic_data, figsize=(num_synthetic_data,1))\n",
    "for i in range(num_synthetic_data):\n",
    "    ax[i].scatter(synthetic_data_live[i].real, synthetic_data_live[i].imag, s=5)\n",
    "    ax[i].axis(\"off\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8e960bc",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}