3_dense_network.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 numpy awkward vector uproot lz4 xxhash pandas matplotlib tqdm torch lightning torchmetrics\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": [
    "# Load the dataset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "jet_features = ['jet_pt', 'jet_eta', 'jet_phi', 'jet_energy', 'jet_nparticles',\n",
    "                'jet_sdmass', 'jet_tau1', 'jet_tau2', 'jet_tau3', 'jet_tau4']\n",
    "label_features = ['label_QCD', 'label_Tbqq']\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_jet_signal = uproot.open(signal_file)['tree'].arrays(filter_name=jet_features + label_features, library='pd')\n",
    "df_jet_signal.head(10)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_jet_background = uproot.open(background_file)['tree'].arrays(filter_name=jet_features + label_features, library='pd')\n",
    "df_jet_background.head(10)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dataset splitting: training / validation / testing\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Number of signal jets:', len(df_jet_signal))\n",
    "print('Number of background jets:', len(df_jet_background))\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this study, we are going to use 70% of the data for training, 10% for validation, and 20% for testing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_slice = slice(0, 20000)\n",
    "val_slice = slice(20000, 30000)\n",
    "train_slice = slice(30000, None)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sig_train, sig_val, sig_test = df_jet_signal[train_slice], df_jet_signal[val_slice], df_jet_signal[test_slice]\n",
    "bkg_train, bkg_val, bkg_test = df_jet_background[train_slice], df_jet_background[val_slice], df_jet_background[test_slice]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_train = pd.concat([sig_train, bkg_train])\n",
    "df_val = pd.concat([sig_val, bkg_val])\n",
    "df_test = pd.concat([sig_test, bkg_test])\n",
    "print(f'Training: {len(df_train)}, validation: {len(df_val)}, testing: {len(df_test)}')\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training a fully-connected neural network\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.utils.data as data\n",
    "import torch.nn.functional as F\n",
    "import lightning as L\n",
    "from torchmetrics.functional import accuracy\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MLPModule(L.LightningModule):\n",
    "    def __init__(self, lr=1e-3):\n",
    "        super().__init__()\n",
    "        self.lr = lr\n",
    "        self.nn = nn.Sequential(\n",
    "            nn.BatchNorm1d(10),\n",
    "            nn.Linear(10, 128),\n",
    "            nn.BatchNorm1d(128),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(128, 128),\n",
    "            nn.BatchNorm1d(128),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(128, 128),\n",
    "            nn.BatchNorm1d(128),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(0.2),\n",
    "            nn.Linear(128, 2))\n",
    "\n",
    "    def forward(self, x):\n",
    "        # in lightning, forward defines the prediction/inference actions\n",
    "        return self.nn(x)\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        # training_step defines the train loop. It is independent of forward\n",
    "        x, y = batch\n",
    "        logits = self.nn(x)\n",
    "        loss = F.cross_entropy(logits, y)\n",
    "        self.log(\"train_loss\", loss)\n",
    "        return loss\n",
    "\n",
    "    def evaluate(self, batch, stage=None):\n",
    "        x, y = batch\n",
    "        logits = self(x)\n",
    "        loss = F.nll_loss(logits, y)\n",
    "        probs = F.softmax(logits)\n",
    "        acc = accuracy(probs, y, task='multiclass', num_classes=2)\n",
    "\n",
    "        if stage:\n",
    "            self.log(f\"{stage}_loss\", loss, prog_bar=True)\n",
    "            self.log(f\"{stage}_acc\", acc, prog_bar=True)\n",
    "        return probs\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        self.evaluate(batch, \"val\")\n",
    "\n",
    "    def test_step(self, batch, batch_idx):\n",
    "        self.evaluate(batch, \"test\")\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        optimizer = torch.optim.Adam(self.parameters(), lr=self.lr)\n",
    "        return optimizer\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class JetDataset(data.Dataset):\n",
    "\n",
    "    def __init__(self, df, jet_features=jet_features, label_features=label_features):\n",
    "        self.x = torch.tensor(df[jet_features].values, dtype=torch.float32)\n",
    "        self.y = torch.tensor(df[label_features].values, dtype=torch.int32).argmax(1)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.y)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.x[idx], self.y[idx]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ds = JetDataset(df_train)\n",
    "val_ds = JetDataset(df_val)\n",
    "test_ds = JetDataset(df_test)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 128\n",
    "num_epochs = 20\n",
    "lr = 1e-2\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = MLPModule()\n",
    "trainer = L.Trainer(\n",
    "    max_epochs=num_epochs,\n",
    "    logger=True\n",
    ")\n",
    "trainer.fit(model,\n",
    "            data.DataLoader(train_ds, batch_size=batch_size, shuffle=True),\n",
    "            data.DataLoader(val_ds, batch_size=batch_size, shuffle=True))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.test(model, data.DataLoader(test_ds, batch_size=batch_size))\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Performance evaluation"
   ]
  },
  {
   "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": [
    "model.eval()"
   ]
  },
  {
   "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": [
    "preds = []\n",
    "for x in data.DataLoader(test_ds, batch_size=128):\n",
    "    preds.append(model.evaluate(x).detach().cpu().numpy())\n",
    "\n",
    "prediction_probs = np.concatenate(preds)\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)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualize the results"
   ]
  },
  {
   "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 = model(JetDataset(df_jet_signal)[idx:idx + 1][0]).softmax(1)[0, 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 = model(JetDataset(df_jet_background)[idx:idx + 1][0]).softmax(1)[0, 1]\n",
    "    visualize(background_table, idx, title=f'q/g jet {idx}, prob(Top)={prob:.4f}', ax=axes[idx % 2][idx // 2])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}