CircleCI update of dev docs (3941).

Author: Circle CI

master/_downloads/06f3dd7c3258690ab730be0a658b5e36/plot_ssw_unif_torch.zip

@@ -0,0 +1,133 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "\n# Semi-discrete OT: a toy 2D problem\n\nThis example shows the :mod:`ot.semidiscrete` solver on a small 2D problem:\na uniform source on $[0, 1]^2$ and 15 random target atoms with uniform\nweights. With so few atoms the Laguerre cells can be drawn by brute force on\na grid.\n\nWe call :func:`ot.semidiscrete.solve_semidiscrete` with its default\narguments: the underlying algorithm is **Projected Averaged SGD**, and the\ndefault ``decreasing_reg=True`` adds the **DRAG** entropic-regularization\nschedule of [90]_, which improves convergence.\n\nFor the returned potential $g$ we report:\n\n- the empirical Laguerre-cell masses (mean and max absolute deviation from\n  $1/15$);\n- the semi-dual objective\n  $\\langle g, b\\rangle + \\mathbb{E}_X[\\varphi_g(X)]$ estimated by\n  Monte Carlo, where the c-transform\n  $\\varphi_g(x) = \\min_j\\big(c(x, y_j) - g_j\\big)$ is computed by\n  :func:`ot.semidiscrete.semidiscrete_c_transform`. The solver **maximises** this\n  objective.\n\n.. [90] Genans, F., Godichon-Baggioni, A., Vialard, F.-X., Wintenberger, O.\n   (2025). *Decreasing Entropic Regularization Averaged Gradient for\n   Semi-Discrete Optimal Transport.* NeurIPS 2025.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "# Author: Ferdinand Genans <genans.ferdinand@gmail.com>\n#\n# License: MIT License\n\n# sphinx_gallery_thumbnail_number = 1\n\nimport numpy as np\nimport matplotlib.pyplot as plt\n\nfrom ot.semidiscrete import (\n    solve_semidiscrete,\n    semidiscrete_atom_weights,\n    semidiscrete_c_transform,\n    semidiscrete_ot_map,\n)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Toy 2D problem\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "rng = np.random.default_rng(42)\n\n\ndef source_sampler(batch_size):\n    return rng.random((batch_size, 2))\n\n\nn_atoms = 15\ntarget_positions = 0.1 + 0.8 * np.random.default_rng(0).random((n_atoms, 2))\n\n\ndef plot_laguerre_cells(target, g, ax, title, resolution=300, alpha=0.55):\n    xs = np.linspace(0, 1, resolution)\n    ys = np.linspace(0, 1, resolution)\n    XX, YY = np.meshgrid(xs, ys)\n    grid = np.stack([XX.ravel(), YY.ravel()], axis=1)\n    labels = semidiscrete_atom_weights(target, grid, g, reg=0.0).argmax(axis=1)\n    image = labels.reshape(resolution, resolution)\n    cmap = plt.get_cmap(\"tab20\", target.shape[0])\n    ax.imshow(\n        image,\n        origin=\"lower\",\n        extent=(0, 1, 0, 1),\n        cmap=cmap,\n        alpha=alpha,\n        vmin=-0.5,\n        vmax=target.shape[0] - 0.5,\n        interpolation=\"nearest\",\n    )\n    # Target points share the colour of their Laguerre cell.\n    ax.scatter(\n        target[:, 0],\n        target[:, 1],\n        s=80,\n        c=[cmap(i) for i in range(target.shape[0])],\n        edgecolor=\"black\",\n        linewidths=1.2,\n        zorder=3,\n    )\n    ax.set_title(title)\n    ax.set_aspect(\"equal\")\n    ax.set_xlim(0, 1)\n    ax.set_ylim(0, 1)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Solve and visualise\n\nA single call to :func:`solve_semidiscrete` runs DRAG with the default\narguments (``decreasing_reg=True``). We show the initial Voronoi cells\n($g = 0$) next to the Laguerre cells at the optimum.\nWith the squared-Euclidean cost (default ``metric='sqeuclidean'``), the cost\nbetween a source point in $[0, 1]^2$ and an atom is\n$\\|x - y\\|^2 \\le 2$. We clip the potential to\n``[-max_cost, max_cost] = [-2, 2]``, the localizing set where an optimal\npotential lies ([90]_, Lemma 1), which speeds up convergence.\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "g_drag = solve_semidiscrete(\n    target_positions,\n    source_sampler,\n    max_iter=20_000,\n    batch_size=32,\n    max_cost=2.0,\n)\n\nfig, axes = plt.subplots(1, 2, figsize=(11, 5.5))\nplot_laguerre_cells(target_positions, np.zeros(n_atoms), axes[0], \"Voronoi (g = 0)\")\nplot_laguerre_cells(target_positions, g_drag, axes[1], \"Approximated OT Laguerre cells\")\nplt.tight_layout()\nplt.show()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Transport map over the Laguerre cells\n\n:func:`semidiscrete_ot_map` with ``reg=0`` is the hard Monge map: every\nsource point is sent to the atom of its Laguerre cell. Overlaying the map\n(arrows on a source grid) on the *faded* cells shows each cell's mass\ncollapsing onto its atom -- a direct illustration of the mapping function.\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "gx = np.linspace(0.04, 0.96, 14)\ngrid = np.stack([a.ravel() for a in np.meshgrid(gx, gx)], axis=1)\nmapped = semidiscrete_ot_map(target_positions, grid, g_drag, reg=0.0)\nlabels = semidiscrete_atom_weights(target_positions, grid, g_drag, reg=0.0).argmax(\n    axis=1\n)\n\ncmap = plt.get_cmap(\"tab20\", n_atoms)\nfig, ax = plt.subplots(figsize=(6.5, 6.5))\nplot_laguerre_cells(\n    target_positions, g_drag, ax, \"Approximated OT map over Laguerre cells\", alpha=0.22\n)\nax.quiver(\n    grid[:, 0],\n    grid[:, 1],\n    mapped[:, 0] - grid[:, 0],\n    mapped[:, 1] - grid[:, 1],\n    angles=\"xy\",\n    scale_units=\"xy\",\n    scale=1,\n    width=0.005,\n    headwidth=4,\n    headlength=5,\n    color=[cmap(i) for i in labels],\n    zorder=2,\n)\nplt.tight_layout()\nplt.show()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Cell masses and Monte Carlo cost\n\nAt the optimum each Laguerre cell should carry mass $1/15$. We report\nthe empirical mass error and the semi-dual objective\n\n\\begin{align}\\mathcal{S}(g) = \\langle g, b\\rangle + \\mathbb{E}_X[\\varphi_g(X)]\\end{align}\n\nestimated by Monte Carlo. The solver maximises $\\mathcal{S}$.\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "def cell_masses(target, g, sampler, n_samples=100_000):\n    labels = semidiscrete_atom_weights(target, sampler(n_samples), g, reg=0.0).argmax(\n        axis=1\n    )\n    counts = np.bincount(labels, minlength=target.shape[0])\n    return counts / n_samples\n\n\ndef mc_cost(target, g, sampler, n_samples=100_000):\n    b = np.full(target.shape[0], 1.0 / target.shape[0])\n    samples = sampler(n_samples)\n    return float(g @ b + semidiscrete_c_transform(target, samples, g, reg=0.0).mean())\n\n\ntarget_mass = 1.0 / n_atoms\nm_drag = cell_masses(target_positions, g_drag, source_sampler)\ncost_drag = mc_cost(target_positions, g_drag, source_sampler)\n\nprint(f\"Target mass per cell: {target_mass:.4f}\")\nprint(\n    f\"DRAG  \u2014  mean abs. mass error: \"\n    f\"{np.mean(np.abs(m_drag - target_mass)):.4f}\"\n    f\"   max: {np.max(np.abs(m_drag - target_mass)):.4f}\"\n    f\"   semi-dual cost (MC): {cost_drag:.5f}\"\n)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Laguerre-cell masses\n\nAt the optimum every cell carries the same mass $1/15$. The bar plot\nshows the empirical mass per cell against this ground truth (dashed line):\nevery cell sits close to the theoretical value.\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "cmap = plt.get_cmap(\"tab20\", n_atoms)\nfig, ax = plt.subplots(figsize=(7.5, 4))\nax.bar(\n    np.arange(n_atoms),\n    m_drag,\n    color=[cmap(i) for i in range(n_atoms)],\n    edgecolor=\"black\",\n    linewidth=0.6,\n)\nax.axhline(\n    target_mass,\n    ls=\"--\",\n    color=\"black\",\n    lw=1.5,\n    label=\"theoretical mass per cell at the optimum\",\n)\nax.set_ylim(0, 1.6 * target_mass)\nax.set_xticks(np.arange(n_atoms))\nax.set_xlabel(\"atom index\")\nax.set_ylabel(\"Laguerre-cell mass\")\nax.set_title(\"Approximated OT: Laguerre-cell masses\")\nax.legend()\nplt.tight_layout()\nplt.show()"
+      ]
+    }
+  ],
+  "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.12.13"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}