{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Seeded Graph Matching\n\nSeeded graph matching means some partial of the matching result is already known, and the known matching\nresults are called \"seeds\". In this example, we show how to exploit such prior with ``pygmtools``.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Author: Runzhong Wang <runzhong.wang@sjtu.edu.cn>\n#         Qi Liu <purewhite@sjtu.edu.cn>\n#\n# License: Mulan PSL v2 License"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<div class=\"alert alert-info\"><h4>Note</h4><p>How to perform seeded graph matching is still an open research problem. In this example, we show a\n    simple yet effective approach that works with ``pygmtools``.</p></div>\n\n<div class=\"alert alert-info\"><h4>Note</h4><p>The following solvers are included in this example:\n\n    * :func:`~pygmtools.classic_solvers.rrwm` (classic solver)\n\n    * :func:`~pygmtools.classic_solvers.ipfp` (classic solver)\n\n    * :func:`~pygmtools.classic_solvers.sm` (classic solver)\n\n    * :func:`~pygmtools.neural_solvers.ngm` (neural network solver)</p></div>\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import jittor as jt # jittor backend\nimport pygmtools as pygm\nimport matplotlib.pyplot as plt # for plotting\nfrom matplotlib.patches import ConnectionPatch # for plotting matching result\nimport networkx as nx # for plotting graphs\npygm.BACKEND = 'jittor' # set default backend for pygmtools\n_ = jt.set_seed(1) # fix random seed\n\njt.flags.use_cuda = jt.has_cuda"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Generate two isomorphic graphs (with seeds)\nIn this example, we assume the first three nodes are already aligned. Firstly, we generate the seed matching\nmatrix:\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "num_nodes = 10\nnum_seeds = 3\nseed_mat = jt.zeros((num_nodes, num_nodes))\nseed_mat[:num_seeds, :num_seeds] = jt.init.eye(num_seeds)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Then we generate the isomorphic graphs:\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "X_gt = seed_mat.clone()\n\nX_gt[jt.arange(num_seeds, num_nodes), jt.arange(num_seeds, num_nodes)[jt.randperm(num_nodes-num_seeds)]] = 1\nA1 = jt.rand(num_nodes, num_nodes)\nA1 = (A1 + A1.t() > 1.) * (A1 + A1.t()) / 2\n\nA1[jt.arange(A1.shape[0]), jt.arange(A1.shape[0])] = 0\nA2 = jt.matmul(jt.matmul(X_gt.t(), A1), X_gt)\nn1 = jt.Var([num_nodes])\nn2 = jt.Var([num_nodes])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Visualize the graphs and seeds\nThe seed matching matrix:\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(4, 4))\nplt.title('Seed Matching Matrix')\nplt.imshow(seed_mat.numpy(), cmap='Blues')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The blue lines denote the matching seeds.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(8, 4))\nG1 = nx.from_numpy_array(A1.numpy())\nG2 = nx.from_numpy_array(A2.numpy())\npos1 = nx.spring_layout(G1)\npos2 = nx.spring_layout(G2)\nax1 = plt.subplot(1, 2, 1)\nplt.title('Graph 1')\nnx.draw_networkx(G1, pos=pos1)\nax2 = plt.subplot(1, 2, 2)\nplt.title('Graph 2')\nnx.draw_networkx(G2, pos=pos2)\nfor i in range(num_seeds):\n    j = jt.argmax(seed_mat[i], dim=-1)[0].item()\n    con = ConnectionPatch(xyA=pos1[i], xyB=pos2[j], coordsA=\"data\", coordsB=\"data\",\n                          axesA=ax1, axesB=ax2, color=\"blue\")\n    plt.gca().add_artist(con)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Now these two graphs look dissimilar because they are not aligned. We then align these two graphs\nby graph matching.\n\n## Build affinity matrix with seed prior\nWe follow the formulation of Quadratic Assignment Problem (QAP):\n\n\\begin{align}&\\max_{\\mathbf{X}} \\ \\texttt{vec}(\\mathbf{X})^\\top \\mathbf{K} \\texttt{vec}(\\mathbf{X})\\\\\n    s.t. \\quad &\\mathbf{X} \\in \\{0, 1\\}^{n_1\\times n_2}, \\ \\mathbf{X}\\mathbf{1} = \\mathbf{1}, \\ \\mathbf{X}^\\top\\mathbf{1} \\leq \\mathbf{1}\\end{align}\n\nwhere the first step is to build the affinity matrix ($\\mathbf{K}$). We firstly build a \"standard\"\naffinity matrix:\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "conn1, edge1 = pygm.utils.dense_to_sparse(A1)\nconn2, edge2 = pygm.utils.dense_to_sparse(A2)\nimport functools\ngaussian_aff = functools.partial(pygm.utils.gaussian_aff_fn, sigma=.1) # set affinity function\nK = pygm.utils.build_aff_mat(None, edge1, conn1, None, edge2, conn2, n1, None, n2, None, edge_aff_fn=gaussian_aff)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The next step is to add the seed matching information as priors to the affinity matrix. The matching priors\nare treated as node affinities and the corresponding node affinity is added by 10 if there is an matching\nprior.\n\n<div class=\"alert alert-info\"><h4>Note</h4><p>The node affinity matrix is transposed because in the graph matching formulation followed by ``pygmtools``,\n    $\\texttt{vec}(\\mathbf{X})$ means column vectorization. The node affinity should also be column-\n    vectorized.</p></div>\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "K += jt.diag(seed_mat.t().reshape(-1) * 10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Visualization of the affinity matrix.\n\n<div class=\"alert alert-info\"><h4>Note</h4><p>In this example, the diagonal elements reflect the matching prior.</p></div>\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(4, 4))\nplt.title(f'Affinity Matrix (size: {K.shape[0]}$\\\\times${K.shape[1]})')\nplt.imshow(K.numpy(), cmap='Blues')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Solve graph matching problem by RRWM solver\nSee :func:`~pygmtools.classic_solvers.rrwm` for the API reference.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "X = pygm.rrwm(K, n1, n2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The output of RRWM is a soft matching matrix. The matching prior is well-preserved:\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(8, 4))\nplt.subplot(1, 2, 1)\nplt.title('RRWM Soft Matching Matrix')\nplt.imshow(X.numpy(), cmap='Blues')\nplt.subplot(1, 2, 2)\nplt.title('Ground Truth Matching Matrix')\nplt.imshow(X_gt.numpy(), cmap='Blues')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Get the discrete matching matrix\nHungarian algorithm is then adopted to reach a discrete matching matrix\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "X = pygm.hungarian(X)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Visualization of the discrete matching matrix:\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(8, 4))\nplt.subplot(1, 2, 1)\nplt.title(f'RRWM Matching Matrix (acc={(X * X_gt).sum()/ X_gt.sum():.2f})')\nplt.imshow(X.numpy(), cmap='Blues')\nplt.subplot(1, 2, 2)\nplt.title('Ground Truth Matching Matrix')\nplt.imshow(X_gt.numpy(), cmap='Blues')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Align the original graphs\nDraw the matching (green lines for correct matching, red lines for wrong matching, blue lines for\nseed matching):\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(8, 4))\nax1 = plt.subplot(1, 2, 1)\nplt.title('Graph 1')\nnx.draw_networkx(G1, pos=pos1)\nax2 = plt.subplot(1, 2, 2)\nplt.title('Graph 2')\nnx.draw_networkx(G2, pos=pos2)\nfor i in range(num_nodes):\n    j = jt.argmax(X[i], dim=-1)[0].item()\n    if seed_mat[i, j] == 1:\n        line_color = \"blue\"\n    elif X_gt[i, j] == 1:\n        line_color = \"green\"\n    else:\n        line_color = \"red\"\n    con = ConnectionPatch(xyA=pos1[i], xyB=pos2[j], coordsA=\"data\", coordsB=\"data\",\n                          axesA=ax1, axesB=ax2, color=line_color)\n    plt.gca().add_artist(con)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Align the nodes:\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "align_A2 = jt.matmul(jt.matmul(X, A2), X.t())\nplt.figure(figsize=(8, 4))\nax1 = plt.subplot(1, 2, 1)\nplt.title('Graph 1')\nnx.draw_networkx(G1, pos=pos1)\nax2 = plt.subplot(1, 2, 2)\nplt.title('Aligned Graph 2')\nalign_pos2 = {}\nfor i in range(num_nodes):\n    j = jt.argmax(X[i], dim=-1)[0].item()\n    align_pos2[j] = pos1[i]\n    if seed_mat[i, j] == 1:\n        line_color = \"blue\"\n    elif X_gt[i, j] == 1:\n        line_color = \"green\"\n    else:\n        line_color = \"red\"\n    con = ConnectionPatch(xyA=pos1[i], xyB=align_pos2[j], coordsA=\"data\", coordsB=\"data\",\n                          axesA=ax1, axesB=ax2, color=line_color)\n    plt.gca().add_artist(con)\nnx.draw_networkx(G2, pos=align_pos2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Other solvers are also available\nOnly the affinity matrix is modified to encode matching priors, thus other graph matching solvers are also\navailable to handle this seeded graph matching setting.\n\n### Classic IPFP solver\nSee :func:`~pygmtools.classic_solvers.ipfp` for the API reference.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "X = pygm.ipfp(K, n1, n2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Visualization of IPFP matching result:\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(8, 4))\nplt.subplot(1, 2, 1)\nplt.title(f'IPFP Matching Matrix (acc={(X * X_gt).sum()/ X_gt.sum():.2f})')\nplt.imshow(X.numpy(), cmap='Blues')\nplt.subplot(1, 2, 2)\nplt.title('Ground Truth Matching Matrix')\nplt.imshow(X_gt.numpy(), cmap='Blues')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Classic SM solver\nSee :func:`~pygmtools.classic_solvers.sm` for the API reference.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "X = pygm.sm(K, n1, n2)\nX = pygm.hungarian(X)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Visualization of SM matching result:\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(8, 4))\nplt.subplot(1, 2, 1)\nplt.title(f'SM Matching Matrix (acc={(X * X_gt).sum()/ X_gt.sum():.2f})')\nplt.imshow(X.numpy(), cmap='Blues')\nplt.subplot(1, 2, 2)\nplt.title('Ground Truth Matching Matrix')\nplt.imshow(X_gt.numpy(), cmap='Blues')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### NGM neural network solver\nSee :func:`~pygmtools.neural_solvers.ngm` for the API reference.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "with jt.no_grad():\n    X = pygm.ngm(K, n1, n2, pretrain='voc')\n    X = pygm.hungarian(X)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Visualization of NGM matching result:\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(8, 4))\nplt.subplot(1, 2, 1)\nplt.title(f'NGM Matching Matrix (acc={(X * X_gt).sum()/ X_gt.sum():.2f})')\nplt.imshow(X.numpy(), cmap='Blues')\nplt.subplot(1, 2, 2)\nplt.title('Ground Truth Matching Matrix')\nplt.imshow(X_gt.numpy(), cmap='Blues')"
      ]
    }
  ],
  "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.8.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}