add stuff

Author: rflamary

examples/plot_quickstart_guide.py

@@ -31,11 +31,12 @@
 
 
 # %%
-# Example data
+# Data generation
 # --------------
 #
-# Data generation
-# ~~~~~~~~~~~~~~~
+# We first generate two sets of samples in 2D that 25 and 50
+# samples respectively located on circles. The weights of the samples are
+# uniform.
 
 # Problem size
 n1 = 25
@@ -64,9 +65,6 @@
 # sphinx_gallery_end_ignore
 
 # %%
-# We illustrate above the simple example of two 2D distributions with 25 and 50
-# samples respectively located on circles. The weights of the samples are
-# uniform.
 #
 # Solving exact Optimal Transport
 # -------------------------------
@@ -189,15 +187,13 @@
 #       P = ot.emd(a, b, C)
 #       loss = ot.emd2(a, b, C) # same as np.sum(P*C) but differentiable wrt a/b
 #
-# .. minigallery:: ot.emd2 ot.emd ot.solve ot.solve_sample
-#
 
 
 # %%
 # Sinkhorn and Regularized OT
 # ---------------------------
 #
-# Solve Entropic Regularized OT with Sinkhorn algorithm
+# Entropic OT with Sinkhorn algorithm
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 #
 
@@ -230,8 +226,8 @@
 # exact OT problem and the OT plan is not sparse.
 
 # %%
-# Solve the Regularized OT problem with other regularizations
-# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+# Quadratic Regularized OT
+# ~~~~~~~~~~~~~~~~~~~~~~~~~
 #
 
 # Use quadratic regularization
@@ -266,3 +262,259 @@
 # quadratic regularization is another common choice for regularized OT and
 # preserves the sparsity of the OT plan.
 #
+# Solve the Regularized OT problem with user-defined regularization
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+#
+
+
+# Define a custom regularization function
+def f(G):
+    return 0.5 * np.sum(G**2)
+
+
+def df(G):
+    return G
+
+
+P_reg = ot.solve_sample(x1, x2, a, b, reg=1e2, reg_type=(f, df)).plan
+
+
+# sphinx_gallery_start_ignore
+pl.figure(1, (3, 3))
+plot2D_samples_mat(x1, x2, P_reg)
+pl.plot(x1[:, 0], x1[:, 1], "ob", label="Source samples", **style)
+pl.plot(x2[:, 0], x2[:, 1], "or", label="Target samples", **style)
+pl.title("Custom reg plan")
+pl.show()
+# sphinx_gallery_end_ignore
+# %%
+#
+# .. note::
+#    The examples above use the new API of POT. The old API is still available
+#    and and the entropic OT plan and loss can be computed with the
+#    :func:`ot.sinkhorn` # and :func:`ot.sinkhorn2` functions as below:
+#
+#    .. code-block:: python
+#
+#      Gs = ot.sinkhorn(a, b, C, reg=1e-1)
+#      loss_sink = ot.sinkhorn2(a, b, C, reg=1e-1)
+#
+#    For quadratic regularization, the :func:`ot.smooth.smooth_ot_dual` function
+#    can be used to compute the solution of the regularized OT problem. For
+#    user-defined regularization, the :func:`ot.optim.cg` function can be used
+#    to solve the regularized OT problem with Conditional Gradient algorithm.
+#
+# Unbalanced and Partial Optimal Transport
+# ----------------------------
+#
+# Solve the Unbalanced OT problem
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+#
+# Unbalanced OT relaxes the marginal constraints and allows for the source and
+# target total weights to be different. The :func:`ot.solve_sample` function can be
+# used to solve the unbalanced OT problem by setting the marginal penalization
+# :code:`unbalanced` parameter to a positive value.
+#
+
+# Solve the unbalanced OT problem with KL penalization
+P_unb_kl = ot.solve_sample(x1, x2, a, b, unbalanced=5e-2).plan
+
+# Unbalanced with KL penalization ad KL regularization
+P_unb_kl_reg = ot.solve_sample(
+    x1, x2, a, b, unbalanced=5e-2, reg=1e-1
+).plan  # also regularized
+
+# Unbalanced with L2 penalization
+P_unb_l2 = ot.solve_sample(x1, x2, a, b, unbalanced=7e1, unbalanced_type="L2").plan
+
+# sphinx_gallery_start_ignore
+pl.figure(1, (9, 3))
+
+pl.subplot(1, 3, 1)
+plot2D_samples_mat(x1, x2, P_unb_kl)
+pl.plot(x1[:, 0], x1[:, 1], "ob", label="Source samples", **style)
+pl.plot(x2[:, 0], x2[:, 1], "or", label="Target samples", **style)
+pl.title("Unbalanced KL plan")
+
+pl.subplot(1, 3, 2)
+plot2D_samples_mat(x1, x2, P_unb_kl_reg)
+pl.plot(x1[:, 0], x1[:, 1], "ob", label="Source samples", **style)
+pl.plot(x2[:, 0], x2[:, 1], "or", label="Target samples", **style)
+pl.title("Unbalanced KL + reg plan")
+
+pl.subplot(1, 3, 3)
+plot2D_samples_mat(x1, x2, P_unb_l2)
+pl.plot(x1[:, 0], x1[:, 1], "ob", label="Source samples", **style)
+pl.plot(x2[:, 0], x2[:, 1], "or", label="Target samples", **style)
+pl.title("Unbalanced L2 plan")
+
+pl.show()
+# sphinx_gallery_end_ignore
+# %%
+# .. note::
+#    Solving the unbalanced OT problem with the old API can be done with the
+#    :func:`ot.unbalanced.sinkhorn_unbalanced` function as below:
+#
+#    .. code-block:: python
+#
+#      G_unb_kl = ot.unbalanced.sinkhorn_unbalanced(a, b, C, eps=reg, alpha=unbalanced)
+#
+# Partial Optimal Transport
+# ~~~~~~~~~~~~~~~~~~~~~~~~~
+#
+
+# Solve the Unbalanced OT problem with TV penalization (equivalent)
+P_part_pen = ot.solve_sample(x1, x2, a, b, unbalanced=3, unbalanced_type="TV").plan
+
+# Solve the Partial OT problem with mass constraints (only old API)
+P_part_const = ot.partial.partial_wasserstein(a, b, C, m=0.5)  # 50% mass transported
+
+# sphinx_gallery_start_ignore
+pl.figure(1, (6, 3))
+
+pl.subplot(1, 2, 1)
+plot2D_samples_mat(x1, x2, P_part_pen)
+pl.plot(x1[:, 0], x1[:, 1], "ob", label="Source samples", **style)
+pl.plot(x2[:, 0], x2[:, 1], "or", label="Target samples", **style)
+pl.title("Partial (Unb. TV) plan")
+
+pl.subplot(1, 2, 2)
+plot2D_samples_mat(x1, x2, P_part_const)
+pl.plot(x1[:, 0], x1[:, 1], "ob", label="Source samples", **style)
+pl.plot(x2[:, 0], x2[:, 1], "or", label="Target samples", **style)
+pl.title("Partial 50% mass plan")
+pl.show()
+
+# sphinx_gallery_end_ignore
+# %%
+#
+# Gromov-Wasserstein and Fused Gromov-Wasserstein
+# -----------------------------------------------
+#
+# Solve the Gromov-Wasserstein problem
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+#
+# The Gromov-Wasserstein distance is a similarity measure between metric
+# measure spaces. So it does not require the samples to be in the same space.
+#
+
+# Define the metric cost matrices in each spaces
+
+C1 = ot.dist(x1, x1, metric="sqeuclidean")
+C2 = ot.dist(x2, x2, metric="sqeuclidean")
+
+C1 /= C1.max()
+C2 /= C2.max()
+
+# Solve the Gromov-Wasserstein problem
+sol_gw = ot.solve_gromov(C1, C2, a=a, b=b)
+P_gw = sol_gw.plan
+loss_gw = sol_gw.value
+loss_gw_linear = sol_gw.value_linear  # linear part of loss
+loss_gw_quad = sol_gw.value_quad  # quadratic part of loss
+
+# Solve the Entropic Gromov-Wasserstein problem
+P_egw = ot.solve_gromov(C1, C2, a=a, b=b, reg=1e-2).plan
+
+# sphinx_gallery_start_ignore
+pl.figure(1, (6, 3))
+
+pl.subplot(1, 2, 1)
+plot2D_samples_mat(x1, x2, P_gw)
+pl.plot(x1[:, 0], x1[:, 1], "ob", label="Source samples", **style)
+pl.plot(x2[:, 0], x2[:, 1], "or", label="Target samples", **style)
+pl.title("GW plan")
+
+pl.subplot(1, 2, 2)
+plot2D_samples_mat(x1, x2, P_egw)
+pl.plot(x1[:, 0], x1[:, 1], "ob", label="Source samples", **style)
+pl.plot(x2[:, 0], x2[:, 1], "or", label="Target samples", **style)
+pl.title("Entropic GW plan")
+pl.show()
+# sphinx_gallery_end_ignore
+# %%
+# .. note::
+#    The Gromov-Wasserstein problem can be solved with the old API using the
+#    :func:`ot.gromov.gromov_wasserstein` function and the Entropic
+#    Gromov-Wasserstein problem can be solved with the
+#    :func:`ot.gromov.entropic_gromov_wasserstein` function.
+#
+#    .. code-block:: python
+#
+#      P_gw = ot.gromov.gromov_wasserstein(C1, C2, a, b)
+#      P_egw = ot.gromov.entropic_gromov_wasserstein(C1, C2, a, b, epsilon=reg)
+#
+#      loss_gw = ot.gromov.gromov_wasserstein2(C1, C2, a, b)
+#      loss_egw = ot.gromov.entropic_gromov_wasserstein2(C1, C2, a, b, epsilon=reg)
+#
+# Fused Gromov-Wasserstein
+# ~~~~~~~~~~~~~~~~~~~~~~~~~
+#
+
+# Cost matrix
+M = C / np.max(C)
+
+# Solve FGW problem with alpha=0.1
+P_fgw = ot.solve_gromov(C1, C2, M, a=a, b=b, alpha=0.1).plan  # C is cost across spaces
+
+# SOlve entropic FGW problem with alpha=0.1
+P_efgw = ot.solve_gromov(C1, C2, M, a=a, b=b, alpha=0.1, reg=1e-3).plan
+
+# sphinx_gallery_start_ignore
+pl.figure(1, (6, 3))
+
+pl.subplot(1, 2, 1)
+plot2D_samples_mat(x1, x2, P_fgw)
+pl.plot(x1[:, 0], x1[:, 1], "ob", label="Source samples", **style)
+pl.plot(x2[:, 0], x2[:, 1], "or", label="Target samples", **style)
+pl.title("FGW plan")
+
+pl.subplot(1, 2, 2)
+plot2D_samples_mat(x1, x2, P_efgw)
+pl.plot(x1[:, 0], x1[:, 1], "ob", label="Source samples", **style)
+pl.plot(x2[:, 0], x2[:, 1], "or", label="Target samples", **style)
+pl.title("Entropic FGW plan")
+pl.show()
+
+# sphinx_gallery_end_ignore
+# %%
+# .. note::
+#    The Fused Gromov-Wasserstein problem can be solved with the old API using
+#    the :func:`ot.gromov.fused_gromov_wasserstein` function and the Entropic
+#    Fused Gromov-Wasserstein problem can be solved with the
+#    :func:`ot.gromov.entropic_fused_gromov_wasserstein` function.
+#
+#    .. code-block:: python
+#
+#      P_fgw = ot.gromov.fused_gromov_wasserstein(C1, C2, M, a, b, alpha=0.1)
+#      P_efgw = ot.gromov.entropic_fused_gromov_wasserstein(C1, C2, M, a, b, alpha=0.1, epsilon=reg)
+#
+#      loss_fgw = ot.gromov.fused_gromov_wasserstein2(C1, C2, M, a, b, alpha=0.1)
+#      loss_efgw = ot.gromov.entropic_fused_gromov_wasserstein2(C1, C2, M, a, b, alpha=0.1, epsilon=reg)
+#
+# Unbalanced Gromov-Wasserstein
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+#
+
+# # Solve the Unbalanced Gromov-Wasserstein problem
+# P_gw_unb = ot.solve_gromov(C1, C2, a=a, b=b, unbalanced=1e-2).plan
+
+# # Solve the Unbalanced Entropic Gromov-Wasserstein problem
+# P_egw_unb = ot.solve_gromov(C1, C2, a=a, b=b, reg=1e-2, reg_type='KL', unbalanced=1e-2).plan
+
+# # sphinx_gallery_start_ignore
+# pl.figure(1, (6, 3))
+
+# pl.subplot(1, 2, 1)
+# plot2D_samples_mat(x1, x2, P_gw_unb)
+# pl.plot(x1[:, 0], x1[:, 1], "ob", label="Source samples", **style)
+# pl.plot(x2[:, 0], x2[:, 1], "or", label="Target samples", **style)
+# pl.title("Unbalanced GW plan")
+
+# pl.subplot(1, 2, 2)
+# plot2D_samples_mat(x1, x2, P_egw_unb)
+# pl.plot(x1[:, 0], x1[:, 1], "ob", label="Source samples", **style)
+# pl.plot(x2[:, 0], x2[:, 1], "or", label="Target samples", **style)
+# pl.title("Unbalanced Entropic GW plan")
+# pl.show()
+# # sphinx_gallery_end_ignore