{ "cells": [ { "cell_type": "markdown", "id": "ff71666d", "metadata": {}, "source": [ "# Transporte entrópico y el algoritmo de Sinkhorn\n", "\n", "Este cuaderno acompaña el **Capítulo 13** de las notas del curso (*Transporte entrópico y el algoritmo de Sinkhorn*). Trabajamos con medidas discretas $a\\in\\Sigma_n$, $b\\in\\Sigma_m$ (pesos positivos) y una matriz de costos $C$, y estudiamos el problema regularizado\n", "\n", "$$\n", "\\mathrm{OT}_\\varepsilon(a,b)=\\min_{\\pi\\in\\Pi(a,b)}\\ \\langle C,\\pi\\rangle+\\varepsilon\\,\\mathrm{KL}(\\pi\\,|\\,a\\otimes b),\n", "\\qquad\n", "\\mathrm{KL}(\\pi\\,|\\,a\\otimes b)=\\sum_{ij}\\pi_{ij}\\log\\frac{\\pi_{ij}}{a_ib_j}.\n", "$$\n", "\n", "Exploramos computacionalmente:\n", "\n", "- La **forma del plan entrópico** $\\pi^\\varepsilon_{ij}=a_ib_je^{(\\varphi^\\varepsilon_i+\\psi^\\varepsilon_j-c_{ij})/\\varepsilon}$ y su comparación con el plan del programa lineal.\n", "- El **algoritmo de Sinkhorn** escrito desde cero, en la forma de escalamientos $(u,v)$ y en la forma de potenciales con la $c$-transformada suavizada; la **monotonía del dual** $\\mathcal D_\\varepsilon$.\n", "- La implementación **estabilizada en el dominio logarítmico**, necesaria para $\\varepsilon$ pequeño.\n", "- Los límites $\\varepsilon\\to0$ (plan óptimo de máxima entropía, potenciales de Kantorovich) y $\\varepsilon\\to\\infty$ (plan producto), con las cotas del capítulo.\n", "- La velocidad de convergencia y su dependencia de $\\varepsilon$.\n", "- **Baricentros entrópicos** con el algoritmo de Benamou–Carlier–Cuturi–Nenna–Peyré, recuperando los baricentros explícitos del capítulo anterior, y baricentros de formas en el plano.\n", "\n", "Usamos [POT](https://pythonot.github.io/) sólo como control: todo lo esencial está implementado en el notebook en pocas líneas." ] }, { "cell_type": "code", "execution_count": null, "id": "8b051f4d", "metadata": {}, "outputs": [], "source": [ "# @title\n", "pip install POT" ] }, { "cell_type": "code", "execution_count": null, "id": "af4d2e44", "metadata": {}, "outputs": [], "source": [ "# @title\n", "import numpy as np\n", "import matplotlib.pyplot as plt\n", "import ot\n", "from scipy.special import logsumexp\n", "from scipy.linalg import sqrtm\n", "\n", "rng = np.random.default_rng(7)\n", "np.set_printoptions(precision=4, suppress=True)" ] }, { "cell_type": "markdown", "id": "04178fb5", "metadata": {}, "source": [ "## 1. El problema entrópico y la forma del plan\n", "\n", "Empezamos con un ejemplo pequeño para poder mirar las matrices: $n=m=6$ puntos en $\\mathbb R$ con pesos aleatorios y costo cuadrático. Resolvemos el problema de Kantorovich exacto (LP) y el entrópico para varios $\\varepsilon$ (con `ot.sinkhorn`, que implementaremos nosotros en la sección siguiente).\n", "\n", "Recordar que POT usa el regularizador $-\\varepsilon H(\\pi)$, $H(\\pi)=-\\sum\\pi_{ij}(\\log\\pi_{ij}-1)$, que difiere de $\\varepsilon\\,\\mathrm{KL}(\\pi|a\\otimes b)$ en una constante sobre $\\Pi(a,b)$: el minimizador $\\pi^\\varepsilon$ es el mismo." ] }, { "cell_type": "code", "execution_count": null, "id": "f975ea41", "metadata": {}, "outputs": [], "source": [ "def KL(pi, a, b):\n", " ab = np.outer(a, b); m = pi > 0\n", " return np.sum(pi[m]*np.log(pi[m]/ab[m]))\n", "\n", "n = m = 6\n", "x = np.sort(rng.uniform(0, 1, n)); y = np.sort(rng.uniform(0, 1, m))\n", "a = rng.dirichlet(np.ones(n)); b = rng.dirichlet(np.ones(m))\n", "C = (x[:, None] - y[None, :])**2\n", "\n", "pi_lp = ot.emd(a, b, C)\n", "print(\"plan óptimo del LP (un vértice: a lo sumo n+m-1 = 11 entradas no nulas):\")\n", "print(pi_lp, \"\\n entradas no nulas:\", np.sum(pi_lp > 1e-12), \" costo:\", np.sum(C*pi_lp))\n", "\n", "for eps in [0.1, 0.01, 0.001]:\n", " pi_eps = ot.sinkhorn(a, b, C, eps, numItermax=100_000, stopThr=1e-12)\n", " print(f\"\\nplan entrópico, eps = {eps}: costo = {np.sum(C*pi_eps):.6f}, KL = {KL(pi_eps, a, b):.4f}, mínimo > 0: {pi_eps.min() > 0}\")\n", " print(pi_eps)" ] }, { "cell_type": "markdown", "id": "d8e6cfd3", "metadata": {}, "source": [ "El plan del LP es disperso (un vértice del politopo); el entrópico tiene todas las entradas positivas y, a medida que $\\varepsilon$ decrece, concentra la masa donde estaba la del LP y su costo se acerca a $\\mathrm{OT}(a,b)$. Verificamos la cota del capítulo\n", "\n", "$$\n", "0\\le\\mathrm{OT}_\\varepsilon(a,b)-\\mathrm{OT}(a,b)\\le\\varepsilon\\log\\frac1{\\min_ia_i},\n", "$$\n", "\n", "y también la desigualdad **trivial** $\\langle C,\\pi^\\varepsilon\\rangle\\ge\\mathrm{OT}(a,b)$ (el plan entrópico es un plan admisible, así que su costo no puede ser menor que el mínimo)." ] }, { "cell_type": "code", "execution_count": null, "id": "3d739614", "metadata": {}, "outputs": [], "source": [ "OT = np.sum(C*pi_lp)\n", "print(f\"{'eps':>8} {'':>12} {'OT_eps':>12} {'OT_eps - OT':>12} {'eps log(1/min a)':>18}\")\n", "for eps in [1, 0.3, 0.1, 0.03, 0.01, 0.003]:\n", " pi_eps = ot.sinkhorn(a, b, C, eps, numItermax=200_000, stopThr=1e-13)\n", " OTe = np.sum(C*pi_eps) + eps*KL(pi_eps, a, b)\n", " print(f\"{eps:>8} {np.sum(C*pi_eps):>12.6f} {OTe:>12.6f} {OTe-OT:>12.6f} {eps*np.log(1/a.min()):>18.6f}\")" ] }, { "cell_type": "markdown", "id": "fb71e737", "metadata": {}, "source": [ "## 2. Sinkhorn desde cero\n", "\n", "### 2.1 Escalamientos\n", "\n", "Con $K^\\varepsilon_{ij}=e^{-c_{ij}/\\varepsilon}$, el plan entrópico es $\\pi^\\varepsilon=\\operatorname{diag}(u^\\varepsilon)K^\\varepsilon\\operatorname{diag}(v^\\varepsilon)$, con $u^\\varepsilon_i=a_ie^{\\varphi^\\varepsilon_i/\\varepsilon}$, $v^\\varepsilon_j=b_je^{\\psi^\\varepsilon_j/\\varepsilon}$. Imponer las marginales alternadamente da el algoritmo:\n", "\n", "$$\n", "u\\leftarrow\\frac{a}{K^\\varepsilon v},\\qquad v\\leftarrow\\frac{b}{(K^\\varepsilon)^Tu}.\n", "$$\n", "\n", "Cada paso es una normalización de filas o de columnas. Lo implementamos midiendo, en cada iteración, el error en la marginal que **no** se acaba de imponer." ] }, { "cell_type": "code", "execution_count": null, "id": "1532bcaa", "metadata": {}, "outputs": [], "source": [ "def sinkhorn_uv(a, b, C, eps, iters=2000, tol=1e-12):\n", " K = np.exp(-C/eps)\n", " v = np.ones_like(b); errs = []\n", " for k in range(iters):\n", " u = a/(K @ v) # ahora las filas suman a\n", " v = b/(K.T @ u) # ahora las columnas suman b (y las filas ya no exactamente)\n", " err = np.sum(np.abs(u*(K @ v) - a))\n", " errs.append(err)\n", " if err < tol: break\n", " return u, v, K, np.array(errs)\n", "\n", "eps = 0.05\n", "u, v, K, errs = sinkhorn_uv(a, b, C, eps)\n", "pi = u[:, None]*K*v[None, :]\n", "print(\"iteraciones:\", len(errs))\n", "print(\"marginales correctas:\", np.allclose(pi.sum(1), a), np.allclose(pi.sum(0), b))\n", "print(\"coincide con POT:\", np.allclose(pi, ot.sinkhorn(a, b, C, eps, numItermax=100_000, stopThr=1e-13), atol=1e-9))" ] }, { "cell_type": "markdown", "id": "22774ea7", "metadata": {}, "source": [ "### 2.2 Potenciales y $c$-transformada suavizada\n", "\n", "La misma iteración, en los potenciales $\\varphi=\\varepsilon\\log(u/a)$, $\\psi=\\varepsilon\\log(v/b)$, es la alternancia de $c$-transformadas suavizadas:\n", "\n", "$$\n", "\\varphi\\leftarrow(\\psi)^{\\bar c,\\varepsilon},\\quad\n", "(\\psi)^{\\bar c,\\varepsilon}_i=-\\varepsilon\\log\\sum_jb_je^{(\\psi_j-c_{ij})/\\varepsilon};\n", "\\qquad\n", "\\psi\\leftarrow(\\varphi)^{c,\\varepsilon},\\quad\n", "(\\varphi)^{c,\\varepsilon}_j=-\\varepsilon\\log\\sum_ia_ie^{(\\varphi_i-c_{ij})/\\varepsilon}.\n", "$$\n", "\n", "Cada paso maximiza exactamente el dual entrópico\n", "\n", "$$\n", "\\mathcal D_\\varepsilon(\\varphi,\\psi)=\\langle a,\\varphi\\rangle+\\langle b,\\psi\\rangle-\\varepsilon\\sum_{ij}a_ib_je^{(\\varphi_i+\\psi_j-c_{ij})/\\varepsilon}+\\varepsilon\n", "$$\n", "\n", "en una de las dos variables, de modo que $\\mathcal D_\\varepsilon$ **crece en cada paso** y converge a $\\mathrm{OT}_\\varepsilon(a,b)$. Implementamos esta versión usando `logsumexp`, que es numéricamente estable: es la implementación **en el dominio logarítmico**, y es la que se usa en la práctica." ] }, { "cell_type": "code", "execution_count": null, "id": "773dc7c6", "metadata": {}, "outputs": [], "source": [ "def ctransf_eps(phi, a, C, eps):\n", " \"\"\"(phi)^{c,eps}_j = -eps log sum_i a_i exp((phi_i - c_ij)/eps), estable.\"\"\"\n", " return -eps*logsumexp((phi[:, None] - C)/eps, b=a[:, None], axis=0)\n", "\n", "def cbar_transf_eps(psi, b, C, eps):\n", " return -eps*logsumexp((psi[None, :] - C)/eps, b=b[None, :], axis=1)\n", "\n", "def dual_eps(phi, psi, a, b, C, eps):\n", " return a @ phi + b @ psi - eps*np.sum(np.outer(a, b)*np.exp((phi[:, None] + psi[None, :] - C)/eps)) + eps\n", "\n", "def sinkhorn_log(a, b, C, eps, iters=5000, tol=1e-12):\n", " psi = np.zeros_like(b); D = []\n", " for k in range(iters):\n", " phi = cbar_transf_eps(psi, b, C, eps)\n", " psi_new = ctransf_eps(phi, a, C, eps)\n", " D.append(dual_eps(phi, psi_new, a, b, C, eps))\n", " if np.max(np.abs(psi_new - psi)) < tol: psi = psi_new; break\n", " psi = psi_new\n", " pi = np.outer(a, b)*np.exp((phi[:, None] + psi[None, :] - C)/eps)\n", " return phi, psi, pi, np.array(D)\n", "\n", "phi, psi, pi_log, D = sinkhorn_log(a, b, C, eps)\n", "OTe = np.sum(C*pi_log) + eps*KL(pi_log, a, b)\n", "print(\"D_eps es no decreciente:\", np.all(np.diff(D) >= -1e-13))\n", "print(\"D_eps final =\", D[-1], \" OT_eps (primal) =\", OTe, \" (dualidad fuerte)\")\n", "print(\"mismo plan que la versión (u,v):\", np.allclose(pi_log, pi, atol=1e-10))\n", "print(\"relaciones psi = (phi)^{c,eps}, phi = (psi)^{cbar,eps}:\",\n", " np.allclose(psi, ctransf_eps(phi, a, C, eps)), np.allclose(phi, cbar_transf_eps(psi, b, C, eps)))\n", "\n", "plt.figure(figsize=(6, 3.5)); plt.semilogy(OTe - D + 1e-17, 'o-', ms=3)\n", "plt.xlabel('iteración'); plt.ylabel(r'$\\mathrm{OT}_\\varepsilon-\\mathcal{D}_\\varepsilon(\\varphi^{(k)},\\psi^{(k)})$'); plt.title('Maximización alternada del dual'); plt.show()" ] }, { "cell_type": "markdown", "id": "1b4cdc30", "metadata": {}, "source": [ "### 2.3 Por qué hace falta el dominio logarítmico\n", "\n", "Con $\\varepsilon$ pequeño, $K^\\varepsilon_{ij}=e^{-c_{ij}/\\varepsilon}$ tiene entradas del orden de $e^{-1/\\varepsilon}$: para $\\varepsilon=0.001$ y costos de orden $1$, eso es $10^{-434}$, que en doble precisión (cuyo menor número positivo es del orden de $10^{-308}$) es exactamente cero. La versión $(u,v)$ divide entonces por cero, mientras que la versión logarítmica sólo maneja números del orden de los costos." ] }, { "cell_type": "code", "execution_count": null, "id": "44771e68", "metadata": {}, "outputs": [], "source": [ "eps_chico = 0.001\n", "C2 = (x[:, None] - (y[None, :] + 2))**2 # el mismo problema con nu trasladada: todos los costos son >= 1\n", "OT2 = ot.emd2(a, b, C2)\n", "with np.errstate(all='ignore'):\n", " u_, v_, K_, errs_ = sinkhorn_uv(a, b, C2, eps_chico, iters=50)\n", " pi_uv = u_[:, None]*K_*v_[None, :]\n", "print(\"versión (u,v): entradas de K^eps no nulas:\", np.sum(K_ > 0), \"de\", K_.size, \"; el plan contiene NaN:\", np.isnan(pi_uv).any())\n", "\n", "phi_, psi_, pi_l, D_ = sinkhorn_log(a, b, C2, eps_chico, iters=20000, tol=1e-10)\n", "print(\"versión log-domain: marginales correctas:\", np.allclose(pi_l.sum(1), a, atol=1e-8), np.allclose(pi_l.sum(0), b, atol=1e-8),\n", " \"; iteraciones:\", len(D_), \"; costo:\", np.sum(C2*pi_l), \" vs OT =\", OT2)" ] }, { "cell_type": "markdown", "id": "c765b8c5", "metadata": {}, "source": [ "## 3. Velocidad de convergencia\n", "\n", "La demostración elemental del capítulo da convergencia sin tasa; la demostración de Franklin–Lorenz da convergencia lineal en la métrica de Hilbert, con una razón que se deteriora cuando $\\varepsilon\\to0$ (la razón se acerca a $1$ como $1-e^{-\\Delta/\\varepsilon}$ aproximadamente, con $\\Delta$ la oscilación del costo). Lo vemos midiendo el error en la marginal a lo largo de las iteraciones, para varios $\\varepsilon$, en un problema algo más grande." ] }, { "cell_type": "code", "execution_count": null, "id": "887faee9", "metadata": {}, "outputs": [], "source": [ "n = m = 200\n", "X = rng.uniform(0, 1, (n, 2)); Y = rng.uniform(0, 1, (m, 2)) + 0.3\n", "a = np.ones(n)/n; b = np.ones(m)/m\n", "C = ot.dist(X, Y) # euclídea al cuadrado\n", "\n", "def sinkhorn_log_errs(a, b, C, eps, iters):\n", " psi = np.zeros_like(b); errs = []\n", " for k in range(iters):\n", " phi = cbar_transf_eps(psi, b, C, eps)\n", " psi = ctransf_eps(phi, a, C, eps)\n", " pi = np.outer(a, b)*np.exp((phi[:, None] + psi[None, :] - C)/eps)\n", " errs.append(np.sum(np.abs(pi.sum(1) - a)))\n", " return np.array(errs)\n", "\n", "plt.figure(figsize=(7, 4))\n", "for eps in [0.3, 0.1, 0.03, 0.01, 0.003]:\n", " e = sinkhorn_log_errs(a, b, C, eps, 400)\n", " plt.semilogy(e, label=fr'$\\varepsilon={eps}$')\n", "plt.xlabel('iteración'); plt.ylabel(r'$\\|\\pi^{(k)}\\mathbf{1}-a\\|_1$'); plt.legend()\n", "plt.title(r'Convergencia lineal, más lenta cuanto menor es $\\varepsilon$'); plt.show()" ] }, { "cell_type": "markdown", "id": "e81cfc0e", "metadata": {}, "source": [ "Para $\\varepsilon$ moderado, el algoritmo converge en decenas de iteraciones; para $\\varepsilon$ pequeño se necesitan miles. Como cada iteración cuesta $O(nm)$, el costo total para $\\varepsilon$ moderado es incomparablemente menor que el del simplex para $n$ grande, y además las operaciones son productos matriz-vector, que se paralelizan trivialmente: esa es la razón del éxito práctico del método." ] }, { "cell_type": "markdown", "id": "03066b20", "metadata": {}, "source": [ "## 4. Los límites $\\varepsilon\\to\\infty$ y $\\varepsilon\\to0$\n", "\n", "### 4.1 $\\varepsilon\\to\\infty$: el plan producto\n", "\n", "Cuando $\\varepsilon\\to\\infty$ el término entrópico domina y $\\pi^\\varepsilon\\to a\\otimes b$, el único minimizador de $\\mathrm{KL}(\\cdot\\,|\\,a\\otimes b)$. En el ejemplo de $6\\times6$:" ] }, { "cell_type": "code", "execution_count": null, "id": "a97ac838", "metadata": {}, "outputs": [], "source": [ "n = m = 6\n", "x = np.sort(rng.uniform(0, 1, n)); y = np.sort(rng.uniform(0, 1, m))\n", "a = rng.dirichlet(np.ones(n)); b = rng.dirichlet(np.ones(m))\n", "C = (x[:, None] - y[None, :])**2\n", "\n", "print(f\"{'eps':>8} {'||pi_eps - a x b||_1':>22}\")\n", "for eps in [0.01, 0.1, 1, 10, 100, 1000]:\n", " _, _, pi_eps, _ = sinkhorn_log(a, b, C, eps)\n", " print(f\"{eps:>8} {np.sum(np.abs(pi_eps - np.outer(a, b))):>22.2e}\")" ] }, { "cell_type": "markdown", "id": "6eb51ac5", "metadata": {}, "source": [ "### 4.2 $\\varepsilon\\to0$: el plan óptimo de máxima entropía\n", "\n", "Cuando $\\varepsilon\\to0$, $\\pi^\\varepsilon$ converge al plan óptimo del problema de Kantorovich que **minimiza $\\mathrm{KL}(\\cdot\\,|\\,a\\otimes b)$ entre todos los planes óptimos**. Si el plan óptimo es único, simplemente converge a él. El caso interesante es el degenerado: elegimos un ejemplo con muchos planes óptimos, $a=b$ uniformes sobre $\\{0,1,2,3\\}$ y $\\{1,2,3,4\\}$ con costo $|x-y|$, donde $W_1=1$ se alcanza en muchos planes (mover cada punto un lugar, o mover sólo el $0$ hasta el $4$, o…)." ] }, { "cell_type": "code", "execution_count": null, "id": "d78a8437", "metadata": {}, "outputs": [], "source": [ "x = np.arange(4.); y = np.arange(1., 5.)\n", "a = b = np.ones(4)/4\n", "C = np.abs(x[:, None] - y[None, :])\n", "OT = ot.emd2(a, b, C)\n", "pi_vertice = ot.emd(a, b, C)\n", "print(\"W_1 =\", OT, \"\\nun plan óptimo (vértice) que devuelve el LP:\\n\", pi_vertice, \"\\n KL =\", KL(pi_vertice, a, b))\n", "\n", "# otros dos planes óptimos escritos a mano\n", "pi_shift = np.eye(4)/4 # i -> i+1 (x_i -> y_i)\n", "pi_salto = np.zeros((4, 4)); pi_salto[0, 3] = 1/4; pi_salto[1, 0] = pi_salto[2, 1] = pi_salto[3, 2] = 1/4 # 0->4, y los demás quietos\n", "for nombre, p in [(\"desplazar todos un lugar\", pi_shift), (\"saltar 0 -> 4\", pi_salto)]:\n", " print(f\"{nombre}: costo = {np.sum(C*p):.4f}, KL = {KL(p, a, b):.4f}\")\n", "\n", "print(\"\\nlímite eps -> 0 del plan entrópico:\")\n", "anterior = None\n", "for eps in [0.3, 0.1, 0.03, 0.01, 0.003, 0.001]:\n", " _, _, pi_eps, _ = sinkhorn_log(a, b, C, eps, iters=50000, tol=1e-13)\n", " cambio = \"\" if anterior is None else f\" ||pi_eps - pi_anterior||_1 = {np.sum(np.abs(pi_eps - anterior)):.2e}\"\n", " print(f\"eps = {eps:<6} costo = {np.sum(C*pi_eps):.6f} KL = {KL(pi_eps, a, b):.5f}{cambio}\")\n", " anterior = pi_eps\n", "print(\"\\npi^0 (numérico):\\n\", pi_eps)" ] }, { "cell_type": "markdown", "id": "f6d56c1c", "metadata": {}, "source": [ "El límite $\\pi^0$ es óptimo (su costo es $W_1=1$), tiene entropía relativa **menor** que la de los vértices (es una combinación convexa de varios de ellos), y no coincide con ninguno de los planes \"naturales\". Es el plan óptimo que Sinkhorn selecciona entre todos los posibles: la regularización actúa como criterio de selección.\n", "\n", "### 4.3 Los potenciales entrópicos convergen a potenciales de Kantorovich\n", "\n", "En un problema **no degenerado**, los potenciales $(\\varphi^\\varepsilon,\\psi^\\varepsilon)$ convergen, salvo la constante aditiva, a los potenciales de Kantorovich del LP. Lo verificamos con puntos aleatorios (el LP es entonces no degenerado con probabilidad 1) comparando con los potenciales duales que devuelve `ot.emd(..., log=True)`, normalizando $\\varphi_1=0$." ] }, { "cell_type": "code", "execution_count": null, "id": "85a1ff8d", "metadata": {}, "outputs": [], "source": [ "n = m = 8\n", "x = rng.uniform(0, 1, n); y = rng.uniform(0, 1, m)\n", "a = rng.dirichlet(np.ones(n)); b = rng.dirichlet(np.ones(m))\n", "C = (x[:, None] - y[None, :])**2\n", "\n", "_, log = ot.emd(a, b, C, log=True)\n", "phi_K, psi_K = log['u'] - log['u'][0], log['v'] + log['u'][0]\n", "print(\"potenciales de Kantorovich (LP), normalizados phi_1 = 0:\\n\", phi_K, \"\\n\", psi_K)\n", "print(\"valor dual =\", a @ phi_K + b @ psi_K, \"= OT =\", ot.emd2(a, b, C))\n", "\n", "print(f\"\\n{'eps':>8} {'||phi_eps - phi_K||_inf':>24} {'||psi_eps - psi_K||_inf':>24} {'D_eps':>10}\")\n", "for eps in [0.1, 0.03, 0.01, 0.003, 0.001]:\n", " phi_e, psi_e, pi_e, D = sinkhorn_log(a, b, C, eps, iters=100000, tol=1e-13)\n", " s = phi_e[0]; phi_e, psi_e = phi_e - s, psi_e + s\n", " print(f\"{eps:>8} {np.max(np.abs(phi_e - phi_K)):>24.2e} {np.max(np.abs(psi_e - psi_K)):>24.2e} {D[-1]:>10.6f}\")" ] }, { "cell_type": "markdown", "id": "793708a2", "metadata": {}, "source": [ "Los potenciales entrópicos convergen a los del LP cuando $\\varepsilon\\to0$, y el valor dual $\\mathcal D_\\varepsilon\\to\\mathrm{OT}$. Esto es lo que dice la observación del capítulo: el límite $\\varepsilon\\to0$ de los potenciales entrópicos da una **construcción alternativa de los potenciales de Kantorovich discretos**, sin pasar por el teorema de dualidad de la programación lineal.\n", "\n", "**Ejercicio.** Repetir con $a=b$ uniformes y $x,y$ equiespaciados (problema degenerado). ¿Convergen los potenciales? ¿A qué? Comparar con la no unicidad de los potenciales de Kantorovich en ese caso." ] }, { "cell_type": "markdown", "id": "8c735116", "metadata": {}, "source": [ "## 5. Baricentros entrópicos\n", "\n", "Sobre una grilla fija $\\{x_1,\\dots,x_n\\}$ y con medidas $b^1,\\dots,b^N$ sobre ella, el **baricentro entrópico** es el minimizador en $a\\in\\Sigma_n$ de $\\mathcal F_\\varepsilon(a)=\\sum_k\\lambda_k\\,\\mathrm{OT}_\\varepsilon(a,b^k)$. El algoritmo de Benamou, Carlier, Cuturi, Nenna y Peyré escribe cada plan como $\\pi^k=\\operatorname{diag}(u^k)K^\\varepsilon\\operatorname{diag}(v^k)$ y alterna:\n", "\n", "1. $u^k\\leftarrow a/(K^\\varepsilon v^k)$ para cada $k$;\n", "2. $v^k\\leftarrow b^k/((K^\\varepsilon)^Tu^k)$ para cada $k$;\n", "3. $a\\leftarrow\\prod_k\\bigl(u^k\\odot K^\\varepsilon v^k\\bigr)^{\\lambda_k}$ (media geométrica ponderada de las primeras marginales).\n", "\n", "Lo implementamos tal cual, en la versión $(u,v)$ y en la versión estabilizada en el dominio logarítmico (la primera falla en cuanto $K^\\varepsilon$ tiene ceros numéricos), y lo usamos para **recuperar numéricamente los baricentros explícitos del capítulo anterior**." ] }, { "cell_type": "code", "execution_count": null, "id": "4a813510", "metadata": {}, "outputs": [], "source": [ "def baricentro_sinkhorn(B, C, lam, eps, iters=2000, tol=1e-10):\n", " \"\"\"B: matriz n x N con las medidas b^k en columnas. Devuelve el baricentro a en la misma grilla.\"\"\"\n", " n, N = B.shape\n", " K = np.exp(-C/eps)\n", " V = np.ones((n, N)); a = np.ones(n)/n\n", " for it in range(iters):\n", " U = a[:, None]/(K @ V)\n", " V = B/(K.T @ U)\n", " a_new = np.prod((U*(K @ V))**lam[None, :], axis=1)\n", " a_new /= a_new.sum() # (la media geométrica no tiene masa 1 exactamente hasta converger)\n", " if np.sum(np.abs(a_new - a)) < tol: a = a_new; break\n", " a = a_new\n", " return a\n", "\n", "def baricentro_sinkhorn_log(B, C, lam, eps, iters=2000, tol=1e-10):\n", " \"\"\"El mismo algoritmo en el dominio logarítmico: f^k = eps log u^k, g^k = eps log v^k.\"\"\"\n", " n, N = B.shape\n", " logB = np.log(B + 1e-300)\n", " Gm = np.zeros((n, N)); loga = np.full(n, -np.log(n))\n", " for it in range(iters):\n", " # paso 1: u^k = a / (K v^k) -> f^k = eps log a - eps LSE_j((g^k_j - c_ij)/eps)\n", " Fm = np.stack([eps*loga - eps*logsumexp((Gm[None, :, k] - C)/eps, axis=1) for k in range(N)], axis=1)\n", " # paso 2: v^k = b^k / (K^T u^k) -> g^k = eps log b^k - eps LSE_i((f^k_i - c_ij)/eps)\n", " Gm = np.stack([eps*logB[:, k] - eps*logsumexp((Fm[:, None, k] - C)/eps, axis=0) for k in range(N)], axis=1)\n", " # paso 3: log(u^k (K v^k))_i = f^k_i/eps + LSE_j((g^k_j - c_ij)/eps); log a = sum_k lam_k (...)\n", " logmarg = np.stack([Fm[:, k]/eps + logsumexp((Gm[None, :, k] - C)/eps, axis=1) for k in range(N)], axis=1)\n", " loga_new = logmarg @ lam; loga_new -= logsumexp(loga_new)\n", " if np.sum(np.abs(np.exp(loga_new) - np.exp(loga))) < tol: loga = loga_new; break\n", " loga = loga_new\n", " return np.exp(loga)" ] }, { "cell_type": "markdown", "id": "ba31f0f9", "metadata": {}, "source": [ "### 5.1 Dimensión uno: el promedio de cuantiles\n", "\n", "Tomamos tres medidas en una grilla de $[-5,8]$ y comparamos el baricentro entrópico, para $\\varepsilon$ decreciente, con el promedio de cuantiles $F^{[-1]}=\\sum\\lambda_kF_k^{[-1]}$, que es el baricentro exacto." ] }, { "cell_type": "code", "execution_count": null, "id": "531dc0db", "metadata": {}, "outputs": [], "source": [ "grilla = np.linspace(-5, 8, 301); h = grilla[1] - grilla[0]\n", "C = (grilla[:, None] - grilla[None, :])**2\n", "lam = np.array([0.4, 0.4, 0.2])\n", "\n", "def densidad_en_grilla(f):\n", " p = f(grilla); return p/p.sum()\n", "from scipy.stats import norm, expon\n", "B = np.stack([densidad_en_grilla(lambda t: norm.pdf(t, -3, 0.6)),\n", " densidad_en_grilla(lambda t: 0.5*norm.pdf(t, 0, 0.3) + 0.5*norm.pdf(t, 1.5, 0.3)),\n", " densidad_en_grilla(lambda t: expon.pdf(t - 3, scale=1.0))], axis=1)\n", "\n", "# baricentro exacto por cuantiles, con las pseudoinversas de las distribuciones continuas\n", "ts = (np.arange(400_000) + 0.5)/400_000\n", "tt = np.linspace(-6, 9, 200_001)\n", "F_mix = 0.5*norm.cdf(tt, 0, 0.3) + 0.5*norm.cdf(tt, 1.5, 0.3)\n", "q_mix = np.interp(ts, F_mix, tt) # inversión numérica de la FDA de la mezcla\n", "q_bar = lam[0]*norm.ppf(ts, -3, 0.6) + lam[1]*q_mix + lam[2]*(expon.ppf(ts, scale=1.0) + 3)\n", "a_exacto, _ = np.histogram(q_bar, bins=np.concatenate([grilla - h/2, [grilla[-1] + h/2]])); a_exacto = a_exacto/a_exacto.sum()\n", "\n", "fig, ax = plt.subplots(figsize=(10, 4))\n", "for k in range(3): ax.fill_between(grilla, B[:, k]/h, alpha=0.3, label=fr'$b^{k+1}$')\n", "ax.plot(grilla, a_exacto/h, 'k', lw=2.5, label='baricentro exacto (cuantiles)')\n", "for eps, ls in [(0.5, ':'), (0.1, '--'), (0.02, '-')]:\n", " a_eps = baricentro_sinkhorn_log(B, C, lam, eps, iters=5000)\n", " ax.plot(grilla, a_eps/h, ls, lw=1.8, label=fr'entrópico, $\\varepsilon={eps}$')\n", " print(f\"eps = {eps}: W_2(baricentro entrópico, exacto) = {np.sqrt(ot.emd2(a_eps, a_exacto, C)):.4f}\")\n", "ax.legend(); ax.set_title('Baricentro entrópico vs. promedio de cuantiles'); plt.show()" ] }, { "cell_type": "markdown", "id": "401daa7f", "metadata": {}, "source": [ "Para $\\varepsilon$ grande el baricentro entrópico está **suavizado** (la entropía penaliza las concentraciones); cuando $\\varepsilon\\to0$ converge al baricentro exacto. Este sesgo de suavizado es el precio de la regularización, y en la práctica se elige $\\varepsilon$ del orden del cuadrado del paso de la grilla.\n", "\n", "### 5.2 Gaussianas en el plano: el punto fijo\n", "\n", "Discretizamos tres gaussianas en una grilla de $61\\times61$ puntos de $[-3,6]^2$ y comparamos la media y la covarianza del baricentro entrópico con la solución $\\bar\\Sigma$ de la ecuación de punto fijo $\\bar\\Sigma=\\sum\\lambda_k(\\bar\\Sigma^{1/2}\\Sigma_k\\bar\\Sigma^{1/2})^{1/2}$." ] }, { "cell_type": "code", "execution_count": null, "id": "50a9e76c", "metadata": {}, "outputs": [], "source": [ "g = np.linspace(-3, 6, 61); GX, GY = np.meshgrid(g, g, indexing='ij'); G = np.c_[GX.ravel(), GY.ravel()]\n", "C = ot.dist(G, G)\n", "\n", "def rot(th): return np.array([[np.cos(th), -np.sin(th)], [np.sin(th), np.cos(th)]])\n", "Sigmas = [rot(th) @ np.diag([2.0, 0.25]) @ rot(th).T for th in [0, np.pi/3, 2*np.pi/3]]\n", "medias = [np.array([0., 0.]), np.array([3., 0.]), np.array([1.5, 2.5])]\n", "lam = np.ones(3)/3\n", "\n", "def gauss_en_grilla(m, S):\n", " d = G - m; p = np.exp(-0.5*np.einsum('ij,jk,ik->i', d, np.linalg.inv(S), d)); return p/p.sum()\n", "B = np.stack([gauss_en_grilla(m, S) for m, S in zip(medias, Sigmas)], axis=1)\n", "\n", "# punto fijo exacto\n", "Sbar = np.mean(Sigmas, axis=0)\n", "for _ in range(200):\n", " Sh = np.real(sqrtm(Sbar)); Sbar = sum(l*np.real(sqrtm(Sh @ S @ Sh)) for l, S in zip(lam, Sigmas))\n", "mbar = sum(l*m for l, m in zip(lam, medias))\n", "\n", "eps = 0.05\n", "a_eps = baricentro_sinkhorn(B, C, lam, eps, iters=3000)\n", "m_emp = a_eps @ G; S_emp = (G - m_emp).T @ ((G - m_emp)*a_eps[:, None])\n", "print(\"media exacta :\", mbar, \"\\nmedia entrópica :\", m_emp)\n", "print(\"covarianza exacta :\\n\", Sbar, \"\\ncovarianza entrópica (eps = %.2f):\\n\" % eps, S_emp)\n", "print(\"paso de la grilla h =\", round(g[1]-g[0], 3))\n", "\n", "fig, ax = plt.subplots(1, 4, figsize=(16, 4))\n", "for k in range(3):\n", " ax[k].contourf(GX, GY, B[:, k].reshape(GX.shape), levels=12, cmap='Blues'); ax[k].set_title(fr'$b^{k+1}$')\n", "ax[3].contourf(GX, GY, a_eps.reshape(GX.shape), levels=12, cmap='Reds'); ax[3].set_title(fr'baricentro entrópico, $\\varepsilon={eps}$')\n", "th = np.linspace(0, 2*np.pi, 200); w, V = np.linalg.eigh(Sbar); el = (V*np.sqrt(w)) @ np.vstack([np.cos(th), np.sin(th)])*2\n", "ax[3].plot(mbar[0] + el[0], mbar[1] + el[1], 'k--', lw=1.5, label='elipse exacta (2 desvíos)'); ax[3].legend(fontsize=8)\n", "for a_ in ax: a_.set_aspect('equal')\n", "plt.show()" ] }, { "cell_type": "markdown", "id": "64692687", "metadata": {}, "source": [ "El baricentro entrópico recupera la media y la covarianza del punto fijo salvo errores del orden del paso de la grilla, del truncamiento de las colas al cuadrado $[-3,6]^2$ y del sesgo entrópico, y la elipse exacta se superpone a las curvas de nivel del baricentro calculado.\n", "\n", "### 5.3 Baricentros de formas\n", "\n", "Cerramos con el ejemplo que hizo populares a los baricentros de Wasserstein: interpolar entre **formas** vistas como medidas uniformes en el plano. Con tres formas y pesos $(\\lambda_1,\\lambda_2,\\lambda_3)$ variando en el triángulo de pesos, se obtiene una familia de formas intermedias. Aquí usamos la implementación de POT `ot.bregman.convolutional_barycenter2d` (Solomon et al., 2015): es el mismo algoritmo de la Sección 5, pero en una grilla regular con costo cuadrático el núcleo $K^\\varepsilon$ es una gaussiana separable, y multiplicar por $K^\\varepsilon$ se reduce a dos convoluciones unidimensionales de costo $O(N^3)$ en vez de $O(N^4)$ para imágenes de $N\\times N$. Eso permite grillas más finas y $\\varepsilon$ más chico. (El costo en POT está normalizado a la grilla $[0,1]^2$.)" ] }, { "cell_type": "code", "execution_count": null, "id": "356c1e6b", "metadata": {}, "outputs": [], "source": [ "N = 50\n", "g = np.linspace(-1, 1, N); GX, GY = np.meshgrid(g, g, indexing='ij'); G = np.c_[GX.ravel(), GY.ravel()]\n", "disco = (GX**2 + GY**2 <= 0.5**2).astype(float)\n", "cuadrado = ((np.abs(GX) <= 0.45) & (np.abs(GY) <= 0.45)).astype(float)\n", "triangulo = ((GY >= -0.5) & (GY <= 0.5 - np.sqrt(3)*np.abs(GX)) ).astype(float)\n", "formas = [disco, cuadrado, triangulo]\n", "A = np.stack([f/f.sum() for f in formas]) # imágenes N x N normalizadas\n", "eps = 0.001\n", "\n", "fig, ax = plt.subplots(5, 5, figsize=(10, 10))\n", "for i in range(5):\n", " for j in range(5):\n", " ax[i, j].axis('off')\n", " if i + j > 4: continue\n", " # coordenadas baricéntricas en el triángulo de pesos\n", " l1, l2 = (4 - i - j)/4, j/4; l3 = 1 - l1 - l2\n", " w = np.array([l1, l2, l3])\n", " bar = ot.bregman.convolutional_barycenter2d(A, eps, weights=w, numItermax=3000, stopThr=1e-8)\n", " ax[i, j].imshow(bar.T, origin='lower', cmap='Greys')\n", " ax[i, j].set_title(f'({l1:.2f}, {l2:.2f}, {l3:.2f})', fontsize=8)\n", "plt.suptitle(r'Baricentros entrópicos de tres formas; pesos $(\\lambda_{\\rm disco},\\lambda_{\\rm cuadrado},\\lambda_{\\rm triángulo})$')\n", "plt.tight_layout(); plt.show()" ] }, { "cell_type": "markdown", "id": "21451fc5", "metadata": {}, "source": [ "Las formas intermedias (levemente difuminadas por la entropía, incluso en los vértices del triángulo) son genuinas interpolaciones geométricas —un disco que se va cuadrando, un cuadrado que se va afinando hacia un triángulo—, no superposiciones de las tres formas originales, que es lo que daría la mezcla $\\sum\\lambda_kb^k$. Esto es exactamente la diferencia entre la geometría de $W_2$ y la geometría lineal de las medidas que recorrió todo el curso.\n", "\n", "## Ejercicios computacionales\n", "\n", "Los enunciados siguientes figuran también en la sección de ejercicios del capítulo correspondiente de las notas.\n", "\n", "1. **Convergencia de valores.** Graficar $\\varepsilon\\mapsto\\mathrm{OT}_\\varepsilon(a,b)-\\mathrm{OT}(a,b)$ en escala log-log para el ejemplo de la Sección 1 y estimar el orden de convergencia. ¿Es $O(\\varepsilon)$ como dice la cota, o mejor? (Para el problema no degenerado, la convergencia real es $O(\\varepsilon\\log(1/\\varepsilon))$ o incluso $O(\\varepsilon)$ con una constante menor que $\\log(1/\\min a_i)$.)\n", "\n", "2. **Sinkhorn como iteración $\\varphi\\mapsto\\varphi^{c\\bar c}$.** Implementar la mejora de potenciales del capítulo de dualidad, $\\varphi\\mapsto(\\varphi^c)^{\\bar c}$, sin regularización, y compararla con la iteración de Sinkhorn para $\\varepsilon$ pequeño partiendo del mismo $\\varphi$. ¿Converge la versión no regularizada? (No en general: puede oscilar; ésa es una de las razones para regularizar.)\n", "\n", "3. **Coste de cada iteración.** Medir el tiempo de una iteración de Sinkhorn para $n=m=500,1000,2000,4000$ y verificar el crecimiento $O(nm)$. Comparar con el tiempo de `ot.emd` para los mismos tamaños.\n", "\n", "4. **Un baricentro con $\\varepsilon\\to0$.** En la Sección 5.1, calcular el baricentro entrópico para $\\varepsilon\\in\\{0.5,0.2,0.1,0.05,0.02,0.01,0.005\\}$ y graficar $W_2$(baricentro entrópico, exacto) en función de $\\varepsilon$, en escala log-log. ¿Cuántas iteraciones hacen falta en cada caso?\n", "\n", "5. **Unicidad del baricentro entrópico.** Verificar numéricamente que el baricentro entrópico del ejemplo de no unicidad del capítulo anterior ($\\mu_1$ con átomos en $(\\pm1,0)$, $\\mu_2$ con átomos en $(0,\\pm1)$, sobre una grilla que contenga los cuatro candidatos) es **único** y simétrico: la entropía rompe la degeneración eligiendo la combinación más \"repartida\"." ] } ], "metadata": { "colab": { "provenance": [] }, "kernelspec": { "display_name": "Python 3", "name": "python3" }, "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 5 }