{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Practical Session: Optimal Transport on the Real Line\n",
    "\n",
    "In this practical session, we will consider the simple case of optimal transport on the real line. We will consider the case where the cost function is the squared Euclidean distance.\n",
    "\n",
    "**Instructions**\n",
    "\n",
    "* You should submit one single Jupyter notebook with the name **`2025-OT-1-Name1-Name2-Name3.ipynb`**, where **Name1**, **Name2**, and **Name3** are the last names of the three members of your trinom.\n",
    "\n",
    "* You need to submit it by email before **November 4th 2025, 13:59:59** with the subject **`2025-OT-1-Name1-Name2-Name3`** at **`samuel.vaiter@cnrs.fr`**. *No other subject allowed (otherwise grade = 0.)*\n",
    "\n",
    "* You are required to write your code in **Python**, using the **NumPy** library, and the theoretical questions should be answered using a **Markdown cell** (these cells support the use of LaTeX equations).\n",
    "\n",
    "* (If this rule is not followed ⇒ **grade = 0**.) You **can** use ChatGPT or other AI assistants, but are **required** to state it clearly, e.g,\n",
    "\n",
    "  > \"I used ChatGPT to help me write the conclusion of question 2\",\n",
    "  > \"ChatGPT solved the whole question 4\", etc.)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We are going to use Numpy and Matplotlib for this practical.\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We consider empirical distributions of the form\n",
    "$$\n",
    "\\mu_n = \\frac{1}{n} \\sum_{i=1}^{n} \\delta_{x_i} \\quad \\text{and} \\quad \\nu_n = \\frac{1}{n} \\sum_{i=1}^{n} \\delta_{y_i}\n",
    "$$\n",
    "for $n \\in \\mathbb{N}^*$ and points $x_i, y_i \\in \\mathbb{R}$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. (a.) Write a function that takes as input the $n$-vectors $(x_i)_{i=1}^n$ and $(y_j)_{j=1}^n$ and returns the optimal transport **plan/coupling**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "# CODE HERE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. (b.) Write a function that takes as input the $n$-vectors $(x_i)_{i=1}^n$ and $(y_j)_{j=1}^n$ and returns the permutation of indices associated to the **Monge map $T$** that characterize the optimal transport plan."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "# CODE HERE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. Let $\\mu = \\frac{3}{4} L|_{[0, 2/3]} + \\frac{3}{2} L|_{[2/3, 1]}$ and $\\nu = \\frac{3}{4} L|_{[1/3, 1]} + \\frac{3}{2} L|_{[0, 1/3]}$ where $L$ is the Lebesgue measure. What is the optimal transport map between $\\mu$ to $\\nu$?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "_Write your answer here_"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3. Draw $n$ independent samples from $\\mu$ and from $\\nu$ and represent the optimal transport plan between the empirical distributions $\\hat{\\mu}_n$ and $\\hat{\\nu}_n$ on $[0, 1]^2$ (plot is via a point cloud). Plot for $n = 20$ and $n = 300$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "# CODE HERE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4. Compute the 2-Wasserstein distance between $\\mu_n$ and $\\nu_n$ and plot it as a function of $n$, for $n = 1$ to $n = 500$. What is the (almost sure) limit of $W_2(\\mu_n, \\nu_n)$?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "# CODE HERE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "_Write your answer here_"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "5. Represent the $W_2$ optimal transport geodesic between $\\mu_n$ and $\\nu_n$ at times $t \\in \\{0, 1/4, 1/2, 3/4, 1\\}$. You can plot for instance the histogram.\n",
    "\n",
    "_Optional_: do an animation from 0 to 1 instead (using FuncAnimation from matplotlib)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# CODE HERE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now consider probability distributions that are given by their discretized density on a uniform grid. More specifically, we consider a fixed uniform grid $(x_i)_{i=1}^n$ on $[0, 1]$ with $x_i = \\frac{i}{n} - \\frac{1}{2n}$ (say, with $n = 200$) and measures of the form \n",
    "$$\n",
    "\\mu_a = \\sum_{i=1}^n a_i \\delta_{x_i}\n",
    "$$\n",
    "with $a \\in \\mathbb{R}^n$ such that $\\sum_{i=1}^n a_i = 1$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "6. Write a function that takes as input two $n$-vectors $a$ and $b$ in $\\mathbb{R}_+^n$ such that $\\sum a_i = \\sum b_i = 1$ and returns the discrete optimal transport plan represented by a matrix $P \\in \\mathbb{R}_+^{n \\times n}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "# CODE HERE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "7. Let $\\mu_a$ and $\\mu_b$ be the discretization of truncated Gaussian distributions on $[0,1]$ of (mean, variance) $(0.2, 0.12)$ and $(0.6, 0.22)$ respectively. Do not forget to normalize the weight vectors after discretization. Represent the optimal transport plan in greyscale on $[0, 1]^2$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "# CODE HERE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "8. Because of the discretization, we see that the optimal transport plan is not deterministic (while it should be in the continuous world). A workaround is to define the **barycentric projection map** \n",
    "\n",
    "$$T(x_i) = \\frac{\\sum_{j=1}^{n} x_j P_{i,j}}{\\sum_{j=1}^{n} P_{i,j}}$$\n",
    "\n",
    "which is well-defined whenever $a_i = \\sum_{j=1}^{n} P_{i,j} > 0$. Using this map, plot the (approximate) geodesic between $\\mu_a$ and $\\mu_b$ at times $t \\in \\{0, 1/4, 1/2, 3/4, 1\\}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# CODE HERE"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}