Visualization of optimal transport

Let $X=\{x_1,...,x_n\}\subset \mathbb R^2$ be a set of $n$ locations. Each location $x_i$ is associated with a mass $a_i>0$. Let $a=(a_1,...,a_n)\in \mathbb R^n_+$. Similarly, define some other locations $Y=\{y_1,....y_m\}\subset \mathbb R^2$ and $b=(b_1,...,b_m)\in \mathbb R^m_+$. We assume that the total mass of $X$ and $Y$ are equal, i.e. $\sum_{i=1}^n a_i=\sum_{i=1}^m b_i$.

Loosely speaking, Optimal transport is an optimal way to move masses from $X$ to $Y$ so that it minimizes the distance. Note that the masses can be splitted or merged during the “migration”. In the following figures, the heat-map represents the “blurring” mass of the plane.

Code: Jupyter Notebook
Reference: Optimal Transport by Gabriel Peyré