{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Discovering Subgraphs\n\nThis example shows how to match a smaller graph to a subset of a larger graph.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Author: Runzhong Wang <runzhong.wang@sjtu.edu.cn>\n#\n# License: Mulan PSL v2 License"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<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 torch # pytorch 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 = 'pytorch' # set default backend for pygmtools\n_ = torch.manual_seed(1) # fix random seed"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Generate the larger graph\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "num_nodes2 = 10\nA2 = torch.rand(num_nodes2, num_nodes2)\nA2 = (A2 + A2.t() > 1.) * (A2 + A2.t()) / 2\ntorch.diagonal(A2)[:] = 0\nn2 = torch.tensor([num_nodes2])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Generate the smaller graph\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "num_nodes1 = 5\nG2 = nx.from_numpy_array(A2.numpy())\npos2 = nx.spring_layout(G2)\npos2_t = torch.tensor([pos2[_] for _ in range(num_nodes2)])\nselected = [0] # build G1 as a cluster in visualization\nunselected = list(range(1, num_nodes2))\nwhile len(selected) < num_nodes1:\n    dist = torch.sum(torch.sum(torch.abs(pos2_t[selected].unsqueeze(1) - pos2_t[unselected].unsqueeze(0)), dim=-1), dim=0)\n    select_id = unselected[torch.argmin(dist).item()] # find the closest node from unselected\n    selected.append(select_id)\n    unselected.remove(select_id)\nselected.sort()\nA1 = A2[selected, :][:, selected]\nX_gt = torch.eye(num_nodes2)[selected, :]\nn1 = torch.tensor([num_nodes1])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Visualize the graphs\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "G1 = nx.from_numpy_array(A1.numpy())\npos1 = {_: pos2[selected[_]] for _ in range(num_nodes1)}\ncolor1 = ['#FF5733' for _ in range(num_nodes1)]\ncolor2 = ['#FF5733' if _ in selected else '#1f78b4' for _ in range(num_nodes2)]\nplt.figure(figsize=(8, 4))\nplt.subplot(1, 2, 1)\nplt.title('Subgraph 1')\nplt.gca().margins(0.4)\nnx.draw_networkx(G1, pos=pos1, node_color=color1)\nplt.subplot(1, 2, 2)\nplt.title('Graph 2')\nnx.draw_networkx(G2, pos=pos2, node_color=color2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We then show how to automatically discover the matching by graph matching.\n\n## Build affinity matrix\nTo match the larger graph and the smaller graph, we 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}$)\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=.001) # 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": [
        "Visualization of the affinity matrix. For graph matching problem with $N_1$ and $N_2$ nodes,\nthe affinity matrix has $N_1N_2\\times N_1N_2$ elements because there are $N_1^2$ and\n$N_2^2$ edges in each graph, respectively.\n\n<div class=\"alert alert-info\"><h4>Note</h4><p>The diagonal elements of the affinity matrix is empty because there is no node features in this example.</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. Visualization:\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": [
        "## Match the subgraph\nDraw the matching:\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(8, 4))\nplt.suptitle(f'RRWM Matching Result (acc={(X * X_gt).sum()/ X_gt.sum():.2f})')\nax1 = plt.subplot(1, 2, 1)\nplt.title('Subgraph 1')\nplt.gca().margins(0.4)\nnx.draw_networkx(G1, pos=pos1, node_color=color1)\nax2 = plt.subplot(1, 2, 2)\nplt.title('Graph 2')\nnx.draw_networkx(G2, pos=pos2, node_color=color2)\nfor i in range(num_nodes1):\n    j = torch.argmax(X[i]).item()\n    con = ConnectionPatch(xyA=pos1[i], xyB=pos2[j], coordsA=\"data\", coordsB=\"data\",\n                          axesA=ax1, axesB=ax2, color=\"green\" if X_gt[i,j] == 1 else \"red\")\n    plt.gca().add_artist(con)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Other solvers are also available\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.suptitle(f'IPFP Matching Result (acc={(X * X_gt).sum()/ X_gt.sum():.2f})')\nax1 = plt.subplot(1, 2, 1)\nplt.title('Subgraph 1')\nplt.gca().margins(0.4)\nnx.draw_networkx(G1, pos=pos1, node_color=color1)\nax2 = plt.subplot(1, 2, 2)\nplt.title('Graph 2')\nnx.draw_networkx(G2, pos=pos2, node_color=color2)\nfor i in range(num_nodes1):\n    j = torch.argmax(X[i]).item()\n    con = ConnectionPatch(xyA=pos1[i], xyB=pos2[j], coordsA=\"data\", coordsB=\"data\",\n                          axesA=ax1, axesB=ax2, color=\"green\" if X_gt[i,j] == 1 else \"red\")\n    plt.gca().add_artist(con)"
      ]
    },
    {
      "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.suptitle(f'SM Matching Result (acc={(X * X_gt).sum()/ X_gt.sum():.2f})')\nax1 = plt.subplot(1, 2, 1)\nplt.title('Subgraph 1')\nplt.gca().margins(0.4)\nnx.draw_networkx(G1, pos=pos1, node_color=color1)\nax2 = plt.subplot(1, 2, 2)\nplt.title('Graph 2')\nnx.draw_networkx(G2, pos=pos2, node_color=color2)\nfor i in range(num_nodes1):\n    j = torch.argmax(X[i]).item()\n    con = ConnectionPatch(xyA=pos1[i], xyB=pos2[j], coordsA=\"data\", coordsB=\"data\",\n                          axesA=ax1, axesB=ax2, color=\"green\" if X_gt[i,j] == 1 else \"red\")\n    plt.gca().add_artist(con)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### NGM neural network solver\nSee :func:`~pygmtools.neural_solvers.ngm` for the API reference.\n\n<div class=\"alert alert-info\"><h4>Note</h4><p>The NGM solvers are pretrained on a different problem setting, so their performance may seem inferior.\n    To improve their performance, you may change the way of building affinity matrices, or try finetuning\n    NGM on the new problem.</p></div>\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "with torch.set_grad_enabled(False):\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.suptitle(f'NGM Matching Result (acc={(X * X_gt).sum()/ X_gt.sum():.2f})')\nax1 = plt.subplot(1, 2, 1)\nplt.title('Subgraph 1')\nplt.gca().margins(0.4)\nnx.draw_networkx(G1, pos=pos1, node_color=color1)\nax2 = plt.subplot(1, 2, 2)\nplt.title('Graph 2')\nnx.draw_networkx(G2, pos=pos2, node_color=color2)\nfor i in range(num_nodes1):\n    j = torch.argmax(X[i]).item()\n    con = ConnectionPatch(xyA=pos1[i], xyB=pos2[j], coordsA=\"data\", coordsB=\"data\",\n                          axesA=ax1, axesB=ax2, color=\"green\" if X_gt[i,j] == 1 else \"red\")\n    plt.gca().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.9.7"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}