{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Row and column alignments with CO-Optimal Transport\n\nThis example is designed to show how to use the CO-Optimal Transport [47]_ in POT.\nCO-Optimal Transport allows to calculate the distance between two **arbitrary-size**\nmatrices, and to align their rows and columns. In this example, we consider two\nrandom matrices $X_1$ and $X_2$ defined by\n$(X_1)_{i,j} = \\cos(\\frac{i}{n_1} \\pi) + \\cos(\\frac{j}{d_1} \\pi) + \\sigma \\mathcal N(0,1)$\nand $(X_2)_{i,j} = \\cos(\\frac{i}{n_2} \\pi) + \\cos(\\frac{j}{d_2} \\pi) + \\sigma \\mathcal N(0,1)$.\n\n.. [49] Redko, I., Vayer, T., Flamary, R., and Courty, N. (2020).\n   [CO-Optimal Transport](https://proceedings.neurips.cc/paper/2020/file/cc384c68ad503482fb24e6d1e3b512ae-Paper.pdf).\n   Advances in Neural Information Processing Systems, 33.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Author: Remi Flamary <remi.flamary@unice.fr>\n#         Quang Huy Tran <quang-huy.tran@univ-ubs.fr>\n# License: MIT License\n\n# sphinx_gallery_thumbnail_number = 2\n\nfrom matplotlib.patches import ConnectionPatch\nimport matplotlib.pylab as pl\nimport numpy as np\nfrom ot.coot import co_optimal_transport as coot\nfrom ot.coot import co_optimal_transport2 as coot2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Generating two random matrices\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "n1 = 20\nn2 = 10\nd1 = 16\nd2 = 8\nsigma = 0.2\n\nX1 = (\n    np.cos(np.arange(n1) * np.pi / n1)[:, None]\n    + np.cos(np.arange(d1) * np.pi / d1)[None, :]\n    + sigma * np.random.randn(n1, d1)\n)\nX2 = (\n    np.cos(np.arange(n2) * np.pi / n2)[:, None]\n    + np.cos(np.arange(d2) * np.pi / d2)[None, :]\n    + sigma * np.random.randn(n2, d2)\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Visualizing the matrices\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "pl.figure(1, (8, 5))\npl.subplot(1, 2, 1)\npl.imshow(X1)\npl.title(\"$X_1$\")\n\npl.subplot(1, 2, 2)\npl.imshow(X2)\npl.title(\"$X_2$\")\n\npl.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Visualizing the alignments of rows and columns, and calculating the CO-Optimal Transport distance\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "pi_sample, pi_feature, log = coot(X1, X2, log=True, verbose=True)\ncoot_distance = coot2(X1, X2)\nprint(\"CO-Optimal Transport distance = {:.5f}\".format(coot_distance))\n\nfig = pl.figure(4, (9, 7))\npl.clf()\n\nax1 = pl.subplot(2, 2, 3)\npl.imshow(X1)\npl.xlabel(\"$X_1$\")\n\nax2 = pl.subplot(2, 2, 2)\nax2.yaxis.tick_right()\npl.imshow(np.transpose(X2))\npl.title(\"Transpose($X_2$)\")\nax2.xaxis.tick_top()\n\nfor i in range(n1):\n    j = np.argmax(pi_sample[i, :])\n    xyA = (d1 - 0.5, i)\n    xyB = (j, d2 - 0.5)\n    con = ConnectionPatch(\n        xyA=xyA, xyB=xyB, coordsA=ax1.transData, coordsB=ax2.transData, color=\"black\"\n    )\n    fig.add_artist(con)\n\nfor i in range(d1):\n    j = np.argmax(pi_feature[i, :])\n    xyA = (i, -0.5)\n    xyB = (-0.5, j)\n    con = ConnectionPatch(\n        xyA=xyA, xyB=xyB, coordsA=ax1.transData, coordsB=ax2.transData, color=\"blue\"\n    )\n    fig.add_artist(con)"
      ]
    }
  ],
  "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
}