CircleCI update of dev docs (4084).

Author: Circle CI

master/_downloads/00990263cc4deb242eda1996f922d200/plot_qsw_3d.py

@@ -0,0 +1,200 @@
+# -*- coding: utf-8 -*-
+"""
+=========================================================
+Quasi-Monte Carlo Sliced Wasserstein in 3D
+=========================================================
+
+This example illustrates the Quasi-Sliced Wasserstein (QSW) and Randomized
+Quasi-Sliced Wasserstein (RQSW) sampling schemes introduced in [95], and
+compares them to the default uniform (Monte Carlo) sampling of slicing
+directions.
+
+Sliced Wasserstein (SWD) approximates the Wasserstein distance by averaging
+1D Wasserstein distances over projections onto random directions
+:math:`\\theta` drawn uniformly on the sphere. By default these directions
+are sampled purely at random (Monte Carlo), which introduces some variance
+in the estimate for a given number of projections.
+
+QSW replaces the random directions with a deterministic, low-discrepancy
+point set on the sphere (generalized spiral points), which covers the
+sphere more evenly than random sampling and reduces the approximation
+error, especially in 3D. Since QSW is deterministic it cannot directly be
+used as an unbiased estimator in stochastic settings (e.g. gradient-based
+optimization) -- RQSW addresses this by applying a random rotation to the
+same point set, which preserves both its low discrepancy and its
+unbiasedness.
+
+We first visualize the three sampling schemes on the sphere, then measure
+how fast each one converges to the true Sliced Wasserstein distance
+between two point clouds -- known here in closed form, with no
+approximation error left except from the number of projections itself.
+
+.. [95] Nguyen, K., Bariletto, N., & Ho, N. (2024). Quasi-Monte Carlo for
+    3D Sliced Wasserstein. International Conference on Learning
+    Representations (ICLR).
+.. [96] Rakhmanov, E. A., Saff, E. B., & Zhou, Y. M. (1994). Minimal
+    Discrete Energy on the Sphere. Mathematical Research Letters, 1(6),
+    647-662.
+"""
+
+# Author: Samuel Vangu <samuelvangu0@gmail.com>
+#
+# License: MIT License
+
+# sphinx_gallery_thumbnail_number = 1
+
+import numpy as np
+import matplotlib.pylab as pl
+from mpl_toolkits.mplot3d import Axes3D  # noqa: F401 (registers the 3D projection)
+
+import ot
+from ot.sliced import get_random_projections, get_projections_spiral
+
+##############################################################################
+# Visualize the three sampling schemes on the sphere
+# ----------------------------------------------------
+# We draw a few hundred directions on :math:`S^2` with each scheme:
+#
+# - ``uniform``: directions are Gaussian vectors normalized to unit norm
+#   (standard Monte Carlo sampling of the sphere).
+# - ``spiral_qmc``: deterministic generalized spiral points -- a simple,
+#   closed-form low-discrepancy point set (Rakhmanov, Saff & Zhou, 1994)
+#   [96]. The same call always returns the same points.
+# - ``randomized_spiral_qmc``: the same spiral point set, rotated by a
+#   random (3, 3) rotation matrix (drawn via QR decomposition of a
+#   Gaussian matrix). The rotation makes the estimator unbiased while
+#   keeping the points as evenly spread out as the deterministic spiral
+#   set.
+
+n_projections = 500
+d = 3
+seed = 42
+
+theta_uniform = get_random_projections(d, n_projections, seed=seed)
+theta_qsw = get_projections_spiral(d, n_projections, randomized=False)
+theta_rqsw = get_projections_spiral(d, n_projections, randomized=True, seed=seed)
+
+fig = pl.figure(1, figsize=(15, 5))
+
+schemes = [
+    (theta_uniform, "Uniform (Monte Carlo)"),
+    (theta_qsw, "QSW (deterministic spiral)"),
+    (theta_rqsw, "RQSW (randomly rotated spiral)"),
+]
+
+for i, (theta, title) in enumerate(schemes):
+    ax = fig.add_subplot(1, 3, i + 1, projection="3d")
+    ax.scatter(theta[0], theta[1], theta[2], c=theta[2], cmap="viridis", s=4, alpha=0.8)
+    ax.set_title(title)
+    ax.set_box_aspect([1, 1, 1])
+    ax.view_init(elev=20, azim=45)
+    ax.set_xticks([])
+    ax.set_yticks([])
+    ax.set_zticks([])
+
+pl.tight_layout()
+pl.show()
+
+# Notice how the uniform sample leaves visible gaps and clusters, while QSW
+# and RQSW spread the points much more evenly over the sphere -- this is
+# exactly the low-discrepancy property that reduces the error of the Sliced
+# Wasserstein estimate.
+
+##############################################################################
+# Convergence to the true Sliced Wasserstein distance
+# ------------------------------------------------------
+# We now compare how fast each sampling scheme converges to the *true*
+# SWD as the number of projections grows. To get a reference value with
+# **zero** approximation error -- not even from a finite number of
+# samples -- we build ``Xt`` as a pure translation of ``Xs`` by a fixed
+# vector :math:`\delta`: ``Xt = Xs + delta``.
+#
+# For a rigid translation, the classical 1D Wasserstein identity
+# :math:`W_2(\mu, \mu + c) = |c|` holds *exactly*, for any distribution
+# shape and any (even very small) sample size -- no law-of-large-numbers
+# argument, no Gaussian assumption, just an algebraic identity of optimal
+# transport on the line. Projected onto any direction :math:`\theta`, this
+# gives :math:`W_2(\theta_\# \mu, \theta_\# \nu) = |\theta^T \delta|`
+# exactly, and averaging the square over :math:`\theta` uniform on
+# :math:`S^{d-1}` gives the closed-form identity
+#
+# .. math::
+#     \mathcal{SWD}_2(\mu, \nu) = \frac{\|\delta\|}{\sqrt{d}}
+#
+# Because this holds regardless of ``Xs``'s shape or size, the *only*
+# remaining source of error in the experiment below is the number of
+# projections -- exactly the quantity we want to study.
+
+rng = np.random.RandomState(0)
+
+n_samples = 200
+delta = np.array([1.5, 1.0, -0.5])
+Xs = rng.uniform(-2, 2, (n_samples, d))
+Xt = Xs + delta
+
+# Exact reference: no approximation at all, at any cost.
+sw_true = np.linalg.norm(delta) / np.sqrt(d)
+
+n_proj_list = [10, 20, 50, 100, 200, 500]
+n_trials = 8
+
+errors_uniform = np.zeros((n_trials, len(n_proj_list)))
+errors_rqsw = np.zeros((n_trials, len(n_proj_list)))
+errors_qsw = np.zeros(len(n_proj_list))
+
+for j, n_proj in enumerate(n_proj_list):
+    for t in range(n_trials):
+        sw_uniform = ot.sliced_wasserstein_distance(
+            Xs, Xt, n_projections=n_proj, sampling_slices="uniform", seed=t
+        )
+        sw_rqsw = ot.sliced_wasserstein_distance(
+            Xs,
+            Xt,
+            n_projections=n_proj,
+            sampling_slices="randomized_spiral_qmc",
+            seed=t,
+        )
+        errors_uniform[t, j] = np.abs(sw_uniform - sw_true)
+        errors_rqsw[t, j] = np.abs(sw_rqsw - sw_true)
+
+    sw_qsw = ot.sliced_wasserstein_distance(
+        Xs, Xt, n_projections=n_proj, sampling_slices="spiral_qmc"
+    )
+    errors_qsw[j] = np.abs(sw_qsw - sw_true)
+
+mean_err_uniform = errors_uniform.mean(axis=0)
+std_err_uniform = errors_uniform.std(axis=0)
+mean_err_rqsw = errors_rqsw.mean(axis=0)
+std_err_rqsw = errors_rqsw.std(axis=0)
+
+pl.figure(2, figsize=(6, 5))
+pl.plot(n_proj_list, mean_err_uniform, "o-", label="Uniform (MC)")
+pl.fill_between(
+    n_proj_list,
+    mean_err_uniform - std_err_uniform,
+    mean_err_uniform + std_err_uniform,
+    alpha=0.3,
+)
+pl.plot(n_proj_list, mean_err_rqsw, "s-", label="RQSW")
+pl.fill_between(
+    n_proj_list,
+    mean_err_rqsw - std_err_rqsw,
+    mean_err_rqsw + std_err_rqsw,
+    alpha=0.3,
+)
+pl.plot(n_proj_list, errors_qsw, "^-", label="QSW (deterministic)")
+pl.xscale("log")
+pl.yscale("log")
+pl.xlabel("Number of projections")
+pl.ylabel("Absolute error to the true SWD")
+pl.title("Convergence of the Sliced Wasserstein estimate (3D)")
+pl.legend()
+pl.show()
+
+# QSW and RQSW reach a given accuracy with fewer projections than uniform
+# sampling, and RQSW keeps the estimator unbiased -- so it is a drop-in
+# replacement for uniform sampling in stochastic optimization settings
+# (e.g. Sliced Wasserstein gradient flows) where a deterministic QSW
+# estimate would not be appropriate.
+
+# %%

master/_downloads/06f3dd7c3258690ab730be0a658b5e36/plot_ssw_unif_torch.zip

@@ -0,0 +1,79 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "\n# Quasi-Monte Carlo Sliced Wasserstein in 3D\n\nThis example illustrates the Quasi-Sliced Wasserstein (QSW) and Randomized\nQuasi-Sliced Wasserstein (RQSW) sampling schemes introduced in [95], and\ncompares them to the default uniform (Monte Carlo) sampling of slicing\ndirections.\n\nSliced Wasserstein (SWD) approximates the Wasserstein distance by averaging\n1D Wasserstein distances over projections onto random directions\n$\\theta$ drawn uniformly on the sphere. By default these directions\nare sampled purely at random (Monte Carlo), which introduces some variance\nin the estimate for a given number of projections.\n\nQSW replaces the random directions with a deterministic, low-discrepancy\npoint set on the sphere (generalized spiral points), which covers the\nsphere more evenly than random sampling and reduces the approximation\nerror, especially in 3D. Since QSW is deterministic it cannot directly be\nused as an unbiased estimator in stochastic settings (e.g. gradient-based\noptimization) -- RQSW addresses this by applying a random rotation to the\nsame point set, which preserves both its low discrepancy and its\nunbiasedness.\n\nWe first visualize the three sampling schemes on the sphere, then measure\nhow fast each one converges to the true Sliced Wasserstein distance\nbetween two point clouds -- known here in closed form, with no\napproximation error left except from the number of projections itself.\n\n.. [95] Nguyen, K., Bariletto, N., & Ho, N. (2024). Quasi-Monte Carlo for\n    3D Sliced Wasserstein. International Conference on Learning\n    Representations (ICLR).\n.. [96] Rakhmanov, E. A., Saff, E. B., & Zhou, Y. M. (1994). Minimal\n    Discrete Energy on the Sphere. Mathematical Research Letters, 1(6),\n    647-662.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "# Author: Samuel Vangu <samuelvangu0@gmail.com>\n#\n# License: MIT License\n\n# sphinx_gallery_thumbnail_number = 1\n\nimport numpy as np\nimport matplotlib.pylab as pl\nfrom mpl_toolkits.mplot3d import Axes3D  # noqa: F401 (registers the 3D projection)\n\nimport ot\nfrom ot.sliced import get_random_projections, get_projections_spiral"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Visualize the three sampling schemes on the sphere\nWe draw a few hundred directions on $S^2$ with each scheme:\n\n- ``uniform``: directions are Gaussian vectors normalized to unit norm\n  (standard Monte Carlo sampling of the sphere).\n- ``spiral_qmc``: deterministic generalized spiral points -- a simple,\n  closed-form low-discrepancy point set (Rakhmanov, Saff & Zhou, 1994)\n  [96]. The same call always returns the same points.\n- ``randomized_spiral_qmc``: the same spiral point set, rotated by a\n  random (3, 3) rotation matrix (drawn via QR decomposition of a\n  Gaussian matrix). The rotation makes the estimator unbiased while\n  keeping the points as evenly spread out as the deterministic spiral\n  set.\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "n_projections = 500\nd = 3\nseed = 42\n\ntheta_uniform = get_random_projections(d, n_projections, seed=seed)\ntheta_qsw = get_projections_spiral(d, n_projections, randomized=False)\ntheta_rqsw = get_projections_spiral(d, n_projections, randomized=True, seed=seed)\n\nfig = pl.figure(1, figsize=(15, 5))\n\nschemes = [\n    (theta_uniform, \"Uniform (Monte Carlo)\"),\n    (theta_qsw, \"QSW (deterministic spiral)\"),\n    (theta_rqsw, \"RQSW (randomly rotated spiral)\"),\n]\n\nfor i, (theta, title) in enumerate(schemes):\n    ax = fig.add_subplot(1, 3, i + 1, projection=\"3d\")\n    ax.scatter(theta[0], theta[1], theta[2], c=theta[2], cmap=\"viridis\", s=4, alpha=0.8)\n    ax.set_title(title)\n    ax.set_box_aspect([1, 1, 1])\n    ax.view_init(elev=20, azim=45)\n    ax.set_xticks([])\n    ax.set_yticks([])\n    ax.set_zticks([])\n\npl.tight_layout()\npl.show()\n\n# Notice how the uniform sample leaves visible gaps and clusters, while QSW\n# and RQSW spread the points much more evenly over the sphere -- this is\n# exactly the low-discrepancy property that reduces the error of the Sliced\n# Wasserstein estimate."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Convergence to the true Sliced Wasserstein distance\nWe now compare how fast each sampling scheme converges to the *true*\nSWD as the number of projections grows. To get a reference value with\n**zero** approximation error -- not even from a finite number of\nsamples -- we build ``Xt`` as a pure translation of ``Xs`` by a fixed\nvector $\\delta$: ``Xt = Xs + delta``.\n\nFor a rigid translation, the classical 1D Wasserstein identity\n$W_2(\\mu, \\mu + c) = |c|$ holds *exactly*, for any distribution\nshape and any (even very small) sample size -- no law-of-large-numbers\nargument, no Gaussian assumption, just an algebraic identity of optimal\ntransport on the line. Projected onto any direction $\\theta$, this\ngives $W_2(\\theta_\\# \\mu, \\theta_\\# \\nu) = |\\theta^T \\delta|$\nexactly, and averaging the square over $\\theta$ uniform on\n$S^{d-1}$ gives the closed-form identity\n\n\\begin{align}\\mathcal{SWD}_2(\\mu, \\nu) = \\frac{\\|\\delta\\|}{\\sqrt{d}}\\end{align}\n\nBecause this holds regardless of ``Xs``'s shape or size, the *only*\nremaining source of error in the experiment below is the number of\nprojections -- exactly the quantity we want to study.\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "rng = np.random.RandomState(0)\n\nn_samples = 200\ndelta = np.array([1.5, 1.0, -0.5])\nXs = rng.uniform(-2, 2, (n_samples, d))\nXt = Xs + delta\n\n# Exact reference: no approximation at all, at any cost.\nsw_true = np.linalg.norm(delta) / np.sqrt(d)\n\nn_proj_list = [10, 20, 50, 100, 200, 500]\nn_trials = 8\n\nerrors_uniform = np.zeros((n_trials, len(n_proj_list)))\nerrors_rqsw = np.zeros((n_trials, len(n_proj_list)))\nerrors_qsw = np.zeros(len(n_proj_list))\n\nfor j, n_proj in enumerate(n_proj_list):\n    for t in range(n_trials):\n        sw_uniform = ot.sliced_wasserstein_distance(\n            Xs, Xt, n_projections=n_proj, sampling_slices=\"uniform\", seed=t\n        )\n        sw_rqsw = ot.sliced_wasserstein_distance(\n            Xs,\n            Xt,\n            n_projections=n_proj,\n            sampling_slices=\"randomized_spiral_qmc\",\n            seed=t,\n        )\n        errors_uniform[t, j] = np.abs(sw_uniform - sw_true)\n        errors_rqsw[t, j] = np.abs(sw_rqsw - sw_true)\n\n    sw_qsw = ot.sliced_wasserstein_distance(\n        Xs, Xt, n_projections=n_proj, sampling_slices=\"spiral_qmc\"\n    )\n    errors_qsw[j] = np.abs(sw_qsw - sw_true)\n\nmean_err_uniform = errors_uniform.mean(axis=0)\nstd_err_uniform = errors_uniform.std(axis=0)\nmean_err_rqsw = errors_rqsw.mean(axis=0)\nstd_err_rqsw = errors_rqsw.std(axis=0)\n\npl.figure(2, figsize=(6, 5))\npl.plot(n_proj_list, mean_err_uniform, \"o-\", label=\"Uniform (MC)\")\npl.fill_between(\n    n_proj_list,\n    mean_err_uniform - std_err_uniform,\n    mean_err_uniform + std_err_uniform,\n    alpha=0.3,\n)\npl.plot(n_proj_list, mean_err_rqsw, \"s-\", label=\"RQSW\")\npl.fill_between(\n    n_proj_list,\n    mean_err_rqsw - std_err_rqsw,\n    mean_err_rqsw + std_err_rqsw,\n    alpha=0.3,\n)\npl.plot(n_proj_list, errors_qsw, \"^-\", label=\"QSW (deterministic)\")\npl.xscale(\"log\")\npl.yscale(\"log\")\npl.xlabel(\"Number of projections\")\npl.ylabel(\"Absolute error to the true SWD\")\npl.title(\"Convergence of the Sliced Wasserstein estimate (3D)\")\npl.legend()\npl.show()\n\n# QSW and RQSW reach a given accuracy with fewer projections than uniform\n# sampling, and RQSW keeps the estimator unbiased -- so it is a drop-in\n# replacement for uniform sampling in stochastic optimization settings\n# (e.g. Sliced Wasserstein gradient flows) where a deterministic QSW\n# estimate would not be appropriate."
+      ]
+    }
+  ],
+  "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.14"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}