4_particle_net.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Preparation\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Install the necessary packages:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [],
   "source": [
    "! pip install torch weaver-core matplotlib\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import awkward as ak\n",
    "import pandas as pd\n",
    "import uproot\n",
    "import matplotlib.pyplot as plt\n",
    "import vector\n",
    "vector.register_awkward()\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here defines some helper functions to visualize a jet:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "\n",
    "typelist = ['ch+', 'ch-', 'nh', 'ph', 'el+', 'el-', 'mu+', 'mu-']\n",
    "\n",
    "\n",
    "def make_subplot(ax, data, force_xylim=None):\n",
    "    # default plotting configuration\n",
    "    color_dict_ = {'ch': 'C0', 'nh': 'mediumpurple', 'ph': 'orange', 'el': 'red', 'mu': 'green'}\n",
    "    color_dict = color_dict_.copy()\n",
    "    color_dict.update({k + '+': color_dict_[k] for k in color_dict_})\n",
    "    color_dict.update({k + '-': color_dict_[k] for k in color_dict_})\n",
    "    if data.get('id') is None:\n",
    "        data['id'] = ['default'] * len(data['pt'])\n",
    "    if data.get('e') is None:\n",
    "        for eta, phi, pt, id, d3d in zip(data['eta'], data['phi'], data['pt'], data['id'], data['d3d']):\n",
    "            ptdraw = np.sqrt(pt) / 200\n",
    "            alpha = 0.3\n",
    "            if id in [4, 5]:\n",
    "                ax.add_patch(mpl.patches.RegularPolygon((eta, phi), 3, radius=ptdraw, clip_on=True,\n",
    "                                                        alpha=alpha, edgecolor='black', **make_color_args(id, d3d)))\n",
    "            elif id in [6, 7]:\n",
    "                ax.add_patch(mpl.patches.RegularPolygon((eta, phi), 3, radius=ptdraw, orientation=np.pi,\n",
    "                                                        clip_on=True, alpha=alpha, edgecolor='black', **make_color_args(id, d3d)))\n",
    "            elif id in [3]:\n",
    "                ax.add_patch(mpl.patches.RegularPolygon((eta, phi), 5, radius=ptdraw,\n",
    "                                                        clip_on=True, alpha=alpha, **make_color_args(id, d3d)))\n",
    "            else:\n",
    "                ax.add_patch(plt.Circle((eta, phi), ptdraw, clip_on=True, alpha=alpha, **make_color_args(id, d3d)))\n",
    "    else:\n",
    "        for eta, phi, pt, e, id, d3d in zip(data['eta'], data['phi'], data['pt'], data['e'], data['id'], data['d3d']):\n",
    "            ax.add_patch(mpl.patches.Wedge((eta, phi), pt / 600., 90, 270,\n",
    "                                           clip_on=True, alpha=alpha, **make_color_args(id, d3d)))\n",
    "            ax.add_patch(mpl.patches.Wedge((eta, phi), e / 600., 270, 90,\n",
    "                                           clip_on=True, alpha=alpha, **make_color_args(id, d3d)))\n",
    "    max_ang = force_xylim if force_xylim else max(max(abs(data['eta'])), max(abs(data['phi'])))\n",
    "    # make square plot centered at (0,0)\n",
    "    ax.set_xlim(-max_ang, max_ang)\n",
    "    ax.set_ylim(-max_ang, max_ang)\n",
    "    ax.set_xlabel(r'$\\Delta\\eta$')\n",
    "    ax.set_ylabel(r'$\\Delta\\phi$')\n",
    "    ax.set_aspect('equal')\n",
    "    return max_ang\n",
    "\n",
    "\n",
    "def make_color_args(id, d3d):\n",
    "    color = color_fader('#74c476', '#081d58', d3d)\n",
    "    if id in [2, 3]:\n",
    "        return {'edgecolor': color, 'linewidth': 1, 'fill': False}\n",
    "    else:\n",
    "        return {'facecolor': color}\n",
    "\n",
    "\n",
    "def color_fader(c1, c2, mix=0):  # fade (linear interpolate) from color c1 (at mix=0) to c2 (mix=1)\n",
    "    mix = min(1., mix)\n",
    "    c1 = np.array(mpl.colors.to_rgb(c1))\n",
    "    c2 = np.array(mpl.colors.to_rgb(c2))\n",
    "    return mpl.colors.to_hex((1 - mix) * c1 + mix * c2)\n",
    "\n",
    "\n",
    "def visualize(arrays, idx=0, title=None, ax=None):\n",
    "    data = {}\n",
    "    data['pt'] = np.hypot(arrays[idx].part_px, arrays[idx].part_py)\n",
    "    data['eta'] = arrays[idx].part_deta\n",
    "    data['phi'] = arrays[idx].part_dphi\n",
    "    data['d3d'] = np.tanh(np.hypot(arrays[idx].part_d0val, arrays[idx].part_dzval))\n",
    "    part_type = np.concatenate([\n",
    "        [(arrays[idx].part_isChargedHadron) & (arrays[idx].part_charge == 1)],\n",
    "        [(arrays[idx].part_isChargedHadron) & (arrays[idx].part_charge == -1)],\n",
    "        [arrays[idx].part_isNeutralHadron],\n",
    "        [arrays[idx].part_isPhoton],\n",
    "        [(arrays[idx].part_isElectron) & (arrays[idx].part_charge == 1)],\n",
    "        [(arrays[idx].part_isElectron) & (arrays[idx].part_charge == -1)],\n",
    "        [(arrays[idx].part_isMuon) & (arrays[idx].part_charge == 1)],\n",
    "        [(arrays[idx].part_isMuon) & (arrays[idx].part_charge == -1)],\n",
    "    ], axis=0)\n",
    "    data['id'] = np.argmax(part_type.T, axis=1)  # better\n",
    "\n",
    "    assert len(data['eta'] == data['id'])\n",
    "    if ax is None:\n",
    "        _, ax = plt.subplots(figsize=(5, 5))\n",
    "    make_subplot(ax, data, force_xylim=0.5)\n",
    "    if title:\n",
    "        ax.set_title(title)\n",
    "    return ax\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Download the dataset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def download(url, fname, chunk_size=1024):\n",
    "    '''https://gist.github.com/yanqd0/c13ed29e29432e3cf3e7c38467f42f51'''\n",
    "    import requests\n",
    "    from tqdm import tqdm\n",
    "\n",
    "    if os.path.dirname(fname):\n",
    "        os.makedirs(os.path.dirname(fname), exist_ok=True)\n",
    "\n",
    "    resp = requests.get(url, stream=True)\n",
    "    total = int(resp.headers.get('content-length', 0))\n",
    "    with open(fname, 'wb') as file, tqdm(\n",
    "        desc=fname,\n",
    "        total=total,\n",
    "        unit='iB',\n",
    "        unit_scale=True,\n",
    "        unit_divisor=1024,\n",
    "    ) as bar:\n",
    "        for data in resp.iter_content(chunk_size=chunk_size):\n",
    "            size = file.write(data)\n",
    "            bar.update(size)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "signal_file = './JetClassMini/TTBar_000.root'\n",
    "background_file = './JetClassMini/ZJetsToNuNu_000.root'\n",
    "\n",
    "if not os.path.exists(signal_file):\n",
    "    download('https://hqu.web.cern.ch/datasets/JetClassMini/TTBar_000.root', signal_file)\n",
    "if not os.path.exists(background_file):\n",
    "    download('https://hqu.web.cern.ch/datasets/JetClassMini/ZJetsToNuNu_000.root', background_file)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train a light-weight ParticleNet model\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the training of more complex networks, we are going to use [weaver](https://github.com/hqucms/weaver-core) to handles all the data loading, preprocessing, and the boilerplate PyTorch training code. \n",
    "\n",
    "To train a neural network using `weaver`, we just need to use the CLI and specify two configuration files:\n",
    "\n",
    "- A YAML _data configuration file_ describing how to process the input data.\n",
    "- A python _model configuration file_ providing the neural network module and the loss function.\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this tutorial, we are going to use the configurations from this github repo:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [],
   "source": [
    "! git clone https://github.com/hqucms/dl4hep.git -b MITP2023"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's have a look at the data config first: [particle_cloud/data/JetClassMini.yaml](https://github.com/hqucms/dl4hep/blob/MITP2023/particle_cloud/data/JetClassMini.yaml)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [],
   "source": [
    "! cat dl4hep/particle_cloud/data/JetClassMini.yaml"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And the network config: [particle_cloud/networks/example_ParticleNet.py](https://github.com/hqucms/dl4hep/blob/MITP2023/particle_cloud/networks/example_ParticleNet.py)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [],
   "source": [
    "! cat dl4hep/particle_cloud/networks/example_ParticleNet.py"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following will run the training of a light-weight ParticleNet model, and then apply the training to the test dataset. The prediction scores will be saved to two files in the current directory:\n",
    "\n",
    "- pred_particle_net_TTBar.root\n",
    "- pred_particle_net_ZJetsToNuNu.root\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [],
   "source": [
    "! weaver \\\n",
    "    --data-train \"TTBar:./JetClassMini/TTBar_*.root\" \"ZJetsToNuNu:./JetClassMini/ZJetsToNuNu_*.root\" \\\n",
    "    --data-test \"TTBar:./JetClassMini/TTBar_*.root\" \"ZJetsToNuNu:./JetClassMini/ZJetsToNuNu_*.root\" \\\n",
    "    --data-config dl4hep/particle_cloud/data/JetClassMini.yaml \\\n",
    "    --network-config dl4hep/particle_cloud/networks/example_ParticleNet.py -o conv_params \"[(7,(32,32,32)),(7,(64,64,64))]\" \\\n",
    "    --model-prefix \"training/JetClassMini/PN/{auto}/net\" \\\n",
    "    --num-workers 1 --fetch-step 1 --batch-size 512 --start-lr 1e-2 --num-epochs 20 \\\n",
    "    --optimizer ranger --log \"logs/JetClassMini_PN_{auto}.log\" --predict-output ./pred_particle_net.root \\\n",
    "    --tensorboard JetClassMini_ParticleNet \\\n",
    "    --gpus 0\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluated the performance of the trained model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import accuracy_score, roc_curve, auc\n",
    "from scipy.interpolate import interp1d\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_test_signal = uproot.open(\"pred_particle_net_TTBar.root\")[\"Events\"].arrays(library=\"pd\")\n",
    "df_test_background = uproot.open(\"pred_particle_net_ZJetsToNuNu.root\")[\"Events\"].arrays(library=\"pd\")\n",
    "df_test = pd.concat([df_test_signal, df_test_background])\n",
    "df_test\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Truth label: Top jet=1, q/g jet=0\n",
    "y_true = df_test[\"label_Tbqq\"]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prediction_probs = df_test[[\"score_label_QCD\", \"score_label_Tbqq\"]].values\n",
    "prediction_probs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(6, 6))\n",
    "plt.hist([prediction_probs[y_true == 1, 1], prediction_probs[y_true == 0, 1]],\n",
    "         label=[\"Top\", \"q/g\"], bins=50, histtype=\"step\", density=True)\n",
    "plt.yscale(\"log\")\n",
    "plt.xlabel(\"Prediction probability\")\n",
    "plt.legend()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fpr = epsilon_B, tpr = epsilon_S\n",
    "fpr, tpr, thresholds = roc_curve(y_true=y_true, y_score=prediction_probs[:, 1])\n",
    "auc_test = auc(fpr, tpr)\n",
    "\n",
    "plt.figure(figsize=(6, 6))\n",
    "plt.plot(fpr, tpr, label=f\"XGBoost (AUC={auc_test:.4f})\")\n",
    "plt.plot([0, 1], [0, 1], ls=\"--\", color=\"k\")\n",
    "plt.xlabel(\"False positive rate\")\n",
    "plt.ylabel(\"True positive rate\")\n",
    "plt.legend()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prediction_class = prediction_probs.argmax(1)\n",
    "prediction_class\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "acc = accuracy_score(df_test[\"label_Tbqq\"], prediction_class)\n",
    "print(f\"Accuracy: {acc}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "background_eff_fn = interp1d(tpr, fpr)\n",
    "background_eff_at_50 = background_eff_fn(0.5)\n",
    "print(f\"Backround rejection at signal efficiency 50%: {1/background_eff_at_50:.0f}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_metrics(y_true, probs):\n",
    "    fpr, tpr, _ = roc_curve(y_true=y_true, y_score=probs[:, 1])\n",
    "    auc_test = auc(fpr, tpr)\n",
    "    acc_test = accuracy_score(y_true, probs.argmax(1))\n",
    "    background_eff_fn = interp1d(tpr, fpr)\n",
    "    background_eff_at_50 = background_eff_fn(0.5)\n",
    "\n",
    "    print(f\"Accuracy: {acc_test:.4f}\")\n",
    "    print(f\"AUC: {auc_test:.4f}\")\n",
    "    print(f\"Backround rejection at 50% signal efficiency: {1/background_eff_at_50:.0f}\")\n",
    "    return fpr, tpr\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "compute_metrics(y_true, prediction_probs)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualize the results\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "signal_table = uproot.open(signal_file)[\"tree\"].arrays()\n",
    "background_table = uproot.open(background_file)[\"tree\"].arrays()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(2, 5, figsize=(25, 10))\n",
    "for idx in range(10):\n",
    "    prob = prediction_probs[idx, 1]\n",
    "    visualize(signal_table, idx, title=f\"Top quark jet {idx}, prob(Top)={prob:.4f}\", ax=axes[idx % 2][idx // 2])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(2, 5, figsize=(25, 10))\n",
    "for idx in range(10):\n",
    "    prob = prediction_probs[idx + 20000, 1]\n",
    "    visualize(background_table, idx, title=f\"q/g jet {idx}, prob(Top)={prob:.4f}\", ax=axes[idx % 2][idx // 2])\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}