.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "auto_examples/barycenters/plot_barycenter_1D.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_barycenters_plot_barycenter_1D.py>`
        to download the full example code.

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

.. _sphx_glr_auto_examples_barycenters_plot_barycenter_1D.py:


==============================
1D Wasserstein barycenter demo
==============================

This example illustrates the computation of regularized Wasserstein Barycenter
as proposed in [3].


[3] Benamou, J. D., Carlier, G., Cuturi, M., Nenna, L., & Peyré, G. (2015).
Iterative Bregman projections for regularized transportation problems
SIAM Journal on Scientific Computing, 37(2), A1111-A1138.

.. GENERATED FROM PYTHON SOURCE LINES 16-31

.. code-block:: Python


    # Author: Remi Flamary <remi.flamary@unice.fr>
    #
    # License: MIT License

    # sphinx_gallery_thumbnail_number = 1

    import numpy as np
    import matplotlib.pyplot as plt
    import ot

    # necessary for 3d plot even if not used
    from mpl_toolkits.mplot3d import Axes3D  # noqa
    from matplotlib.collections import PolyCollection








.. GENERATED FROM PYTHON SOURCE LINES 32-34

Generate data
-------------

.. GENERATED FROM PYTHON SOURCE LINES 36-54

.. code-block:: Python


    n = 100  # nb bins

    # bin positions
    x = np.arange(n, dtype=np.float64)

    # Gaussian distributions
    a1 = ot.datasets.make_1D_gauss(n, m=20, s=5)  # m= mean, s= std
    a2 = ot.datasets.make_1D_gauss(n, m=60, s=8)

    # creating matrix A containing all distributions
    A = np.vstack((a1, a2)).T
    n_distributions = A.shape[1]

    # loss matrix + normalization
    M = ot.utils.dist0(n)
    M /= M.max()








.. GENERATED FROM PYTHON SOURCE LINES 55-57

Barycenter computation
----------------------

.. GENERATED FROM PYTHON SOURCE LINES 59-81

.. code-block:: Python


    alpha = 0.2  # 0<=alpha<=1
    weights = np.array([1 - alpha, alpha])

    # l2bary
    bary_l2 = A.dot(weights)

    # wasserstein
    reg = 1e-3
    bary_wass = ot.bregman.barycenter(A, M, reg, weights)

    f, (ax1, ax2) = plt.subplots(2, 1, tight_layout=True, num=1)
    ax1.plot(x, A, color="black")
    ax1.set_title("Distributions")

    ax2.plot(x, bary_l2, "r", label="l2")
    ax2.plot(x, bary_wass, "g", label="Wasserstein")
    ax2.set_title("Barycenters")

    plt.legend()
    plt.show()




.. image-sg:: /auto_examples/barycenters/images/sphx_glr_plot_barycenter_1D_001.png
   :alt: Distributions, Barycenters
   :srcset: /auto_examples/barycenters/images/sphx_glr_plot_barycenter_1D_001.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 82-84

Barycentric interpolation
-------------------------

.. GENERATED FROM PYTHON SOURCE LINES 86-101

.. code-block:: Python


    n_alpha = 11
    alpha_list = np.linspace(0, 1, n_alpha)


    B_l2 = np.zeros((n, n_alpha))

    B_wass = np.copy(B_l2)

    for i in range(n_alpha):
        alpha = alpha_list[i]
        weights = np.array([1 - alpha, alpha])
        B_l2[:, i] = A.dot(weights)
        B_wass[:, i] = ot.bregman.barycenter(A, M, reg, weights)








.. GENERATED FROM PYTHON SOURCE LINES 102-148

.. code-block:: Python

    plt.figure(2)

    cmap = plt.get_cmap("viridis")
    verts = []
    zs = alpha_list
    for i, z in enumerate(zs):
        ys = B_l2[:, i]
        verts.append(list(zip(x, ys)))

    ax = plt.gcf().add_subplot(projection="3d")

    poly = PolyCollection(verts, facecolors=[cmap(a) for a in alpha_list])
    poly.set_alpha(0.7)
    ax.add_collection3d(poly, zs=zs, zdir="y")
    ax.set_xlabel("x")
    ax.set_xlim3d(0, n)
    ax.set_ylabel("$\\alpha$")
    ax.set_ylim3d(0, 1)
    ax.set_zlabel("")
    ax.set_zlim3d(0, B_l2.max() * 1.01)
    plt.title("Barycenter interpolation with l2")
    plt.tight_layout()

    plt.figure(3)
    cmap = plt.get_cmap("viridis")
    verts = []
    zs = alpha_list
    for i, z in enumerate(zs):
        ys = B_wass[:, i]
        verts.append(list(zip(x, ys)))

    ax = plt.gcf().add_subplot(projection="3d")

    poly = PolyCollection(verts, facecolors=[cmap(a) for a in alpha_list])
    poly.set_alpha(0.7)
    ax.add_collection3d(poly, zs=zs, zdir="y")
    ax.set_xlabel("x")
    ax.set_xlim3d(0, n)
    ax.set_ylabel("$\\alpha$")
    ax.set_ylim3d(0, 1)
    ax.set_zlabel("")
    ax.set_zlim3d(0, B_l2.max() * 1.01)
    plt.title("Barycenter interpolation with Wasserstein")
    plt.tight_layout()

    plt.show()



.. rst-class:: sphx-glr-horizontal


    *

      .. image-sg:: /auto_examples/barycenters/images/sphx_glr_plot_barycenter_1D_002.png
         :alt: Barycenter interpolation with l2
         :srcset: /auto_examples/barycenters/images/sphx_glr_plot_barycenter_1D_002.png
         :class: sphx-glr-multi-img

    *

      .. image-sg:: /auto_examples/barycenters/images/sphx_glr_plot_barycenter_1D_003.png
         :alt: Barycenter interpolation with Wasserstein
         :srcset: /auto_examples/barycenters/images/sphx_glr_plot_barycenter_1D_003.png
         :class: sphx-glr-multi-img






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

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


.. _sphx_glr_download_auto_examples_barycenters_plot_barycenter_1D.py:

.. only:: html

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

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

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

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

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

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

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


.. only:: html

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

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