{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "19423e09",
   "metadata": {},
   "source": [
    "# Baricentros de Wasserstein\n",
    "\n",
    "Este cuaderno acompaña el **Capítulo 12** de las notas del curso (*Baricentros de Wasserstein*). Dadas medidas $\\mu_1,\\dots,\\mu_N\\in\\mathcal P_2(\\mathbb R^d)$ y pesos $\\lambda_i>0$ con $\\sum\\lambda_i=1$, un **baricentro** es un minimizador de\n",
    "\n",
    "$$\n",
    "\\mathcal F(\\nu)=\\sum_{i=1}^N\\lambda_i\\,W_2^2(\\mu_i,\\nu),\\qquad \\nu\\in\\mathcal P_2(\\mathbb R^d).\n",
    "$$\n",
    "\n",
    "Exploramos computacionalmente:\n",
    "\n",
    "- Los **casos ancla**: masas de Dirac y $N=2$, donde el baricentro es un punto de la geodésica.\n",
    "- El **caso unidimensional**: el baricentro es el promedio de las pseudoinversas, $F_\\nu^{[-1]}=\\sum\\lambda_iF_{\\mu_i}^{[-1]}$.\n",
    "- El **caso gaussiano**: la ecuación de punto fijo para la covarianza y su resolución iterativa.\n",
    "- El **criterio de optimalidad** $\\sum\\lambda_iT_i=\\mathrm{id}$ y el algoritmo de soporte libre que lo usa como iteración.\n",
    "- Un ejemplo de **no unicidad** con medidas discretas en el plano.\n",
    "- La **formulación multimarginal** resuelta como programa lineal.\n",
    "\n",
    "Usamos [POT](https://pythonot.github.io/) para los problemas de transporte y `scipy` para el programa lineal multimarginal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f55ce8a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title\n",
    "pip install POT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "445e459f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import ot\n",
    "from scipy.linalg import sqrtm\n",
    "from scipy.optimize import linprog\n",
    "\n",
    "rng = np.random.default_rng(1)\n",
    "\n",
    "def W2_discreta(X, a, Y, b):\n",
    "    \"\"\"W_2 entre dos medidas discretas (soportes X, Y; pesos a, b).\"\"\"\n",
    "    M = ot.dist(X, Y)            # euclídea al cuadrado por defecto\n",
    "    return np.sqrt(ot.emd2(a, b, M, numItermax=10_000_000))\n",
    "\n",
    "def F_bar(nu_X, nu_w, medidas, lam):\n",
    "    \"\"\"Funcional de baricentro F(nu) = sum lam_i W_2^2(mu_i, nu) para medidas discretas.\"\"\"\n",
    "    return sum(l*W2_discreta(X, a, nu_X, nu_w)**2 for l, (X, a) in zip(lam, medidas))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0dbeabf3",
   "metadata": {},
   "source": [
    "## 1. Casos ancla\n",
    "\n",
    "### 1.1 Masas de Dirac\n",
    "\n",
    "Si $\\mu_i=\\delta_{x_i}$, entonces $W_2^2(\\delta_{x_i},\\nu)=\\int|x_i-y|^2\\,d\\nu(y)$ y\n",
    "\n",
    "$$\n",
    "\\mathcal F(\\nu)=\\int\\sum_i\\lambda_i|x_i-y|^2\\,d\\nu(y)\\ \\ge\\ \\sum_i\\lambda_i|x_i-\\bar x|^2,\\qquad \\bar x=\\sum_i\\lambda_ix_i,\n",
    "$$\n",
    "\n",
    "con igualdad si y sólo si $\\nu=\\delta_{\\bar x}$: el baricentro de Wasserstein de masas de Dirac es la masa de Dirac en el baricentro euclídeo. Verificamos comparando $\\mathcal F(\\delta_{\\bar x})$ con $\\mathcal F$ evaluado en otros candidatos."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c68d66a",
   "metadata": {},
   "outputs": [],
   "source": [
    "xs = np.array([[0., 0.], [4., 0.], [1., 3.]])\n",
    "lam = np.array([0.5, 0.3, 0.2])\n",
    "xbar = lam @ xs\n",
    "medidas = [(x[None, :], np.array([1.])) for x in xs]\n",
    "\n",
    "print(\"baricentro euclídeo:\", xbar)\n",
    "print(\"F(delta_xbar)              =\", F_bar(xbar[None, :], np.array([1.]), medidas, lam))\n",
    "print(\"F(delta en otro punto)     =\", F_bar(np.array([[1., 1.]]), np.array([1.]), medidas, lam))\n",
    "print(\"F(medida repartida en x_i) =\", F_bar(xs, lam, medidas, lam))\n",
    "print(\"valor teórico sum lam_i |x_i - xbar|^2 =\", np.sum(lam*np.sum((xs - xbar)**2, axis=1)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37ffbd61",
   "metadata": {},
   "source": [
    "### 1.2 Dos medidas: el baricentro está sobre la geodésica\n",
    "\n",
    "Para $N=2$ y pesos $(1-t,t)$, el baricentro es $\\mu_t=(e_t)_\\#\\pi$, el punto de la geodésica de McCann entre $\\mu_1$ y $\\mu_2$, y el valor mínimo es\n",
    "\n",
    "$$\n",
    "\\min_\\nu\\ (1-t)\\,W_2^2(\\mu_1,\\nu)+t\\,W_2^2(\\nu,\\mu_2)=t(1-t)\\,W_2^2(\\mu_1,\\mu_2).\n",
    "$$\n",
    "\n",
    "(La desigualdad $\\ge$ es Cauchy–Schwarz sobre la desigualdad triangular; la igualdad se alcanza en $\\mu_t$.) Lo verificamos con dos gaussianas en el plano, para las que todo es explícito: la geodésica es $N(m_t,\\Sigma_t)$ con $\\Sigma_t=B_t\\Sigma_1B_t$, $B_t=(1-t)I+tA$, donde $T(x)=m_2+A(x-m_1)$ es el mapa óptimo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce2a5c6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def W2_gauss(m0, S0, m1, S1):\n",
    "    S0h = np.real(sqrtm(S0)); C = np.real(sqrtm(S0h @ S1 @ S0h))\n",
    "    return np.sqrt(np.sum((m0 - m1)**2) + np.trace(S0 + S1 - 2*C))\n",
    "\n",
    "def mapa_gauss(S0, S1):\n",
    "    \"\"\"Matriz A (simétrica definida positiva) con A S0 A = S1.\"\"\"\n",
    "    S0h = np.real(sqrtm(S0)); S0hi = np.linalg.inv(S0h)\n",
    "    return S0hi @ np.real(sqrtm(S0h @ S1 @ S0h)) @ S0hi\n",
    "\n",
    "m1, S1 = np.array([0., 0.]), np.array([[2.0, 0.5], [0.5, 0.5]])\n",
    "m2, S2 = np.array([4., 2.]), np.array([[0.5, -0.3], [-0.3, 1.5]])\n",
    "A = mapa_gauss(S1, S2); I = np.eye(2)\n",
    "W12 = W2_gauss(m1, S1, m2, S2)\n",
    "\n",
    "print(f\"{'t':>5} {'F(mu_t)':>10} {'t(1-t)W^2':>10}\")\n",
    "for t in [0.2, 0.5, 0.7]:\n",
    "    Bt = (1-t)*I + t*A\n",
    "    mt, St = (1-t)*m1 + t*m2, Bt @ S1 @ Bt\n",
    "    F = (1-t)*W2_gauss(m1, S1, mt, St)**2 + t*W2_gauss(mt, St, m2, S2)**2\n",
    "    print(f\"{t:>5} {F:>10.5f} {t*(1-t)*W12**2:>10.5f}\")\n",
    "\n",
    "# ... y que cualquier otra gaussiana da un valor mayor\n",
    "t = 0.5; Bt = (1-t)*I + t*A; mt, St = (1-t)*m1 + t*m2, Bt @ S1 @ Bt\n",
    "for nombre, (m, S) in {\"promedio de medias y covarianzas\": ((m1+m2)/2, (S1+S2)/2),\n",
    "                       \"mu_t con otra covarianza\": (mt, I)}.items():\n",
    "    F = (1-t)*W2_gauss(m1, S1, m, S)**2 + t*W2_gauss(m, S, m2, S2)**2\n",
    "    print(f\"F({nombre}) = {F:.5f}  >  {t*(1-t)*W12**2:.5f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "884967a6",
   "metadata": {},
   "source": [
    "## 2. Dimensión uno: promedio de cuantiles\n",
    "\n",
    "En $\\mathbb R$ el baricentro se calcula explícitamente: es la medida $\\nu$ cuya pseudoinversa es el promedio ponderado de las pseudoinversas,\n",
    "\n",
    "$$\n",
    "F_\\nu^{[-1]}=\\sum_{i=1}^N\\lambda_i\\,F_{\\mu_i}^{[-1]}.\n",
    "$$\n",
    "\n",
    "La razón es que en dimensión uno la isometría $\\mu\\mapsto F_\\mu^{[-1]}\\in L^2(0,1)$ convierte $W_2$ en la norma de $L^2$, y el baricentro en $L^2$ de $N$ funciones es su promedio. Para medidas empíricas con el mismo número de átomos equiponderados esto es simplemente promediar los datos ordenados.\n",
    "\n",
    "Comparamos con el baricentro calculado **sin usar la estructura unidimensional**: fijamos una grilla fina como soporte y resolvemos el programa lineal en los pesos con `ot.lp.barycenter`. Ese es el problema de baricentro con soporte fijo, que es un programa lineal porque $\\nu\\mapsto W_2^2(\\mu_i,\\nu)$ es el valor de un LP cuyas restricciones dependen linealmente de $\\nu$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a99c76e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 2000\n",
    "mus = [rng.normal(-3, 0.6, n),\n",
    "       np.concatenate([rng.normal(0, 0.3, n//2), rng.normal(1.5, 0.3, n//2)]),\n",
    "       rng.exponential(1.0, n) + 3]\n",
    "lam = np.array([0.4, 0.4, 0.2])\n",
    "\n",
    "# baricentro por cuantiles: promedio ponderado de los datos ordenados\n",
    "bar_q = sum(l*np.sort(m) for l, m in zip(lam, mus))\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(9, 3.8))\n",
    "bins = np.linspace(-5, 8, 130)\n",
    "for k, m in enumerate(mus):\n",
    "    ax.hist(m, bins=bins, density=True, alpha=0.4, label=fr'$\\mu_{k+1}$ ($\\lambda={lam[k]}$)')\n",
    "ax.hist(bar_q, bins=bins, density=True, histtype='step', lw=2.5, color='k', label='baricentro (cuantiles)')\n",
    "ax.legend(); ax.set_title('Baricentro en dimensión uno'); plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "188519e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# el mismo baricentro por programación lineal con soporte fijo en una grilla\n",
    "grilla = np.linspace(-5, 8, 261)\n",
    "def histograma(m, grilla):\n",
    "    h, _ = np.histogram(m, bins=np.concatenate([[grilla[0]-1e9], (grilla[1:]+grilla[:-1])/2, [grilla[-1]+1e9]]))\n",
    "    return h/h.sum()\n",
    "\n",
    "Amat = np.stack([histograma(m, grilla) for m in mus], axis=1)     # columnas: las mu_i discretizadas\n",
    "M = ot.dist(grilla[:, None], grilla[:, None])                    # costo cuadrático en la grilla\n",
    "bar_lp = ot.lp.barycenter(Amat, M, weights=lam)\n",
    "\n",
    "plt.figure(figsize=(9, 3.8))\n",
    "plt.plot(grilla, histograma(bar_q, grilla), 'k', lw=2.5, label='cuantiles (discretizado)')\n",
    "plt.plot(grilla, bar_lp, 'C3--', lw=2, label='LP con soporte en la grilla')\n",
    "plt.legend(); plt.title('Dos cálculos del mismo baricentro'); plt.show()\n",
    "\n",
    "print(\"W_2 entre ambos baricentros (discretizados):\",\n",
    "      W2_discreta(grilla[:, None], histograma(bar_q, grilla), grilla[:, None], bar_lp))\n",
    "print(\"paso de la grilla:\", grilla[1]-grilla[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c00b659",
   "metadata": {},
   "source": [
    "Las dos curvas coinciden salvo por el error de discretización (la distancia entre ambas es del orden del paso de la grilla). Observar que el baricentro **no** es la mezcla $\\sum\\lambda_i\\mu_i$: tiene un solo bloque de masa, más concentrado, y hereda de $\\mu_2$ una bimodalidad atenuada. Es el mismo fenómeno de la interpolación por desplazamiento frente a la lineal.\n",
    "\n",
    "**Ejercicio.** Verificar numéricamente que el baricentro de $N$ gaussianas unidimensionales $N(m_i,\\sigma_i^2)$ es la gaussiana $N\\bigl(\\sum\\lambda_im_i,(\\sum\\lambda_i\\sigma_i)^2\\bigr)$, y probarlo a mano con la fórmula de los cuantiles."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7eb0d5e",
   "metadata": {},
   "source": [
    "## 3. Gaussianas: la ecuación de punto fijo\n",
    "\n",
    "Si $\\mu_i=N(m_i,\\Sigma_i)$ con $\\Sigma_i$ definidas positivas, el baricentro es $N(\\bar m,\\bar\\Sigma)$ con $\\bar m=\\sum\\lambda_im_i$ y $\\bar\\Sigma$ la única solución definida positiva de\n",
    "\n",
    "$$\n",
    "\\bar\\Sigma=\\sum_{i=1}^N\\lambda_i\\bigl(\\bar\\Sigma^{1/2}\\Sigma_i\\bar\\Sigma^{1/2}\\bigr)^{1/2}.\n",
    "$$\n",
    "\n",
    "En las notas se prueba que una solución de esta ecuación **es** baricentro, mediante el criterio $\\sum\\lambda_iT_i=\\mathrm{id}$: los mapas óptimos $T_i(x)=m_i+A_i(x-\\bar m)$ de $\\bar\\nu$ a $\\mu_i$ tienen $A_i=\\bar\\Sigma^{-1/2}(\\bar\\Sigma^{1/2}\\Sigma_i\\bar\\Sigma^{1/2})^{1/2}\\bar\\Sigma^{-1/2}$, y la ecuación dice exactamente que $\\sum\\lambda_iA_i=I$. La existencia y unicidad de la solución, y la convergencia de la iteración\n",
    "\n",
    "$$\n",
    "\\Sigma_{k+1}=\\sum_i\\lambda_i\\bigl(\\Sigma_k^{1/2}\\Sigma_i\\Sigma_k^{1/2}\\bigr)^{1/2}\n",
    "$$\n",
    "\n",
    "están demostradas en Álvarez-Esteban, del Barrio, Cuesta-Albertos y Matrán (2016). Implementamos la iteración y comprobamos ambas cosas."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f190e25",
   "metadata": {},
   "outputs": [],
   "source": [
    "def baricentro_gauss(Sigmas, lam, iters=200, tol=1e-12):\n",
    "    S = np.mean(Sigmas, axis=0)\n",
    "    hist = []\n",
    "    for k in range(iters):\n",
    "        Sh = np.real(sqrtm(S))\n",
    "        S_new = sum(l*np.real(sqrtm(Sh @ Si @ Sh)) for l, Si in zip(lam, Sigmas))\n",
    "        hist.append(np.linalg.norm(S_new - S))\n",
    "        S = S_new\n",
    "        if hist[-1] < tol: break\n",
    "    return S, hist\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.5, 0.3]) @ 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.array([1/3, 1/3, 1/3])\n",
    "\n",
    "Sbar, hist = baricentro_gauss(Sigmas, lam)\n",
    "mbar = sum(l*m for l, m in zip(lam, medias))\n",
    "\n",
    "# 1) criterio de optimalidad: sum lam_i A_i = I\n",
    "As = [mapa_gauss(Sbar, Si) for Si in Sigmas]\n",
    "print(\"sum lam_i A_i =\\n\", sum(l*Ai for l, Ai in zip(lam, As)))\n",
    "\n",
    "# 2) comparación con POT\n",
    "m_pot, S_pot = ot.gaussian.bures_wasserstein_barycenter(np.array(medias), np.array(Sigmas), weights=lam)\n",
    "print(\"\\n||Sigma_bar - Sigma_POT|| =\", np.linalg.norm(Sbar - S_pot), \"   ||m_bar - m_POT|| =\", np.linalg.norm(mbar - m_pot))\n",
    "\n",
    "plt.figure(figsize=(5, 3.5)); plt.semilogy(hist, 'o-'); plt.xlabel('iteración'); plt.ylabel(r'$\\|\\Sigma_{k+1}-\\Sigma_k\\|$')\n",
    "plt.title('Convergencia del punto fijo'); plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19a4583d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def elipse(m, S, ax, **kw):\n",
    "    w, V = np.linalg.eigh(S); th = np.linspace(0, 2*np.pi, 200)\n",
    "    pts = (V*np.sqrt(w)) @ np.vstack([np.cos(th), np.sin(th)])*2\n",
    "    ax.plot(m[0]+pts[0], m[1]+pts[1], **kw)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(7, 5.5))\n",
    "for k, (m, S) in enumerate(zip(medias, Sigmas)):\n",
    "    elipse(m, S, ax, color=f'C{k}', lw=2, label=fr'$\\mu_{k+1}$')\n",
    "elipse(mbar, Sbar, ax, color='k', lw=3, label='baricentro')\n",
    "# comparación: la \"mezcla\" tiene covarianza mucho mayor\n",
    "S_mezcla = sum(l*(S + np.outer(m-mbar, m-mbar)) for l, m, S in zip(lam, medias, Sigmas))\n",
    "elipse(mbar, S_mezcla, ax, color='gray', ls='--', lw=1.5, label=r'mezcla $\\sum\\lambda_i\\mu_i$ (2 desvíos)')\n",
    "ax.axis('equal'); ax.legend(); ax.set_title('Baricentro de tres gaussianas anisótropas'); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d49471a2",
   "metadata": {},
   "source": [
    "El baricentro de tres elipses alargadas en distintas direcciones es un círculo (exactamente, por la simetría de rotación en $120^\\circ$ del ejemplo): el baricentro promedia las **formas**, no las densidades. La mezcla, en cambio, es una medida trimodal cuya covarianza (línea gris) es mucho mayor.\n",
    "\n",
    "**Ejercicio.** Si las $\\Sigma_i$ conmutan entre sí (por ejemplo, si todas son diagonales), la ecuación de punto fijo se resuelve a mano: $\\bar\\Sigma^{1/2}=\\sum\\lambda_i\\Sigma_i^{1/2}$. Verificarlo numéricamente y compararlo con el caso unidimensional de la sección anterior."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "624593d3",
   "metadata": {},
   "source": [
    "## 4. El criterio $\\sum\\lambda_iT_i=\\mathrm{id}$ y el algoritmo de soporte libre\n",
    "\n",
    "La proposición del capítulo dice: si $T_i$ es un mapa óptimo de $\\nu$ a $\\mu_i$ para cada $i$ y $\\sum_i\\lambda_iT_i=\\mathrm{id}$ $\\nu$-c.t.p., entonces $\\nu$ es baricentro. Esto sugiere un algoritmo cuando se busca un baricentro con soporte en $n$ puntos libres $\\{y_k\\}$ con pesos fijos $1/n$: dado el soporte actual, calcular los planes óptimos $\\pi_i$ de $\\nu$ a $\\mu_i$, definir $T_i(y_k)$ como el **baricentro condicional** $\\frac{1}{\\nu_k}\\sum_j\\pi_i(y_k,x_j)\\,x_j$, y mover cada punto a $y_k\\leftarrow\\sum_i\\lambda_iT_i(y_k)$. En un punto fijo se cumple el criterio (con los mapas $T_i$ reemplazados por baricentros condicionales, que es lo que corresponde cuando los planes no son mapas).\n",
    "\n",
    "Este algoritmo (Cuturi–Doucet, 2014) está implementado en POT como `ot.lp.free_support_barycenter`. Lo probamos con tres medidas empíricas en el plano, y verificamos el residuo $\\|\\sum\\lambda_iT_i(y_k)-y_k\\|$ al final."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfb4dceb",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_i, n_bar = 300, 300\n",
    "# tres nubes: un anillo, una gaussiana alargada, una uniforme en un cuadrado\n",
    "th = rng.uniform(0, 2*np.pi, n_i)\n",
    "X1 = np.c_[2*np.cos(th), 2*np.sin(th)] + rng.normal(0, 0.1, (n_i, 2)) + np.array([-4, 0])\n",
    "X2 = rng.multivariate_normal([4, 0], [[0.2, 0], [0, 2.0]], n_i)\n",
    "X3 = rng.uniform(-1, 1, (n_i, 2)) + np.array([0, 4])\n",
    "medidas = [X1, X2, X3]; pesos = [np.ones(n_i)/n_i]*3\n",
    "lam = np.array([1/3, 1/3, 1/3])\n",
    "\n",
    "Y0 = rng.normal(0, 1, (n_bar, 2)) + np.array([0, 1.5])\n",
    "Y, log = ot.lp.free_support_barycenter(medidas, pesos, Y0, weights=lam, numItermax=200, stopThr=1e-9, log=True)\n",
    "b = np.ones(n_bar)/n_bar\n",
    "\n",
    "# verificación del criterio con baricentros condicionales\n",
    "res = np.zeros_like(Y)\n",
    "for l, Xi, ai in zip(lam, medidas, pesos):\n",
    "    P = ot.emd(b, ai, ot.dist(Y, Xi))\n",
    "    Ti = (P @ Xi)/b[:, None]\n",
    "    res += l*Ti\n",
    "print(\"iteraciones:\", len(log['displacement_square_norms']))\n",
    "print(\"residuo máximo ||sum lam_i T_i(y_k) - y_k|| :\", np.max(np.linalg.norm(res - Y, axis=1)))\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(7.5, 6))\n",
    "for k, Xi in enumerate(medidas): ax.scatter(*Xi.T, s=8, alpha=0.5, label=fr'$\\mu_{k+1}$')\n",
    "ax.scatter(*Y0.T, s=8, color='gray', alpha=0.4, label='soporte inicial')\n",
    "ax.scatter(*Y.T, s=12, color='k', label='baricentro (soporte libre)')\n",
    "ax.axis('equal'); ax.legend(); ax.set_title(r'Algoritmo de soporte libre: iteración de $\\sum\\lambda_iT_i$'); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4e31712",
   "metadata": {},
   "source": [
    "El algoritmo converge en pocas iteraciones a una configuración que satisface el criterio con precisión de máquina. Dos advertencias: el número de puntos del soporte y sus pesos están fijos de antemano, de modo que el resultado es un baricentro **restringido** a esa clase, y el problema restringido no es convexo en las posiciones, así que distintas inicializaciones pueden dar distintos puntos fijos. Con eso en mente, es una herramienta práctica y rápida.\n",
    "\n",
    "Como control, comparamos el valor $\\mathcal F(Y)$ con el que se obtiene en las gaussianas de la Sección 3 al resolver el problema con soporte libre a partir de muestras: debe aproximar el valor exacto $\\sum\\lambda_iW_2^2(\\mu_i,\\bar\\nu)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c948fe3",
   "metadata": {},
   "outputs": [],
   "source": [
    "muestras = [rng.multivariate_normal(m, S, 400) for m, S in zip(medias, Sigmas)]\n",
    "pesos = [np.ones(400)/400]*3\n",
    "Yg = ot.lp.free_support_barycenter(muestras, pesos, rng.normal(0, 1, (400, 2)) + mbar, weights=lam, numItermax=200)\n",
    "F_libre = F_bar(Yg, np.ones(400)/400, list(zip(muestras, pesos)), lam)\n",
    "F_exacto = sum(l*W2_gauss(m, S, mbar, Sbar)**2 for l, m, S in zip(lam, medias, Sigmas))\n",
    "print(\"F(baricentro de soporte libre, sobre muestras):\", F_libre)\n",
    "print(\"F exacto (fórmula gaussiana)                  :\", F_exacto)\n",
    "print(\"media y covarianza empíricas del baricentro calculado:\\n\", Yg.mean(0), \"\\n\", np.cov(Yg.T))\n",
    "print(\"vs. exactas:\\n\", mbar, \"\\n\", Sbar)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "628f7d73",
   "metadata": {},
   "source": [
    "## 5. No unicidad\n",
    "\n",
    "Sin absoluta continuidad de alguna de las $\\mu_i$ el baricentro puede no ser único. El ejemplo de las notas: $\\mu_1=\\frac12(\\delta_{(-1,0)}+\\delta_{(1,0)})$, $\\mu_2=\\frac12(\\delta_{(0,-1)}+\\delta_{(0,1)})$, $\\lambda_1=\\lambda_2=\\frac12$. Los dos planes de transporte entre $\\mu_1$ y $\\mu_2$ (los dos emparejamientos) tienen el mismo costo, luego ambos son óptimos, y cada uno produce un baricentro distinto como punto medio de la geodésica:\n",
    "\n",
    "$$\n",
    "\\nu_a=\\tfrac12\\bigl(\\delta_{(-\\frac12,-\\frac12)}+\\delta_{(\\frac12,\\frac12)}\\bigr),\\qquad\n",
    "\\nu_b=\\tfrac12\\bigl(\\delta_{(-\\frac12,\\frac12)}+\\delta_{(\\frac12,-\\frac12)}\\bigr),\n",
    "$$\n",
    "\n",
    "y también toda combinación convexa $(1-s)\\nu_a+s\\nu_b$, porque $\\mathcal F$ es convexa. Verificamos que todas alcanzan el valor mínimo $\\frac14W_2^2(\\mu_1,\\mu_2)=\\frac14\\cdot2=\\frac12$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acfba521",
   "metadata": {},
   "outputs": [],
   "source": [
    "mu1 = (np.array([[-1., 0.], [1., 0.]]), np.array([.5, .5]))\n",
    "mu2 = (np.array([[0., -1.], [0., 1.]]), np.array([.5, .5]))\n",
    "lam = np.array([.5, .5])\n",
    "nu_a = np.array([[-.5, -.5], [.5, .5]]); nu_b = np.array([[-.5, .5], [.5, -.5]])\n",
    "\n",
    "print(\"W_2^2(mu1, mu2) =\", W2_discreta(*mu1, *mu2)**2, \"  -> mínimo teórico t(1-t)W^2 =\", 0.25*W2_discreta(*mu1, *mu2)**2)\n",
    "print(\"F(nu_a) =\", F_bar(nu_a, np.array([.5, .5]), [mu1, mu2], lam))\n",
    "print(\"F(nu_b) =\", F_bar(nu_b, np.array([.5, .5]), [mu1, mu2], lam))\n",
    "for s in [0.25, 0.5, 0.8]:\n",
    "    print(f\"F((1-s) nu_a + s nu_b), s={s} :\", F_bar(np.vstack([nu_a, nu_b]), np.array([(1-s)/2, (1-s)/2, s/2, s/2]), [mu1, mu2], lam))\n",
    "print(\"F(delta_0) =\", F_bar(np.zeros((1, 2)), np.array([1.]), [mu1, mu2], lam), \" (el baricentro euclídeo de los átomos NO es baricentro)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5700e3a0",
   "metadata": {},
   "source": [
    "Si en cambio reemplazamos $\\mu_1$ por una medida absolutamente continua (con soporte compacto), el teorema de unicidad se aplica. Ilustramos el mecanismo con el LP de soporte fijo: sobre una grilla que contiene los cuatro candidatos, calculamos el baricentro de $\\mu_1,\\mu_2$ y vemos qué devuelve el solver (uno cualquiera de los óptimos); después \"regularizamos\" $\\mu_1$ difundiendo cada átomo en un pequeño cuadrado y observamos que el baricentro se vuelve único... y distinto de $\\nu_a$ y $\\nu_b$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f21453b",
   "metadata": {},
   "outputs": [],
   "source": [
    "g = np.linspace(-1.5, 1.5, 31); GX, GY = np.meshgrid(g, g); G = np.c_[GX.ravel(), GY.ravel()]\n",
    "Mg = ot.dist(G, G)\n",
    "def en_grilla(X, w):\n",
    "    h = np.zeros(len(G))\n",
    "    for x, wi in zip(X, w): h[np.argmin(np.sum((G - x)**2, axis=1))] += wi\n",
    "    return h\n",
    "\n",
    "Amat = np.stack([en_grilla(*mu1), en_grilla(*mu2)], axis=1)\n",
    "bar = ot.lp.barycenter(Amat, Mg, weights=lam)\n",
    "print(\"Caso discreto: masa del baricentro (LP) en los cuatro candidatos:\")\n",
    "for nombre, X in [(\"nu_a\", nu_a), (\"nu_b\", nu_b)]:\n",
    "    print(f\"  {nombre}: {[round(float(bar[np.argmin(np.sum((G-x)**2,1))]), 3) for x in X]}\")\n",
    "\n",
    "# mu1 difundida: cada átomo se reparte uniformemente en un cuadrado de lado 0.4\n",
    "mu1_ac = np.zeros(len(G))\n",
    "for x in mu1[0]:\n",
    "    caja = (np.abs(G[:, 0]-x[0]) <= 0.2) & (np.abs(G[:, 1]-x[1]) <= 0.2)\n",
    "    mu1_ac[caja] += 0.5/caja.sum()\n",
    "bar_ac = ot.lp.barycenter(np.stack([mu1_ac, en_grilla(*mu2)], axis=1), Mg, weights=lam)\n",
    "\n",
    "fig, ax = plt.subplots(1, 2, figsize=(11, 5))\n",
    "for a_, b_, tit in [(ax[0], bar, r'$\\mu_1$ discreta: el LP elige uno de los baricentros'),\n",
    "                    (ax[1], bar_ac, r'$\\mu_1$ difundida: baricentro único')]:\n",
    "    a_.scatter(*G[b_ > 1e-9].T, s=3000*b_[b_ > 1e-9], color='k', label='baricentro')\n",
    "    a_.scatter(*mu1[0].T, marker='o', s=120, facecolor='none', edgecolor='C0', lw=2, label=r'$\\mu_1$ (átomos)')\n",
    "    a_.scatter(*mu2[0].T, marker='s', s=120, facecolor='none', edgecolor='C1', lw=2, label=r'$\\mu_2$')\n",
    "    a_.set_xlim(-1.5, 1.5); a_.set_ylim(-1.5, 1.5); a_.set_aspect('equal'); a_.set_title(tit)\n",
    "ax[0].legend(loc='upper center', bbox_to_anchor=(1.1, -0.08), ncol=3, fontsize=9)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92257072",
   "metadata": {},
   "source": [
    "Con $\\mu_1$ difundida, el baricentro reparte la masa en las cuatro regiones (es la interpolación de McCann del plan óptimo entre la medida difundida y $\\mu_2$, que ahora **es** único): la asimetría entre $\\nu_a$ y $\\nu_b$ desapareció porque cada cuadradito de $\\mu_1$ se parte en dos mitades que viajan a átomos distintos de $\\mu_2$.\n",
    "\n",
    "**Ejercicio.** ¿Qué pasa con la unicidad si la medida difundida es $\\mu_2$ en lugar de $\\mu_1$? ¿Y si se difunden las dos? Explicar con el teorema de unicidad del capítulo."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0101842",
   "metadata": {},
   "source": [
    "## 6. La formulación multimarginal\n",
    "\n",
    "El problema de baricentro equivale al problema de transporte **multimarginal**\n",
    "\n",
    "$$\n",
    "\\min\\Bigl\\{\\int c(x_1,\\dots,x_N)\\,d\\gamma:\\ \\gamma\\in\\Pi(\\mu_1,\\dots,\\mu_N)\\Bigr\\},\\qquad\n",
    "c(x_1,\\dots,x_N)=\\sum_i\\lambda_i\\Bigl|x_i-\\sum_j\\lambda_jx_j\\Bigr|^2 ,\n",
    "$$\n",
    "\n",
    "y si $\\gamma^*$ es óptimo, el baricentro es $\\bar\\nu=(B)_\\#\\gamma^*$ con $B(x_1,\\dots,x_N)=\\sum_j\\lambda_jx_j$. Para medidas discretas el problema multimarginal es un programa lineal con $n_1n_2\\cdots n_N$ variables: sólo es tratable para $N$ y $n_i$ chicos, pero en ese régimen se puede resolver exactamente con `scipy.optimize.linprog` y verificar las dos afirmaciones del teorema: la igualdad de los valores óptimos, y que la imagen de $\\gamma^*$ por $B$ es un baricentro.\n",
    "\n",
    "Tomamos $N=3$ medidas con $n=5$ átomos cada una ($125$ variables)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7f22816",
   "metadata": {},
   "outputs": [],
   "source": [
    "N, n = 3, 5\n",
    "Xs = [rng.normal(c, 0.7, (n, 2)) for c in [np.array([-2, 0]), np.array([2, 0]), np.array([0, 3])]]\n",
    "ws = [rng.dirichlet(np.ones(n)) for _ in range(N)]\n",
    "lam = np.array([0.5, 0.25, 0.25])\n",
    "\n",
    "# costo multimarginal c(x1,x2,x3) sobre la grilla de índices\n",
    "idx = np.stack(np.meshgrid(*[np.arange(n)]*N, indexing='ij'), axis=-1).reshape(-1, N)   # (n^N, N)\n",
    "pts = np.stack([Xs[i][idx[:, i]] for i in range(N)], axis=1)                              # (n^N, N, 2)\n",
    "B = np.einsum('i,kij->kj', lam, pts)                                                       # baricentros euclídeos\n",
    "c = np.einsum('i,ki->k', lam, np.sum((pts - B[:, None, :])**2, axis=2))\n",
    "\n",
    "# restricciones de marginales: para cada i y cada j, sum de gamma sobre los índices con idx[:, i]==j\n",
    "A_eq, b_eq = [], []\n",
    "for i in range(N):\n",
    "    for j in range(n):\n",
    "        A_eq.append((idx[:, i] == j).astype(float)); b_eq.append(ws[i][j])\n",
    "res = linprog(c, A_eq=np.array(A_eq), b_eq=np.array(b_eq), bounds=(0, None), method='highs')\n",
    "gamma = res.x\n",
    "print(\"valor óptimo multimarginal:\", res.fun, \"   (variables no nulas:\", np.sum(gamma > 1e-10), \")\")\n",
    "\n",
    "# baricentro = imagen de gamma por B (agrupando los puntos B que coinciden)\n",
    "sop = gamma > 1e-10\n",
    "nu_X, nu_w = B[sop], gamma[sop]\n",
    "medidas = list(zip(Xs, ws))\n",
    "print(\"F(B_# gamma*)              :\", F_bar(nu_X, nu_w, medidas, lam))\n",
    "\n",
    "# comparación con otros candidatos\n",
    "Y_libre = ot.lp.free_support_barycenter(Xs, ws, rng.normal(0, 1, (30, 2)), weights=lam, numItermax=300)\n",
    "print(\"F(soporte libre, 30 puntos):\", F_bar(Y_libre, np.ones(30)/30, medidas, lam))\n",
    "print(\"F(delta en la media global):\", F_bar((sum(l*(w @ X) for l, X, w in zip(lam, Xs, ws)))[None, :], np.array([1.]), medidas, lam))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "289bd958",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(7.5, 6))\n",
    "for i in range(N): ax.scatter(*Xs[i].T, s=600*ws[i], alpha=0.7, label=fr'$\\mu_{i+1}$ ($\\lambda={lam[i]}$)')\n",
    "for k in np.where(sop)[0]:\n",
    "    for i in range(N):\n",
    "        ax.plot([pts[k, i, 0], B[k, 0]], [pts[k, i, 1], B[k, 1]], color='gray', lw=6*gamma[k], alpha=0.5)\n",
    "ax.scatter(*nu_X.T, s=600*nu_w, color='k', zorder=3, label=r'baricentro $B_\\#\\gamma^*$')\n",
    "ax.axis('equal'); ax.legend(); ax.set_title('Plan multimarginal óptimo: cada terna se colapsa en su baricentro euclídeo'); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9021c972",
   "metadata": {},
   "source": [
    "El valor óptimo del programa lineal multimarginal coincide con $\\mathcal F(B_\\#\\gamma^*)$, y este valor es menor o igual que el de cualquier otro candidato, como afirma el teorema. Notar que el soporte del baricentro tiene a lo sumo $\\sum_in_i-N+1$ puntos (un vértice del politopo multimarginal), muchos menos que los $n^N$ posibles.\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. **Rigidez en $\\mathbb R$.** Para $N$ medidas empíricas en $\\mathbb R$ con $n$ átomos cada una, mostrar que el plan multimarginal óptimo es el comonótono (empareja los $k$-ésimos datos ordenados de cada medida) y que $B_\\#\\gamma^*$ es exactamente el promedio de cuantiles de la Sección 2. Verificarlo con `linprog` para $n$ chico.\n",
    "\n",
    "2. **Traslaciones.** Probar que si se traslada $\\mu_i$ por $v_i$, el baricentro se traslada por $\\sum\\lambda_iv_i$ y $\\mathcal F$ cambia en $\\sum\\lambda_i|v_i-\\bar v|^2$. Deducir que para calcular baricentros basta centrar todas las medidas.\n",
    "\n",
    "3. **Convexidad del funcional.** Elegir dos medidas $\\nu_0,\\nu_1$ (por ejemplo, dos de las gaussianas de la Sección 3) y graficar $s\\mapsto\\mathcal F((1-s)\\nu_0+s\\nu_1)$ y $s\\mapsto\\mathcal F(\\nu_s)$, donde $\\nu_s$ es la geodésica de $W_2$ entre ellas. La primera es convexa (lema del capítulo). ¿Lo es la segunda? (No tiene por qué: $W_2^2(\\mu,\\cdot)$ no es convexo a lo largo de geodésicas en general; es *semicóncavo*.)\n",
    "\n",
    "4. **Baricentros de formas.** Discretizar dos o tres figuras planas (un disco, un cuadrado, un triángulo) como medidas uniformes en una grilla y calcular su baricentro con `ot.lp.barycenter`. El tamaño del LP crece rápido con la grilla; en el último capítulo veremos la versión entrópica, que escala mucho mejor."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}