{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Convolutional Wasserstein Barycenter example\n\nThis example is designed to illustrate how the Convolutional Wasserstein\nBarycenter function of POT works.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Author: Nicolas Courty <ncourty@irisa.fr>\n#\n# License: MIT License\nimport os\nfrom pathlib import Path\n\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport ot"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Data preparation\n\nThe four distributions are constructed from 4 simple images\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "this_file = os.path.realpath(\"__file__\")\ndata_path = os.path.join(Path(this_file).parent.parent.parent, \"data\")\n\nf1 = 1 - plt.imread(os.path.join(data_path, \"redcross.png\"))[::2, ::2, 2]\nf2 = 1 - plt.imread(os.path.join(data_path, \"tooth.png\"))[::2, ::2, 2]\nf3 = 1 - plt.imread(os.path.join(data_path, \"heart.png\"))[::2, ::2, 2]\nf4 = 1 - plt.imread(os.path.join(data_path, \"duck.png\"))[::2, ::2, 2]\n\nf1 = f1 / np.sum(f1)\nf2 = f2 / np.sum(f2)\nf3 = f3 / np.sum(f3)\nf4 = f4 / np.sum(f4)\nA = np.array([f1, f2, f3, f4])\n\nnb_images = 5\n\n# those are the four corners coordinates that will be interpolated by bilinear\n# interpolation\nv1 = np.array((1, 0, 0, 0))\nv2 = np.array((0, 1, 0, 0))\nv3 = np.array((0, 0, 1, 0))\nv4 = np.array((0, 0, 0, 1))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Barycenter computation and visualization\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fig, axes = plt.subplots(nb_images, nb_images, figsize=(7, 7))\nplt.suptitle(\"Convolutional Wasserstein Barycenters in POT\")\ncm = \"Blues\"\n# regularization parameter\nreg = 0.004\nfor i in range(nb_images):\n    for j in range(nb_images):\n        tx = float(i) / (nb_images - 1)\n        ty = float(j) / (nb_images - 1)\n\n        # weights are constructed by bilinear interpolation\n        tmp1 = (1 - tx) * v1 + tx * v2\n        tmp2 = (1 - tx) * v3 + tx * v4\n        weights = (1 - ty) * tmp1 + ty * tmp2\n\n        if i == 0 and j == 0:\n            axes[i, j].imshow(f1, cmap=cm)\n        elif i == 0 and j == (nb_images - 1):\n            axes[i, j].imshow(f3, cmap=cm)\n        elif i == (nb_images - 1) and j == 0:\n            axes[i, j].imshow(f2, cmap=cm)\n        elif i == (nb_images - 1) and j == (nb_images - 1):\n            axes[i, j].imshow(f4, cmap=cm)\n        else:\n            # call to barycenter computation\n            axes[i, j].imshow(\n                ot.bregman.convolutional_barycenter2d(A, reg, weights), cmap=cm\n            )\n        axes[i, j].axis(\"off\")\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.10.18"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}