.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "auto_examples/unbalanced-partial/plot_conv_sinkhorn_ti.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_auto_examples_unbalanced-partial_plot_conv_sinkhorn_ti.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_auto_examples_unbalanced-partial_plot_conv_sinkhorn_ti.py:


===============================================================
Translation Invariant Sinkhorn for Unbalanced Optimal Transport
===============================================================

This examples illustrates the better convergence of the translation
invariance Sinkhorn algorithm proposed in [73] compared to the classical
Sinkhorn algorithm.

[73] Séjourné, T., Vialard, F. X., & Peyré, G. (2022).
Faster unbalanced optimal transport: Translation invariant sinkhorn and 1-d frank-wolfe.
In International Conference on Artificial Intelligence and Statistics (pp. 4995-5021). PMLR.

.. GENERATED FROM PYTHON SOURCE LINES 16-24

.. code-block:: Python


    # Author: Clément Bonet <clement.bonet@ensae.fr>
    # License: MIT License

    import numpy as np
    import matplotlib.pylab as pl
    import ot








.. GENERATED FROM PYTHON SOURCE LINES 25-27

Setting parameters
-------------

.. GENERATED FROM PYTHON SOURCE LINES 29-46

.. code-block:: Python


    n_iter = 50  # nb iters
    n = 40  # nb samples

    num_iter_max = 100
    n_noise = 10

    reg = 0.005
    reg_m_kl = 0.05

    mu_s = np.array([-1, -1])
    cov_s = np.array([[1, 0], [0, 1]])

    mu_t = np.array([4, 4])
    cov_t = np.array([[1, -0.8], [-0.8, 1]])









.. GENERATED FROM PYTHON SOURCE LINES 47-49

Compute entropic kl-regularized UOT with Sinkhorn and Translation Invariant Sinkhorn
-----------

.. GENERATED FROM PYTHON SOURCE LINES 49-97

.. code-block:: Python


    err_sinkhorn_uot = np.empty((n_iter, num_iter_max))
    err_sinkhorn_uot_ti = np.empty((n_iter, num_iter_max))


    for seed in range(n_iter):
        np.random.seed(seed)
        xs = ot.datasets.make_2D_samples_gauss(n, mu_s, cov_s)
        xt = ot.datasets.make_2D_samples_gauss(n, mu_t, cov_t)

        xs = np.concatenate((xs, (np.random.rand(n_noise, 2) - 4)), axis=0)
        xt = np.concatenate((xt, (np.random.rand(n_noise, 2) + 6)), axis=0)

        n = n + n_noise

        a, b = np.ones((n,)) / n, np.ones((n,)) / n  # uniform distribution on samples

        # loss matrix
        M = ot.dist(xs, xt)
        M /= M.max()

        entropic_kl_uot, log_uot = ot.unbalanced.sinkhorn_unbalanced(
            a,
            b,
            M,
            reg,
            reg_m_kl,
            reg_type="kl",
            log=True,
            numItermax=num_iter_max,
            stopThr=0,
        )
        entropic_kl_uot_ti, log_uot_ti = ot.unbalanced.sinkhorn_unbalanced(
            a,
            b,
            M,
            reg,
            reg_m_kl,
            reg_type="kl",
            method="sinkhorn_translation_invariant",
            log=True,
            numItermax=num_iter_max,
            stopThr=0,
        )

        err_sinkhorn_uot[seed] = log_uot["err"]
        err_sinkhorn_uot_ti[seed] = log_uot_ti["err"]








.. GENERATED FROM PYTHON SOURCE LINES 98-100

Plot the results
----------------

.. GENERATED FROM PYTHON SOURCE LINES 100-123

.. code-block:: Python


    mean_sinkh = np.mean(err_sinkhorn_uot, axis=0)
    std_sinkh = np.std(err_sinkhorn_uot, axis=0)

    mean_sinkh_ti = np.mean(err_sinkhorn_uot_ti, axis=0)
    std_sinkh_ti = np.std(err_sinkhorn_uot_ti, axis=0)

    absc = list(range(num_iter_max))

    pl.plot(absc, mean_sinkh, label="Sinkhorn")
    pl.fill_between(absc, mean_sinkh - 2 * std_sinkh, mean_sinkh + 2 * std_sinkh, alpha=0.5)

    pl.plot(absc, mean_sinkh_ti, label="Translation Invariant Sinkhorn")
    pl.fill_between(
        absc, mean_sinkh_ti - 2 * std_sinkh_ti, mean_sinkh_ti + 2 * std_sinkh_ti, alpha=0.5
    )

    pl.yscale("log")
    pl.legend()
    pl.xlabel("Number of Iterations")
    pl.ylabel(r"$\|u-v\|_\infty$")
    pl.grid(True)
    pl.show()



.. image-sg:: /auto_examples/unbalanced-partial/images/sphx_glr_plot_conv_sinkhorn_ti_001.png
   :alt: plot conv sinkhorn ti
   :srcset: /auto_examples/unbalanced-partial/images/sphx_glr_plot_conv_sinkhorn_ti_001.png
   :class: sphx-glr-single-img






.. rst-class:: sphx-glr-timing

   **Total running time of the script:** (0 minutes 1.998 seconds)


.. _sphx_glr_download_auto_examples_unbalanced-partial_plot_conv_sinkhorn_ti.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: plot_conv_sinkhorn_ti.ipynb <plot_conv_sinkhorn_ti.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: plot_conv_sinkhorn_ti.py <plot_conv_sinkhorn_ti.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: plot_conv_sinkhorn_ti.zip <plot_conv_sinkhorn_ti.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_