Jets_clf_GNN_PointNet_task4a_Dey.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "Jets_clf_GNN_PointNet_task4_Dey.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sharmita01/gsoc_ml4sc_cern/blob/main/Jets_clf_GNN_PointNet_task4a_Dey.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "saWwVV2BZqs-"
      },
      "source": [
        "## Jet classification PointNet architecture: task 4, architecture I\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fxsEpzXeKlYg",
        "outputId": "e4d2af76-d164-470e-806a-8c5d4abde977"
      },
      "source": [
        "# Install required packages.\n",
        "!pip install -q torch-scatter -f https://pytorch-geometric.com/whl/torch-1.8.0+cu101.html\n",
        "!pip install -q torch-sparse -f https://pytorch-geometric.com/whl/torch-1.8.0+cu101.html\n",
        "!pip install -q torch-cluster -f https://pytorch-geometric.com/whl/torch-1.8.0+cu101.html\n",
        "!pip install -q torch-geometric\n",
        "\n",
        "# Helper functions for visualization.\n",
        "%matplotlib inline\n",
        "import matplotlib.pyplot as plt\n",
        "from mpl_toolkits.mplot3d import Axes3D\n",
        "\n",
        "def visualize_mesh(pos, face):\n",
        "    fig = plt.figure()\n",
        "    ax = fig.gca(projection='3d')\n",
        "    ax.axes.xaxis.set_ticklabels([])\n",
        "    ax.axes.yaxis.set_ticklabels([])\n",
        "    ax.axes.zaxis.set_ticklabels([])\n",
        "    ax.plot_trisurf(pos[:, 0], pos[:, 1], pos[:, 2], triangles=data.face.t(), antialiased=False)\n",
        "    plt.show()\n",
        "\n",
        "\n",
        "def visualize_points(pos, edge_index=None, index=None):\n",
        "    fig = plt.figure(figsize=(4, 4))\n",
        "    if edge_index is not None:\n",
        "        for (src, dst) in edge_index.t().tolist():\n",
        "             src = pos[src].tolist()\n",
        "             dst = pos[dst].tolist()\n",
        "             plt.plot([src[0], dst[0]], [src[1], dst[1]], linewidth=1, color='black')\n",
        "    if index is None:\n",
        "        plt.scatter(pos[:, 0], pos[:, 1], s=50, zorder=1000)\n",
        "    else:\n",
        "       mask = torch.zeros(pos.size(0), dtype=torch.bool)\n",
        "       mask[index] = True\n",
        "       plt.scatter(pos[~mask, 0], pos[~mask, 1], s=50, color='lightgray', zorder=1000)\n",
        "       plt.scatter(pos[mask, 0], pos[mask, 1], s=50, zorder=1000)\n",
        "    plt.axis('off')\n",
        "    plt.show()"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\u001b[K     |████████████████████████████████| 2.6MB 31.6MB/s \n",
            "\u001b[K     |████████████████████████████████| 1.5MB 43.5MB/s \n",
            "\u001b[K     |████████████████████████████████| 1.0MB 25.1MB/s \n",
            "\u001b[K     |████████████████████████████████| 194kB 17.9MB/s \n",
            "\u001b[K     |████████████████████████████████| 235kB 38.9MB/s \n",
            "\u001b[K     |████████████████████████████████| 2.2MB 53.2MB/s \n",
            "\u001b[K     |████████████████████████████████| 51kB 9.0MB/s \n",
            "\u001b[?25h  Building wheel for torch-geometric (setup.py) ... \u001b[?25l\u001b[?25hdone\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mPknydxiIXO3"
      },
      "source": [
        "Parts of the notebook has been adapted from the GNN PyTorch Geometric tutotials"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QO6XlgDeAgYd"
      },
      "source": [
        "##Get data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zJVSJ8llAjsE",
        "outputId": "6fcbfe7f-0faf-4c89-cb7e-aac19728c166"
      },
      "source": [
        "#mount Google drive\n",
        "from google.colab import drive\n",
        "import numpy as np\n",
        "drive.mount(\"/content/drive\")\n",
        "\n",
        "data = np.load('/content/drive/My Drive/graph_NN_datastets_jet_clf/QG_jets_1.npz')"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Mounted at /content/drive\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jusvwHIeNuRe"
      },
      "source": [
        "## Training inputs and training labels\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5SdsjLLdLH6g"
      },
      "source": [
        "data_X = data['X']\n",
        "data_y = data['y']"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1WYf0fPsUWxY",
        "outputId": "252d58df-8276-4cdb-96f8-8de6f0d90d9e"
      },
      "source": [
        "print (data_X.shape)\n",
        "print (data_y.shape)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "(100000, 134, 4)\n",
            "(100000,)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bO9T5gGOIzFH"
      },
      "source": [
        "##Constructing edge indexes for point cloud data\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 265
        },
        "id": "Euh4CL7v6Ede",
        "outputId": "94bb6cd9-9230-46c4-cd05-6327b91f0b9d"
      },
      "source": [
        "import torch\n",
        "from torch_cluster import knn_graph\n",
        "\n",
        "data_x_0_edge_idx = knn_graph(torch.tensor(data_X[0]), k=16)\n",
        "print(data_x_0_edge_idx.shape)\n",
        "visualize_points(data_X[0], edge_index=data_x_0_edge_idx)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "torch.Size([2, 2187])\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAO0AAADnCAYAAADy1tHpAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAgAElEQVR4nO2deZicVZX/P6e6unpLhzSQQEgCYU/Cvm/GsC8GEZXFFRkBxd3nmXHEcXSi/nTQUccZNxQZQUUNi6KSsAQiJKwGCAmBhBAgIQmBbJ2tu7qqq+r+/jjv7fft6qrq6q5u4A3n8zz9ZKmqd6l+v/ece+4554pzDsMw4kPizb4AwzAGhonWMGKGidYwYoaJ1jBihonWMGKGidYwYoaJ1jBihonWMGKGidYwYoaJ1jBihonWMGJGspYPT7xm1iagDRDAAe0rr52+21BcmGEYpZHBFgxMvGZWDqgr8VJ+5bXTaxoMDMMoz6Dc48DClhIsQN3Ea2a9NvhLMgyjEoOd07b18/qYQR7XMIx+GKxopcbXDcMYJIMVbX8TYausN4xhYrCibe/n9fWDPK5hGP1g0WPDiBmDTq4IhPk6oSvsgNdNsIYxvAza0hqG8eZgaYyGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETNMtIYRM0y0hhEzTLSGETOSQ3GQidfMagUuBQ4AVgAzV147fftQHNswjN6Ic66mA0y8ZtY7gNmo1W4BOoAC8K6V105/qOYrNAyjFzWJNrCwa4HWEi9vB/Zaee30HYM+gWEYfah1TntphWMkgtcNwxhCahXtAahLXIqWbQvu+IGIfFlEdqnxPIZhBNQq2hXoHLYPhe4u173xlVbgWmCLiKRF5H4RuVxE6ms8r2G8balVtDPRoFNfnCO3bf0RwNHAXwEHnA78GsiKyOsicquIvFtEpMbrMIy3DTWJNljWeT/Q5Qr5AkAhm6aQ6WTLgzf9sGvl0/cAOOfe45xrBg4EbgA2AWOA9wJ/AXIi8rKI/J+InFzLNRnGzk6t0eNeyz0u341zjg23fZOuVU+fh853fw58wDk3t9eJRcYCXwbeB0wAuoOfJiCHut73AL9wzi0b9EUaxk7GoEVbabmnkOlkzU8v21bIpncRkWnArcBnnXO3lLwIDVT9M/BBYH8gA2wFUkAb0AU8B/wNFfG6QV20YewE1OIel1/uEaFl8tSRInKEc+5B4EzghyLyuVJvd85tdc593Tl3IGqdvwZsAHZBBfwKmr31eeBVEdkmIg+KyOdFpNQasWHstNQi2rLLPYlUE8lRY0EDUDjnFgPvAD4rIt+pFHhyzqWdc993zh2GWtrPoVZ3CmrVnwLuRUX8HWCbiGwUkdki8lGLTBs7O7WItvxyTzZNbtvGrcDeIjIewDm3EjgFOAO4QUT6zXt2zuWdc9c7545HBfwRYAdwPnAi8DLwbTS4tRdwPZARkbUiMlNEzrPItLGzUYtoKy73dD7/0M3Bv/4S/rfbiC777AncISLN1Z7MKTOdc9Occ43AdGAN8EXgS6jV/x4wFbgDOBKdA+dE5CUR+ZWIHD/AezSMtxyDFm2w3PMuYHshmwbC5Z71t86gkN52DDofPVpEetxo51wH8B502ed+EdltMOd3zt3tnDvPOTcCOBlYjM55H0It8Ww0Kn0m8DBwNvCYiGRFZImIfF9EDhjUzRvGm8hQVPmM2HTX/2xLjhoruS3r6Fg6H9fdBdAJ3AR8CrjXOXdOrxOr23otcAFwjnPulZouJDzuIehS0nnA7sDrwJ3BuV5Gl5g+grrXY4A08Cw6//6lc+71obgOwxguahYtQKKhudAyaaok28aSa19H54sLaN7/ONrOuOpH7fdf/4WOZfNx2XS9cy7f5wJEvogu95znnFtS88X0PvY+wL8CF6Jz3s3o2u93nXOLRKQJ+ChwCXAsGq3eBixEl6luDDwDw3jLMCT1tIVM53xESKSaKHRnkGQKl8uSqG+g0J0pUMgnNs3+0T0dzz98bsmLEPkg8CPgIufc/JouqAwiMhr4F+BiYCIa0HoA+EGwLEXgql+FivwwoBnYCDwO/B6YWWrgMYw3kuGsp+1FIdOJK+RaV//PB0vW14rIWagwrnLO3THoi6qCYI79RdRNPhhN3ngE+F/gby74UkRkInA1Onc/GKgHXgXmATcCc9xQuCqGMQCGs562NyLsWHzfzHIvO+fmAOcCPxORT9R4XRVxznU4577tnJsMNKAu9G7An4AuEXlYRD4KrHLOXeOcO9w51wAch86Pj0cDXd0iskJEfiEixwznNRuGZzjraXufKNUELn9Opfc4554E3gl8WUS+/kassTrnup1zP3HOHYVa0isBQdd+u0VkoYh8VkTqnXNPOueuds4dELz3XNR1fhewQEQyIvKMiHxXRPYb7ms33p4MWz1tMYVsmlz7q3Uicq+InC4idaXe55xbgSZhXIha3ZLvGw6CteDfOudORi3w+9Clqe+hFnipiHxVRFqC997nnPuwc24CKuKPASuBjwMvikiHiDwuIv8WzKkNo2be0Dntmp9e5peDPFlgNRqtfRi4MxAtIjISdVe3AR9yznXxBlCus2RQ+PDPwKnACFSctwD/5ZzbVHycYN58GRr4OhqNTG8FnkQj0zc559LDfT/GzsfwR4+zaXCO9bd/c2tm9RLfdiaDWiYHeEtaQC1/AV2aeRF4Gjg0eP1859yWmi62inuhis6SInIUuhZ8FrArGpy6A/iec25VqWMHlvYTaGLJoWgJ4gbgMeBm4DaLTBvVMDTrtKmmQsvkqZLcdQJ1zSMpZNMkUk3kO7eR27zaJ1w4VAD54CcN3IVGbt+FluB1o8ssdWhNbSH4ex0q8DXAEuDR4LNPDlX0drCdJYOsqmvQLKw90CWiu4BrnXPPlTtf8LlPokkgB6EFEGuBB4H/A/5ukWmjFEOyw4Dr7qJ701pajzyX5oNPYeQx76b54FNoPfJcujetxXV3+RzlBJr434RanUY0vfAraBXPF9Gc5o2oW7wBtXb+4R2PVgt9GQ0A5UVki4gsEpHfiMjHRWSPQd7GoDpLOudWOOeudM7tiSZw3Iy60M8G13ZbqZzn4HNfcs4d6pxLoRHpu4CTgDloEGy5iPxMRI4Y5D0ZOyFDlhE1/tM3SaKhb/5/IdPJ2us/SfMBx+eTbXuR27y2LsiQakfTDE8Evo9a2C855+YEUeND0YqgM9EigA2o4PdCrXADmiq5FbXW9air2oxa7A3AC8ACNAvqQedcd7l7mPCFP/y4rmnkZyvc5u9WXjv9o9V+JyUK+zuB+cB/O+fu7eezApwDXI7e+9jgnpah7vt15dxwY+dnSETbesQ5hbYzrpJEqqnPa4XujP7FOZdINUoh2wWuwPpbZ5BZ+9wLaLDmcTRSey0a/PlX59wzPRepNbLHowK+GDgE7WTxIirUieiDLcB61M3cjlryfdDBIYla7TVorvFjwXvfAVwy4vCzR7ad8QkSqcZyt9kFjB5M8/UgXfIzaHR5CirAx4GfoHPZir+E4P4vAT4EnICuKXegRRJ/Bn7lnGsf6HUZ8WRIRNt26uVulxMvHtBngmhywXV3bUbdwmtQt/hq4KtoEsPXnXNr+1y0yGnA7cA/ULHug2YpPY8K9FB0njgSFdtKYBE6R56ENpgb0euYqSbGf+5mEslUyet1znWIyBdWXjv9hgHdaN9rr0OXhK4Cjgr+eyFaC/x/1QSjgm4dHwMuCo4xEtgCPIFGtH/jnMvUcp3GW5chmdPm2te5HotaRNlBQVvSJFDrmEKtxr8Av0RTBjcAi0XkW8UtZZxzf0fXcicBf0AF+ns0mHUaOvf9a3C8W1CLexHqqh5FkWCBtMumN3c+X37rIRFpcbnswWXfUCVlCvvTwI/R1rLPiMi/iEhDhWNsDxJCTnXO7YJOGf4bvf//RdeUXxORP4vI+6wRwM7FkIi288UFSBkLVe55ibSk2YAK7Wtopc1zwb+/ggpsH2C5iHwq2u3CObcUFe6HUct8C/BpdB74ZzSz6j9Rl3Iceq8+cl1ME7Br5pUl+NrgYgrZNJvn/PxLIrJVRBaLyG9F5Kqgq+SgqFDYPwNIB4Gob0g/OzQ459Y5577pnDvWOdeEDno3o4PZH9FGAKuCYN20wV6v8dZgSETbvP9xuFy25GvlLG0hmya35TWHWsvH0bYxs1C38VvAfcAuzrnL0If5IuAZEbkgYjleRV3N6WiSxkY0ewlgFSpQv1wkhMtHJelYNh/KTxdcx7KHlqPrt1vRueWP0EZz3SLyqmizue+LyNmD6VVVorD/GeAL6A4Nq0Tkv6sZJJxzy51z/+ycOySITJ+Cfp/vAOYG17tMRH4c1B8bMeLNntP6DKkuNDh0NDovvRGdp12FWtD/QBMuzgN+gM5NlwfvT6Aldsej66QNqDvs0IyrBlS8ErzX37AQJnT00DBuCmMunoHU1WclWZ/qSQ7RwBnowLAbcI1z7nuBOM9AEy2ORa2bD3x1opbzOTRyPNsNoodzpcJ+n0E2gGMJui7+MVTEe6Lf01J00PyFc271QK/ReOMYmujx0ecX2k6/QkoFcfzxizOkti+c7cBJbstr+Y6l8xIum74nNfagb6f2PGBWasy+27o3rGrYsWRuwWU729G2MY+iiQ9T0KjvODQvOCWppt21CH8vl9uyLtexdF7aZdMj/SWgAvV/FtBloO8C16HdK3p/KammjW2nXXFPctSeu3YuffC87IZVJ2XXLW8I3j+J0ILn0If9IeA2IgkRgUU8H5gGHA7sjaYyFoB24CU0ADUXFXNVm3D3V9hfzTGKjlePTiE+gA58u6K1xovQNNIbnHNbB3pcY/gYmjTGbOd8qW8qO38FXfrZvnCWo5CX1qOm05P2mO2CRILMmqWZxgmHJEkkMiKJ5kJ3Ng8F1t/yH/nMmmcdajEdKph1wb9HN4ybImMu+YYDQZeUgkFh0T2dFHLNufZ1dCyb74J14e+hOx40o4PAeNQiRnGoKA5Fl5j84PAF4H50/fQGdLlpf9RV7kSj2AmCABra5vWPUasVWLnjUEt3IjoA7BncSwa1oEvR6cJdwOOVloOqKewfKMH8+Z/QJbgj0YFyMxqZngncbJHpN5c3rmAgm6b9gRtd27SPlUzCKP+5rsLa6z7+q0J626vomq4veXOSasqM//RNjaWO55xDRCh0Z/JSl8wA56z63gUPicgoNFVwLFBcebMZeCUo0/MiewW1aH9FM53Wokn/7yN0r7egD/hEdD31VDQY5Jutr0KFeAdaZN8rySOIjp+LutlHA/uiFi+BZoa9glq+ecCsMstgVRX2DxQRGYdmr70b8PXHr6MFHr+p5djG4KhVtNeic62q6HxxAY0TDqVUEkY5Ctku1z73+u4di+5JoUkJoNlPjDj8bNrOuKra423ftuCO/dvn/urP6Hz5sOJToe7zyc65x/x/ish0dP54KDrfPpowU+tUVFidwZ9TnXNPRD7bhG4ydj7qeu4dXHs7Os+di1rjkjnKIjIJtcqnBOcfj3oJOXRu7TO+7gPmegsYuLyfBK4I7jOPWsrrgN/VIjIRmYKK+By0EiqBBgH/jrrS5dfNjCGh1r18NqJrjf1SyGXpemUJzfsdPeBzbX30FrbM+02f/x817WNUGwBzznVsefCml7Y9ftsONPJbHDlfDiScbk3SCxEpoMUJxxX9fxOa3P8BdECpRyuTbkaFtNg5Vyj6zAHoevEZqBB3Q0W4BhXWnWiWVMk65UCQp6E52z7wNZow8PUqGtR7CE15XIpa308F70+gUekb0KBT2dTOahCRd6CDw2lo7KFAuHnadYMJvBmVqUW0V6JpeGWTAKI453B5fT7KZR2VopBN037/L9mxeE6f1wZoadm+cPbqzff+bA/6DjRZ9KG/wDk3q/hzIvISsLdzruSuCCJyNuo+d6FW/EnUNW5Drel9wH3OuZdLfLYOdY3fhxYL7ItGv7ejlnQeOpfsb347FrXKp6KBr32Ca3GoZV+JbqmyFZ1XH4/+7pYDvwN+VG6gqJZgOvEedJA4BY3mZ9CB42+oiG3ztBqpRbQDco09fq5ZLWWK54Eg9fDTN1HNHDkQv9uxeE7xyQvonLPBOTeu1GdF5Hz0oZtcznIE1UX/QB/U+uD9nyd0pc9ALeF9aEBrrnNuQ4VjfRB1QY8kjHC/hlryu9DOkCU/HzmOAMegYj4JDXyNRcWaRcWcQAeYevop7B8oQVbXh9EKqePQQWx7cA8+Ml1V1NwIqdXS/ogqe0R5qhVtUWEBhEs2PW8BXMO4KXVjLp7RE40ud/xK4veXhqYTbkcDUq+hAaCX0LnsH4AFzrkTyh0gEMmf0Absfi35WOdcOnjtEEIRvzM49v2okOdXcIkFtVwXB587AE3FTAfHeBjNArunmvlqEPg6O7iWY9Dgng98+e85jQ5C/+qc+0d/x6wGEWlDXekLgSOCe9iEzsv/CPy+Vnf97cCw7E87WFw+hyvk6Xz+YTKrn4nuVgBh+mEdKqoRBG6u1DfSMnmqpkVKIt96zLvBFep0SalPckQ5Cqj1yaDzUx8ZridM2PDX0YlGdTeiy0+r0HncUtSKvBf4H3SuugM4uriULpibHocK+Ew0wPUkoSVeUOkBDpZmLkKDXEcTpmpuQues96EiWFnppouOeUBwvHNQyzyScKBMo8Gzh9B63/tqXfoRkQlowGw6GplOoYPlQ+juFLMtMt2XodgJ/j7nXMNQ5KS7fI41P7ucQmdtXWWkvlE7afTdqqTX6QgfSP/3bWhAaAuaeTQaHZSShCLeELwvQ9hZoxl1MVsIItuEgSl//IdQUS5HH/6no0kLIjICrZ31lnhfdGnKW+LnqijhOwwNip2GimAUOhC9gi45/QX460DEFgjrR2jGV2twLy74LtKEga9H0B5fz1Z77DLX/0n0/vcPzrEKjUxfH43qv50ZiuSKe1BXq2YK2S7a7/9FyaDTMBDNlFoT/H0UKjyfOfUa6nouR4NM30RFuwSNlO5OKOo8WuO6Cc3Y2oZa2pNRQYMOBl7kPic6i7rk7ej652p0z6F1wbUchO402Ego4PurSTUM5pQXoMGhE9Alp1RwHUtRMfwxWrvcz/GKC/uzqLeRQaPgPvC1BZ0fLwzOMXsw9b5BccMVaFbZBPT7fAG4G418Lx/oMXcGhkK0d6HRzyGh3PJOQPG8dqCv10L02N5F3oTOKRegXSVGoss4B6AP2WjUjY8WKfi9gl5CRdqJWuS9g8+MReeXrahQfbO7DCoSQauS0qjw5qHzwYVVzmcnommLZ6JruLsFx1+LRpfvBG7pL0BUobB/FjqdOBG19mOD+8iig9my4H33AA9V6/4GUfYL0b2XTkQDfn7ztL+hIn5bbJ42FKL9jHPuJ0PhHldY3rkRfTgqnaSAzh9Hog9Qssz7ywnbu33Qf/VTtOig0nvywbV0Bv/ePXgtgwq+hVDU+eD6N6KWfwWaDvlM8P5JqNXdF52/7o0Kzgsb9DtIo5ZuA+q6riScby8sfrCDINdZwPvRYNd+6KCwI/jcPLTl68PlBBYI6gq0wOPI4L97CvuDazyHMONrf3RgqkO9jNXBfT6IutjVeBFN6NKS3zxtFOGAeBvw61qXsN6qDIVoWwvdXdsS9WXbtFRNlRHeDGFQqBKb0Ie6mMlo14tJqOg2ow/OPsHrxRHqTsL5axN9C+grXaun2hGtUOJz/scPSt6FXoE+6E8Ef56OBqamoplTW4Prrguuu5lwjp1FXfkthO15Xiacb69Fg0PnoSLcM/jcejTQdi8a5Opj2YJB4BK0tvmE4JzPoYGlH0fn06K7MJwfXPNh6GA0Ah3ANqL53U+gU4J7K83FRTdPu5Jw87QWws3T/oB6DztFZHpIRLv1sVu3jTzuQpAEkqgL835zWaSuHpfPkUiWLy8t5LKQz1UT4a2GSi5yPnjdB5aOQxMXJqJCngZ8r651d2mZ/M6rE82j9sltXpPqWDY/57LpNDqnGlXm+DtQy9ZBKO4RqFiSVLbe0fay0QyqJCq6gYg+G1yLL2RoRefUD6AP/9LguqagFm9vNLfaz0mbCOfbGcIlsM7gtVHBTyM6z38ZLb74E3C3K2qXIyLnosUWU4PvYgUqoh+Wqh4KrPY0NE5yPJpHPRoVvw98PYcGvmYBS0p5AME04JPoGvWk4PPr0BLJG6lyeeytyFBEj2cXsunWRKpJRYqQXvU03RtXk9u8ms4VC9jtXV+kef9jyx6n86UnSa/4B8mRuxNU5eDKdJCogWIxZ9GH/CH04TwTSDfsfVhhj4u/MYK6+m7R7T86KOSSG/787WfSLz6xP5ogAOq6Poa6n2eiLpovB4zW6XahFu11wqWhApq55L8UP68t59L7Y3YHx+sKrt/XCaeCczdQvcj9PNnPy1eiRQmPodZpEmqxJqMu8wTU4u6Kir6RMJBH5JzdhMtOf0KDRqucc05ETgS+hLrJu6BR7T+hTd4rZkoFFU3TUUH74oxogG8l6ho/gLrY7UWfPwp1389CpxgEn7kfjUw/QUwYlnXaYjd3xFHTaTv9ipLpi4XuDIhAIU+pdVVJNdEyaSp+w+phEjRQOcOq0J3Jrf355b8ppLfvQDOd0uiD3xL8/Xl05N+KBnoORx9g73I2oS74HoSWKsomdEloBRqdXoQK5d3oQzoW/a7764hRQD0CP7d9DZ3rNQfnHYMKPEV14vZeQA4dLHYE17oWFfsm1LsYjYr8EHTu3lB0fJ+8sgmdb28Pvovx6CCwAU0FrbqwP3DFjyTM+JpC38DX8+ggdC8a+MoHnzsN7XoyDXXLc+jv6i7g5865l6q5hjeDYcmIigaUGsZNYcwlMyhXb1spg2nDn7/D6Pf+W5jtlO/GFQpsuP3/0bVqYSf6ID6LPig5+rcyFaPLlXKZK+VAV8BXDvlzZlBXczP6AHejD9hhqJXNBv/fSCgqH6negAaoXgz+3B7czz7B5w8lLAUE/S68W17KNfeBsgwqRiFMIvG9g7z1j34mSrnv0h87i/5efJP6+sg9dQV/JoJzFp+nAx0Ulkf+fBZY5PrZAykIUp1DmLRyAL0DX2tQT2AeGi1fg8YDPoxGpkej3/kSdG37+v5SRt9Ihi33eOujt7D1sVvLWi7nXE9fqUR935qDQrZLxVriNecc62d+ja5VT/f8F7rOuWuFSy5uLdNHwP1VDVVYjnKRP/1xo2KtRPFn/LV6L0OSbWOlHy8jag27Iz8OFcQIQuF3ELrUTeig2xT8u1z/LC9Cf3yfXFFPOLh00zubTFAhRn/6i8r7eyDyfn/uPOEc37v2fr79OmHK6TJUkMtKzK8nooGvd6KD3ATCwNcmdEB8EhXzHmhm2zGEm6c9hUbSb+xv4BhOhsnSapIESFnL5XJZ0quX0Lxv+VI9l+9G6kp7g4VclrXXXUHz/se5Kh5qb5V7Uex6U1dP26mXl6sa6sht2/DVtT//p9VoAOsa1Crmgu/AP2CFyI8XYh2VH1g/X4WgaqphwqHbxlz0H630dOTocuBk/a0zMpk1z/rjDaQxX6lfdKlBxV97PvKeBNXPlaNLZ0SuMYs++H5u/2Lwf97TmITGC7w1rvY80R+RVFOiZdLURPA7LXQsm9/psulSKafL0Ej4q2ivrHPRwNdBhEUfXcHr3l0fiwbvmhng5mnldmOs4j77MCxzWm8JGyceUdFyVSqKryRYfV3zlHGFknPh/vAN3MK2N2m8ASm1y4DL57p2LLp7wuY5120KgiqvAGOdc/XQk310YvBzOPrLGY9af99orhP1CAqE0eXoRL8AZCXV1Fh2bq3xAhfsj5QjzJdOEwap0oTJGKngGtrQ31UD1Vk9CMXgLV10U7R6wp0PfUDMC7yk6KqIT3gL6ncr9MtVTcExM8Gf/sHw95kAGhvGT0mOuWhGXa/fad9notQau4+6b0Fd5VdRgTehrvI41CqPCt6/PTi3DwAmUC3MA36NZqz1CKva3Rirpdbo8VnoBL8PhUwn7Q/eVNZyFbJpOl94jJZJ7ygpzkIui4hUFG6580aCYL4BW5ScpJqSZd32fDc4cDgSyVSlwcC7tDnCh9o/2N6NzEZ+/JyxCRWxn7f2iRhXmlu7Qr5rx+I5t2++5yevE0Z0/XJNS3COFOUtcVTsXehAkgvupy5ybd5l9hHi6CZqLnKffgnNu66lzpdvGH9IYcxF/9HQR1C3zXCZNc/VnJlTMYiYTbs1P72s3WXT0ev03kq10xjo7Un5gcrHDvwau/9+1gF/S409+HdjL/vBXQxwN8ZKlCzqHgD7oKNGHxdZkvU0jJtUto+w1DfSfLAK1g8cItLzy9xwx38y+n3/Trl94MuW+NUlaTn0dHYsnA2lH6Jky6SpGrEuRSKpx81lcblutj91J1sfmVkq4UMCy5FMto1NDmVkO9k2tmxhvyTqGgvpbR+mryvuBwwf4Y0OGj4YlCS0kP7HB2iiQaviL6d4vu2tXfFygBezd1uagVWJxtZFe1z6rXdHG9r7+xtzyTcLa35y2VrZsVIAABJBSURBVFKX7fQW1Q8a/S2B9aLi71TEjTj0jO7tT925NXLfxQKOTjnKCVki7y22Jomi940Drk6N3udqv/VrCfxujAPaaqZW0R5AmXpaqaunZfI0cHkNKtUlSdQlewlUgoQLLz6X66b9gV/T+fwjNB9wPJ0vPE7L5KklxVkubTKRTLHr6VfRvX5lWTe5oiiC4/oHrPWo6Wx9ZGaf95Vyr9tOv3JIEkRy7eso94t2znW0TJn2tS3zfjMLta7+p5XQ5R6BBk9ag9dGoL8nb0Gb6VvlBL0bu0cfXO/29ocXc/SBPqD5oJMOcIV8afU5V9cy6ZRDay0SqfQ7TdQ3JupG7DrYLVBrotJ1ob+T/Qd8zJquSCfUJS0tgCSCgSsRnqZSjrIr5EiOHM24T/wyFEN3BpIpcIVe2VaViuklWc+YS2bQ/sBNZRI2Kn++98GElslTey31SKqJMRfP6OWK9ViOi2f0l4rZLx3L5ru2068s/Yxn0y2v3vCZHwI/HPQJ3mAqCircHqYmKg10upvFm9PlJrfltXwh25VIpBpL/T470IDcgKh1W5CZ9E67K4mIVCWQRKqJkcdeSKKhuefLT9Q36Gcl0XOs6J9lz1nfRNvpV7DLiRfTdsZVjP/0TTSMm4Kkmmg96l1Vt7wp9VD144rRMnlqVccuh8umZf2tMyhkOnv2Fipk0xQynay/dUZNA0Kl00Z+hhQvqFLUKKieooyOZfO7ym7p4hwdS+f7paihuL/i7ymLLjf9BV1VOBxtEii7nfu5tkSqsdyctYBqaEDUZGlXXjt9+8RrZr2rkOmcL8n6AQeNitE0yKFB3W91caNWsP3Bm8oLrtQ1lXio3gjLkVn7HGt+ehktk6c6LeZ/zXUsnedcd5efy0bXhqN/Rl3bBH1dXSLvL374/Gf8PDn6f4P+1XQsm0/b6VeWflEFBb2XmfwSE/Sdm/tAWz2RuajLpuvW3zYjXy567Lq7fKS7En5Ongl+fHCuJThnN5pltRwtx5xDP5uVe41QPno88P2OhyJnOpFqcruedTUth5weuMSDw+W6e+a5w0Ehm6Zr9bMV86D7fKZE5VHlzKku1z73V107Ft3dQCiA6JfSjqYobkTnmM2EyxrFwRj/cEZFEw06RSO40bmof+iLI6VRr2ggUVNP8QDh7yu6G2Ei8tNDySW28lF5H333iRbRKG10EzX/XfREcqW+MdEyeWoiOWosJbqW5NEgme9OuRwV555oYcIENAZQiLznaWoo5PdMvGbWCDTotD/qEs8cjGBhiEQrIgVJNUmlTZkr4ZzDZdNsXzib1qOnlwvADKiLYzk6X3oy1zh+SrLcOVy+O1zqwbH+9m9tyrzyjBcWQD7R0LJg3Gd/e0wimeozwvRTXlgo8dMrOSB4X1R41SQ2+M8X6H2MchY2KjKfuZSIvO4i//aJH9FG8f79ldIYo1lLeaBJ6hubWya/M5UctadvA1QI1pu9KKG3pfPnTBGuL5dbY/bn24aumS5G63MXoGuvZ6P5xkcTNrKrQyPta9GUxYfRHRze0h0xhlS0u5zyIUYed+GAxOWcY9uCP7P1od+DSMW0x3LHdc5BvhtJpjQ1sq6+dD5zNk37A7+mbdrl5dbzaH/gxnyydbe6Mr2leoTRMH6KjLlohlRhOaJWMUtoQXwZnrd+fhmhmbAdTYbQkkbfE80a8sL3Au1CM4+2BOfwkeUWKjcGKE6lrES0KKEdTQGsQ6PUo4JzFmc2OcKsqA3BdSaC9+4S3LOvUip1fd2E9cTL0HLAhWgT+XbAl/WdjOYd+7I+v5NiNvjs82iXyZ4Cgiru9y3FkIi2ccIhhegDHF3W8SLq+beP/hbyuEKeDbd9M5pDXNaN2jL/ZtrOuLJscUH7gzeSbN2d3PZNtE37WKVsIlJj9ivtqoUL/VGrl0d/2Un04WpAH9gtUt9U1zJ56ii1HK+5jqXzcN1dL6Lzl7Vo2tsEdGQfiyZBCGFZYLHb589dLPAEYWKGT4jwubgpervR5fDH3YHOqZKoUHyyfvFnvUX2LqjPNfbJA9ngOrLBtfvrSwbn6Qh+/JywuCggih8ANqIrEk+hlvIJ4IUy9bK+c2R0yxSfRxzdldC7tjtNf+WhKIIfW+jOvFoysT/Xzea515N+6UlGnfJBkruOB1cg+9oKujfrHlKllmSiLVGjFq9xnyMZ/f6vIYkEUldf1rqVnT/d/s3tmdVLrgY+I6mm/XZ/95duLmQ7z8isXnJgx3MPrnfZ9Bj04fLVIK+idZspVLgbULfK75AXFXc0fzjq2vrXfAVPmt7F+D5JwQdYooPGQOYDvtXMOjQn9iU0inkMmopXXCpXjLdm21EXc1Nwz2vRfN21wXF8PesYqtsSxrutW4LjPIOWKi4EnqqiF1UbWnp3GtoreV80JdNvTrYadW3no67tyiquKdYMRRH8HOdcY8myu1yWzXNvYMfC3jttDCAo0feCywg6QhZISX0jLVOmZRr2PmxpZtWiIzuWPbTZZdPj0OTwOeio7svjJlA+qb9UrqqnQFiQ7vNyoXfGUan5X59Ed3onMPgBIJom6FMTfffGbcH7WglzistZWi/IbWgQZCHa1dBXxaxDhXgE6lLuhw5WfsOvShbcX8+64FjPo8JcALw8gMZt9eiAcBZakHEQOjD4xP3XCJvC3U0/26TszAx7s3JXyNOx9EEyryyhY9l8gPI5ov33iOqPLvoWl6cJd4OPLiWUvWRC17XSw+rnqvnI+724ooKLlqglCOtj29GBw88FfeZStCNENN+3P6vrRbkc7QbxW+fc68FuAkfTW5A+Z9kX1Uczo4rP4/tkvYxa8MVoRPXJwZanie5s73cDnIKm/DXTuzfUU2hXiTk7a4O2wfKGbQsS2QG+bIR4kIXm/VJFdYlvdObnZX7e6dPyutC6TUG7MnjLBaGL693lNOF8zlvy4oBLqYCQd5mhdLK/I2wPsxTtkvF3NPATFeQ4NPiyC73d7WKXu9Sg4MU5G90sq6buDSVaxOxDWCnj991dTIV9d42+1JJcUTbvuBRepCOPfU/ZtdihSkyIUi5HeMMd33m5a+XTv0Qbb09G50nRSGKKMLrbiDY/8wEi35/JL75HA0J+ech3gowupXgLnCMMPvnXvQeQJezbtBS1PCMIBTmacCeCqCChtyi70YFmHRp8ag7ucXfCgNZLaIO0P6EWbVCR1KAs8XTCXlkHBufxzdjWBfdyG+raPvV2dW2Hglot7f8Srl9Whc96khLruS6XJb1qEenljw6kYiaPjtQtaJi/h4rlWuVdcR9l9YKMLuT7Wk+/LUg9YWcIvyVInt5z0VJrrL4X8hZUWN2EtZu7UNpFLp7zdqCC9r2Nn0eTNlajQZtz0EDUoHbcK0XQW+lotGDc78K3V3DtvqHbC4R7EtW834/Rl1pEOxZ9YAZMucynntarlQNTxfOuaCZOL3EMoOdTHhXPC6hVWINGOlehLvMU1IIcTNg4rL85pw8mecvseyFFBemttRekF2UGnfOuJ+xx/CzBUohPmwvWJaejbVFOQt3PRnRQeQFNLpiJbuY1oF+0iIwLjv1OVPx7E3ab9JHgp4Nz3PV26e7/VqAW93g66voMzNJm02x/6k5aj5pOdF03WlQQ1lrOYM1P+ljDHnEEc9W6cnPV/nKEGyYc/syOxXOmo0sbR6Pu3Yno/jd7oYItzs2N/t33K4ru6CeR/xPCVEVvsfNotHUToZVchlrJJyo9/CIyCfh3ETkdHUh2RQcGv4v8t4E/DSRwE2mCFu3+v3twvZ1osPE54LfoeuegN9gyhoZa57QDEmyA2/rITNn6yExaJk+l6aCTadzniJLustQ3sctJl7Jl3k19XqumnjW3bcM2l8s2SzLV5z4L3Zl85tWlE9BIZaWE51IWNNpUPLom6wjdXm8lX0Ct5JNoY+2q5o0i0kK4leWx6Hy2HhX7s8CP0S7/VaXcBa7tCahrewLq2u6JDirRRPibqCIR3njzqEW0K+i9nWNFIllHBdfdVQewY/Eckm170bz/cSU/IyKMPO5Ctj7aq3OEk1STlK1nvWQGa6+78vpCeltzx7N/v6Bt2uXJUgMChXxdx5K5o/q+0CsJHXrPS33ygbeSPli0CF0CGdTu6RFBXYq6oweiyzG+g/9c4A7U0vUrehHZh3C7Dd91sDgR/n6GIBHeeOOpRbQz0dG+omg1LzjP9ifv9OLrtU6aa1+Hy2VLWloAhysuQpeK9awIzQeecNWOxXNw2TTrb51RJmXxG91BmZtfP41uXuVbci5Hs22eBJYOVcRTRPZAo9Zno0shvquC3yvnd+gWlP113W9F1zunUToRfg1qlW8gBonwRnUMWrRBneAFlGns5hERSCYZedyF5LZvpOPZub3mnR3L5rPr2Z8q+/lEMuWXgXxT7RHJtrF1VdSz5gEya5+r07rUd+aSu43vzm9dv3XHs39/3mU6VhBs7oxayT77ygwFQbDoXMJg0X6oS7oDdZ1vCX4e7WdXunfQe3+bPdAlpgxhIvwv0S0kH4ljIrxRHUORe3xWIZe913dOrFiNk8viSmy0NWraxxh5wkXlK3OKEi4qRoW7M4Vtj96ycOujM+cQ7ipXMul8OAgS2T+EBnYOQS1fHg3oPIF2tL+9XM6tiByEBvl8Inx0J7nNqEv+FOra3r0zJcIb1TFURfC5lslT61oOP5vGcZP7fX/xGukA1lMLgEiqScq9n0G2pRwMQeT1fWi0+Vh07liPzhufQ+eif3TO9Vm3KkqE9xtKva0T4Y3qqLWxGwCuuyuwhEJq9MRK3eeUoI9Sx7KHetILty+cRevR04GSrUJ8TnECWOey6Q8nGpq7GcIWHtUgIscCH0DnkAehwZ0MGtiZjwaLZkWjriJSLyJn0n8i/J28zRPhjeoYqiL4HFAnqSbGf+YmEqmSFrAXO5bMpfnAE/sEiLY/NQtwvoInHwlcPQNMi0Y6h7KFR4l72g0V6HnAUejyiKDBokXossgfovmyFRLhc4R7xfhE+HvezP1gjPgypKIFaJhw6I4x7//6CEmV3iUP0C7+CFLX19BH3GG/O8BjwCXOudU1X2j56xc0weC9qOD2Q9egO9Bg0cPoxkvzgi1B9iDMFjqCvonwqwgT4WdbIrwxlAy5aAE34qjpsuuZn8hLoq5sGVzZLS518y63Y/GcJcAHhyMDJ9g97UNoBtRhaHJ/AQ0WPYW6qreh7uuZwc8xhInwScJdyZeirU/uAp4219YYboZkTluEJEfuTiXBAuV3CEg10nrMBb/fvujejwzJxWgFynvRzZlPQINFKXQ9dim6TPLH4P98Ivy/oWvQPhF+I2pxf48mws+1RHjjzWI4REtu6+tdLp8TqUv27UHTPx2pMfv+fbDnFpGjCINFk9CqmSzqsj4KXIvmFJ+EJsJ/BvhK8HGfCD8PTYSf9VbaTNgwYAhEO/GaWa0jDj9bokn76Rce/4E79eNfLTVnrYKqu64HyyaXosGfowh7N21EM4H+EhzvIDRY9UHgI1givBFjhqJH1OxCNt0ajQBv/Nt/LShkuw4c8/6vjfLR4XJE5rYV9+wMgkWnoUn07yDc3LcTLadrRxP2RxMmwmcJE+GfwBLhjZ2AYekRFUSAC4Ab/7nffSZR33igKxQ+L4lEqTzlLuCn6PyyZ8lGRCagwaKzUDd2d1TU29AgkKCZQq2EifAvo03LHkBrPC0R3tjpqMU9vpRyjc80eUJ2LJ4zfvUPL9oX+FLD+Cl1Yy75Vk6SqUxQdtZjWVd99/zH0ayin8t3OQFdQknRu4Dcd5KoR3NtfUf42ZYIb7ydGJYeUYlUEyOOOPf2HYvnzEbXMedn1jw3LVHfsBW4tLt93UmdLzy659ZH/ri7y3TehVpM38nBd3CIJsL/A0uENwyg9nrakp0rXC7rdiy6+/1oFtPJ6Fror1d99/xj0SSEaEPuApot5DvLWyK8YVSgFtHOAq4v+UpdvXSuWLANrVJ5NPJKN1qn+hiWCG8Yg2JYekS5XJbmA09o3LHonqXA39ClF0uEN4whoJad4Mv2iErUN7DbuZ/7gXPuUOfcV5xzj5lgDWNoqEW0K9AIcCk60IoWwzCGmFpEO5PeO4tHqTqryTCMgTFo0a68dvp2NH1wO6HF7Qj+PWyF6IbxdmcoekQNWyG6YRh9GZJ6WsMw3jhqmdMahvEmYKI1jJhhojWMmGGiNYyYYaI1jJhhojWMmGGiNYyYYaI1jJhhojWMmGGiNYyYYaI1jJjx/wHHCWnINvzhSQAAAABJRU5ErkJggg==\n",
            "text/plain": [
              "<Figure size 288x288 with 1 Axes>"
            ]
          },
          "metadata": {
            "tags": [],
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TkV7DDBxLmUR"
      },
      "source": [
        "from torch.nn import Sequential, Linear, ReLU\n",
        "from torch_geometric.nn import MessagePassing\n",
        "\n",
        "\n",
        "class PointNetLayer(MessagePassing):\n",
        "    def __init__(self, in_channels, out_channels):\n",
        "        # Message passing with \"max\" aggregation.\n",
        "        super(PointNetLayer, self).__init__('max')\n",
        "        # the number of input features = hidden node dim + point dim, where point dim = 4, i.e., ( pt, rapidity, azimuthal angle, and pdgid.)\n",
        "        print (in_channels)\n",
        "        self.fc = Sequential(Linear(in_channels + 4, out_channels),\n",
        "                              ReLU(),\n",
        "                              Linear(out_channels, out_channels))\n",
        "        \n",
        "    def forward(self, h, pos, edge_index):\n",
        "        # propagating messages.\n",
        "        return self.propagate(edge_index, h=h, pos=pos)\n",
        "    \n",
        "    def message(self, h_j, pos_j, pos_i):\n",
        "        #print('h_j',h_j.shape)\n",
        "        #print ('pos_j',pos_j.shape)\n",
        "        #print ('pos_i',pos_i.shape)\n",
        "        # h_j : features of neighboring nodes --> shape [num_edges, in_channels]\n",
        "        # pos_j: position of neighboring nodes --> shape [num_edges, 4]\n",
        "        # pos_i: position of central nodes --> shape [num_edges, 4]\n",
        "\n",
        "        input = pos_j - pos_i  # spatial relation.\n",
        "\n",
        "        if h_j is not None:\n",
        "            input = torch.cat([h_j, input], dim=-1)\n",
        "\n",
        "        return self.fc(input)  "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Idjr2nAeLrrp"
      },
      "source": [
        "import torch\n",
        "import torch.nn.functional as F\n",
        "from torch_cluster import knn_graph\n",
        "from torch_geometric.nn import global_max_pool\n",
        "\n",
        "\n",
        "class PointNet(torch.nn.Module):\n",
        "    def __init__(self):\n",
        "        super(PointNet, self).__init__()\n",
        "\n",
        "        torch.manual_seed(42)\n",
        "        self.conv1 = PointNetLayer(4, 64)\n",
        "        self.conv2 = PointNetLayer(64, 32)\n",
        "        self.classifier = Linear(32, 2) #num_classes = 2, i.e, Quark or Gluon\n",
        "        \n",
        "    def forward(self, pos, batch):\n",
        "        # Compute the kNN graph:\n",
        "        edge_index = torch.tensor(knn_graph(pos, k=16, batch=batch, loop=True))\n",
        "\n",
        "        # 3. bipartite message passing.\n",
        "        h = self.conv1(h=pos, pos=pos, edge_index=edge_index)\n",
        "        h = h.relu()\n",
        "        h = self.conv2(h=h, pos=pos, edge_index=edge_index)\n",
        "        h = h.relu()\n",
        "\n",
        "        # 4. Global Pooling.\n",
        "        h = global_max_pool(h, batch) \n",
        "        \n",
        "        # 5. Classifier.\n",
        "        return self.classifier(h)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GTii7UeAOnrv"
      },
      "source": [
        "## Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hUbRw_BhLuXr",
        "outputId": "bcd95078-364c-42d0-8fee-c8806b0dd8f7"
      },
      "source": [
        "from torch_geometric.data import Data, DataLoader\n",
        "\n",
        "train_data_list =  [Data(pos=torch.tensor(data_X[i]), \n",
        "                         y=torch.tensor(data_y[i])) for i in range(8000)]\n",
        "print(len(train_data_list))\n",
        "test_data_list = [Data(pos=torch.tensor(data_X[i]), \n",
        "                       y=torch.tensor(data_y[i])) for i in range(8000, 10000)]\n",
        "\n",
        "train_loader = DataLoader(train_data_list, batch_size=10, shuffle=True)\n",
        "test_loader = DataLoader(test_data_list, batch_size=10)\n",
        "\n",
        "model = PointNet() #initialize model\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=0.01) #optimizer\n",
        "loss_func = torch.nn.CrossEntropyLoss()  # loss function.\n",
        "\n",
        "def train(model, optimizer, loader):\n",
        "    model.train()\n",
        "    \n",
        "    total_loss = 0\n",
        "    for data_ in loader:\n",
        "        optimizer.zero_grad()  \n",
        "        #print(data_.batch)\n",
        "        pred = model(data_.pos.float(), data_.batch)  # forward pass.\n",
        "        loss = loss_func(pred, data_.y.long())  # compute loss\n",
        "        loss.backward()  # backward pass.\n",
        "        optimizer.step()  # Update model params.\n",
        "        total_loss += loss.item() * data_.num_graphs\n",
        "\n",
        "    return total_loss / len(train_loader.dataset)\n",
        "\n",
        "\n",
        "def test(model, loader):\n",
        "    model.eval()\n",
        "    with torch.no_grad():\n",
        "\n",
        "      total_correct = 0\n",
        "      for data in loader:\n",
        "          pred = model(data.pos.float(), data.batch)\n",
        "          pred = pred.argmax(dim=-1)\n",
        "          total_correct += int((pred == data.y).sum())\n",
        "\n",
        "    return total_correct / len(loader.dataset)\n",
        "\n",
        "for epoch in range(1, 20):\n",
        "    loss = train(model, optimizer, train_loader)\n",
        "    test_acc = test(model, test_loader)\n",
        "    print(f'Epoch: {epoch:02d}, Loss: {loss:.4f}, Test Accuracy: {test_acc:.4f}')"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "8000\n",
            "4\n",
            "64\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/ipykernel_launcher.py:18: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Epoch: 01, Loss: 0.8783, Test Accuracy: 0.7445\n",
            "Epoch: 02, Loss: 0.5384, Test Accuracy: 0.7150\n",
            "Epoch: 03, Loss: 0.5584, Test Accuracy: 0.7285\n",
            "Epoch: 04, Loss: 0.5492, Test Accuracy: 0.7305\n",
            "Epoch: 05, Loss: 0.5476, Test Accuracy: 0.7415\n",
            "Epoch: 06, Loss: 0.5555, Test Accuracy: 0.7410\n",
            "Epoch: 07, Loss: 0.5548, Test Accuracy: 0.7490\n",
            "Epoch: 08, Loss: 0.5543, Test Accuracy: 0.7245\n",
            "Epoch: 09, Loss: 0.5737, Test Accuracy: 0.7165\n",
            "Epoch: 10, Loss: 0.5652, Test Accuracy: 0.7185\n",
            "Epoch: 11, Loss: 0.5671, Test Accuracy: 0.7085\n",
            "Epoch: 12, Loss: 0.5589, Test Accuracy: 0.7350\n",
            "Epoch: 13, Loss: 0.5565, Test Accuracy: 0.6935\n",
            "Epoch: 14, Loss: 0.5630, Test Accuracy: 0.7425\n",
            "Epoch: 15, Loss: 0.5462, Test Accuracy: 0.7330\n",
            "Epoch: 16, Loss: 0.5602, Test Accuracy: 0.7230\n",
            "Epoch: 17, Loss: 0.5696, Test Accuracy: 0.7365\n",
            "Epoch: 18, Loss: 0.5590, Test Accuracy: 0.7310\n",
            "Epoch: 19, Loss: 0.5705, Test Accuracy: 0.7350\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8ergUM-XVl7p"
      },
      "source": [
        "The max multiplicity, M=134 is considered as the nodes for constructing the edge indices,  where 4 is the dimensionality representation of each node, i.e.,  pt, rapidity, \n",
        "azimuthal angle, and pdgid."
      ]
    }
  ]
}