analysis_plots.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn import metrics\n",
    "from sklearn.metrics import confusion_matrix\n",
    "from sklearn.metrics import auc, roc_curve\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "\n",
    "import h5py\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "import tqdm\n",
    "import gc\n",
    "import seaborn as sns\n",
    "cols=[\"#DB4437\", \"#4285F4\", \"#F4B400\", \"#0F9D58\", \"purple\", \"goldenrod\", \"peru\", \"coral\",\"turquoise\",'gray','navy','m','darkgreen','fuchsia','steelblue'] "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Style the plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ROOT = {\n",
    "    # \"font.sans-serif\": [\"TeX Gyre Heros\", \"Helvetica\", \"Arial\"],\n",
    "    \"font.family\": \"sans-serif\",\n",
    "    \"mathtext.fontset\": \"custom\",\n",
    "    \"mathtext.rm\": \"TeX Gyre Heros\",\n",
    "    \"mathtext.bf\": \"TeX Gyre Heros:bold\",\n",
    "    \"mathtext.sf\": \"TeX Gyre Heros\",\n",
    "    \"mathtext.it\": \"TeX Gyre Heros:italic\",\n",
    "    \"mathtext.tt\": \"TeX Gyre Heros\",\n",
    "    \"mathtext.cal\": \"TeX Gyre Heros\",\n",
    "    \"mathtext.default\": \"regular\",\n",
    "    \"figure.figsize\": (10.0, 10.0),\n",
    "    \"font.size\": 26,\n",
    "    #\"text.usetex\": True,\n",
    "    \"axes.labelsize\": \"medium\",\n",
    "    \"axes.unicode_minus\": False,\n",
    "    \"xtick.labelsize\": \"small\",\n",
    "    \"ytick.labelsize\": \"small\",\n",
    "    \"legend.fontsize\": \"small\",\n",
    "    \"legend.handlelength\": 1.5,\n",
    "    \"legend.borderpad\": 0.5,\n",
    "    \"xtick.direction\": \"in\",\n",
    "    \"xtick.major.size\": 12,\n",
    "    \"xtick.minor.size\": 6,\n",
    "    \"xtick.major.pad\": 6,\n",
    "    \"xtick.top\": True,\n",
    "    \"xtick.major.top\": True,\n",
    "    \"xtick.major.bottom\": True,\n",
    "    \"xtick.minor.top\": True,\n",
    "    \"xtick.minor.bottom\": True,\n",
    "    \"xtick.minor.visible\": True,\n",
    "    \"ytick.direction\": \"in\",\n",
    "    \"ytick.major.size\": 12,\n",
    "    \"ytick.minor.size\": 6.0,\n",
    "    \"ytick.right\": True,\n",
    "    \"ytick.major.left\": True,\n",
    "    \"ytick.major.right\": True,\n",
    "    \"ytick.minor.left\": True,\n",
    "    \"ytick.minor.right\": True,\n",
    "    \"ytick.minor.visible\": True,\n",
    "    \"grid.alpha\": 0.8,\n",
    "    \"grid.linestyle\": \":\",\n",
    "    \"axes.linewidth\": 2,\n",
    "    \"savefig.transparent\": False,\n",
    "}\n",
    "plt.style.use(ROOT)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with h5py.File('results_ji_latch_norm_100epochs_v2.h5', 'r') as file:\n",
    "    print(file.keys())\n",
    "    vars_test=np.array(file['vars_test'])\n",
    "    top_vars=np.array(file['top_vars'])\n",
    "    labels=np.array(file['labels'])\n",
    "    activations=np.array(file['activations'])\n",
    "    activations_ood=np.array(file['activations_ood'])\n",
    "    top_distance=np.array(file['top_ml_distance'])\n",
    "    qcd_distance=np.array(file['qcd_ml_distance'])\n",
    "    wz_distance = np.array(file['wz_ml_distance'])\n",
    "    logits=np.array(file['logits'])\n",
    "    logits_ood=np.array(file['logits_ood'])\n",
    "    logits_sm=np.array(file['logits_sm'])\n",
    "    logits_sm_ood=np.array(file['logits_sm_ood'])\n",
    "    images=np.array(file['images'])\n",
    "    images_ood=np.array(file['images_ood'])\n",
    "    top_distance_lg=np.array(file['top_distance_lg'])\n",
    "    qcd_distance_lg=np.array(file['qcd_distance_lg'])\n",
    "    wz_distance_lg=np.array(file['wz_distance_lg'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.decomposition import PCA, KernelPCA\n",
    "pca = KernelPCA(n_components=2)\n",
    "pca = pca.fit(activations)\n",
    "activations_pca = pca.transform(activations)\n",
    "activations_ood_pca = pca.transform(activations_ood)\n",
    "activations_pca = activations\n",
    "activations_ood_pca = activations_ood"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "act_cov = np.cov(activations_pca.T)\n",
    "qcd_act_conv = np.cov(activations_pca[labels==0].T)\n",
    "qcd_act_cov_inv = np.linalg.inv(act_cov)\n",
    "qcd_act_mean = np.mean(activations_pca[labels==0],axis=0)\n",
    "\n",
    "wz_act_conv = np.cov(activations_pca[labels==1].T)\n",
    "wz_act_cov_inv = np.linalg.inv(act_cov)\n",
    "wz_act_mean = np.mean(activations_pca[labels==1],axis=0)\n",
    "\n",
    "top_qdistance = np.array([(activations_ood_pca[i]-qcd_act_mean)@qcd_act_cov_inv@(activations_ood_pca[i]-qcd_act_mean).T for i in range(len(activations_ood_pca))])\n",
    "qcd_qdistance = np.array([(activations_pca[labels==0][i]-qcd_act_mean)@qcd_act_cov_inv@(activations_pca[labels==0][i]-qcd_act_mean).T for i in range(len(activations_pca[labels==0]))])\n",
    "wz_qdistance = np.array([(activations_pca[labels==1][i]-wz_act_mean)@wz_act_cov_inv@(activations_pca[labels==1][i]-wz_act_mean).T for i in range(len(activations_pca[labels==1]))])\n",
    "\n",
    "top_wzdistance = np.array([(activations_ood_pca[i]-wz_act_mean)@wz_act_cov_inv@(activations_ood_pca[i]-wz_act_mean).T for i in range(len(activations_ood_pca))])\n",
    "qcd_wzdistance = np.array([(activations_pca[labels==0][i]-wz_act_mean)@wz_act_cov_inv@(activations_pca[labels==0][i]-wz_act_mean).T for i in range(len(activations_pca[labels==0]))])\n",
    "wz_wzdistance = np.array([(activations_pca[labels==1][i]-wz_act_mean)@wz_act_cov_inv@(activations_pca[labels==1][i]-wz_act_mean).T for i in range(len(activations_pca[labels==1]))])\n",
    "\n",
    "_ = plt.hist(top_qdistance, bins=100, alpha=0.5, label='top')\n",
    "_ = plt.hist(qcd_qdistance, bins=100, alpha=0.5, label='qcd')\n",
    "_ = plt.hist(wz_qdistance, bins=100, alpha=0.5, label='wz')\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "_ = plt.hist(top_wzdistance, bins=100, alpha=0.5, label='top')\n",
    "_ = plt.hist(qcd_wzdistance, bins=100, alpha=0.5, label='qcd')\n",
    "_ = plt.hist(wz_wzdistance, bins=100, alpha=0.5, label='wz')\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "top_distance = np.concatenate([(top_qdistance[:,None]), (top_wzdistance[:,None])], axis=1)\n",
    "qcd_distance = np.concatenate([(qcd_qdistance[:,None]), (qcd_wzdistance[:,None])], axis=1)\n",
    "wz_distance = np.concatenate([(wz_qdistance[:,None]), (wz_wzdistance[:,None])], axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plots for Latent representation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(20, 20))\n",
    "spec = fig.add_gridspec(ncols=2, nrows=2, )\n",
    "\n",
    "ax1 = fig.add_subplot(spec[0, 0])\n",
    "temp = activations_pca[labels == 0][:,:2]\n",
    "_ = ax1.hist2d(temp[:, 0], temp[:, 1], bins=100, cmap='Blues')\n",
    "ax1.set_title('QCD')\n",
    "\n",
    "ax2 = fig.add_subplot(spec[0, 1],sharex=ax1,sharey=ax1)\n",
    "temp = activations_pca[labels == 1][:,:2]\n",
    "_ = ax2.hist2d(temp[:, 0], temp[:, 1], bins=100, cmap='Reds')\n",
    "ax2.set_title('W/Z')\n",
    "\n",
    "ax3 = fig.add_subplot(spec[1, 0],sharex=ax1,sharey=ax1)\n",
    "temp = activations_ood_pca[:,:2]\n",
    "_ = ax3.hist2d(temp[:, 0], temp[:, 1], bins=100, cmap='Greys')\n",
    "ax3.set_title('Top')\n",
    "# plt.show()\n",
    "\n",
    "ax4 = fig.add_subplot(spec[1, 1],sharex=ax1,sharey=ax1)\n",
    "mask = (top_vars[:, 48] > 150) * (top_vars[:, 48] < 200)\n",
    "temp = activations_ood_pca[mask][:,:2]\n",
    "_ = ax4.hist2d(temp[:, 0], temp[:, 1], bins=100, cmap='Greys')\n",
    "ax4.set_title('Top')\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import seaborn as sns\n",
    "import matplotlib.lines as  mlines\n",
    "\n",
    "fig = plt.figure(figsize=(10, 10))\n",
    "sns.kdeplot(x=activations_pca[labels == 0][:,0], y=activations_pca[labels == 0][:,1], fill=False, label='QCD', levels=5, color=cols[0],alpha=0.5)\n",
    "sns.kdeplot(x=activations_pca[labels == 1][:,0], y=activations_pca[labels == 1][:,1], fill=False, label='W/Z', levels=5, color=cols[1],alpha=0.5)\n",
    "# sns.kdeplot(x=activations_ood_pca[:,0], y=activations_ood_pca[:,1], fill=False, label='Top', levels=5, color='grey',alpha=0.5)\n",
    "sns.kdeplot(x=activations_ood_pca[mask][:,0], y=activations_ood_pca[mask][:,1], fill=False, label='Top (150 < M < 200)', levels=5, color=cols[2],alpha=0.5)\n",
    "handles = [mlines.Line2D([], [],color=cols[0], label=\"QCD\"),\n",
    "           mlines.Line2D([], [],color=cols[1], label=\"W/Z\"),\n",
    "           mlines.Line2D([], [],color=cols[2], label=\"Top [OOD]\")]\n",
    "plt.legend(handles=handles)\n",
    "plt.xlabel('Represntation dim. #1')\n",
    "plt.ylabel('Represntation dim. #2')\n",
    "# plt.show()\n",
    "# plt.savefig('PCA_representation.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = np.linspace(-0, 200, 100)\n",
    "fig = plt.figure()\n",
    "_ = plt.hist(np.max(top_distance, axis=1), bins=bins, alpha=0.5,\n",
    "             label='top', density=True, histtype='step', lw=2)\n",
    "_ = plt.hist(np.max(qcd_distance, axis=1), bins=bins, alpha=0.5,\n",
    "             label='qcd', density=True, histtype='step', lw=2)\n",
    "_ = plt.hist(np.max(wz_distance, axis=1), bins=bins, alpha=0.5,\n",
    "             label='wz', density=True, histtype='step', lw=2)\n",
    "plt.yscale('log')\n",
    "plt.ylim(0.0001, 1)\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "fig = plt.figure()\n",
    "bins = np.linspace(-0, 200, 100)\n",
    "_ = plt.hist(np.min(top_distance[mask], axis=1), bins=bins, alpha=0.5,\n",
    "             label='Top [OOD Sample]', density=True, histtype='step', lw=2)\n",
    "_ = plt.hist(np.min(qcd_distance, axis=1), bins=bins, alpha=0.5,\n",
    "             label='QCD', density=True, histtype='step', lw=2)\n",
    "_ = plt.hist(np.min(wz_distance, axis=1), bins=bins, alpha=0.5,\n",
    "             label='WZ', density=True, histtype='step', lw=2)\n",
    "plt.yscale('log')\n",
    "plt.ylim(0.0001, 1)\n",
    "plt.legend()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred = np.concatenate(\n",
    "    [np.max(top_distance[mask], axis=1), np.max(qcd_distance, axis=1)])\n",
    "y = np.concatenate([np.ones(len(top_distance[mask])),\n",
    "                   np.zeros(len(qcd_distance))])\n",
    "fpr, tpr, thresholds = roc_curve(y, pred)\n",
    "auc_roc = auc(fpr, tpr)\n",
    "plt.plot(fpr, tpr, label='ROC curve (area = %0.2f)' % auc_roc)\n",
    "plt.legend(loc='lower right')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "top_distance = np.concatenate([(top_qdistance[:,None]), (top_wzdistance[:,None])], axis=1)\n",
    "qcd_distance = np.concatenate([(qcd_qdistance[:,None]), (qcd_wzdistance[:,None])], axis=1)\n",
    "wz_distance = np.concatenate([(wz_qdistance[:,None]), (wz_wzdistance[:,None])], axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with h5py.File('results_vae.h5','r') as file:\n",
    "    vars_test_vae=np.array(file['vars_test'])\n",
    "    top_vars_vae=np.array(file['top_vars'])\n",
    "    qcd_loss=np.array(file['qcd_loss'])\n",
    "    top_loss=np.array(file['top_loss'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qcd_distance.shape\n",
    "# vars_test[labels == 0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mask = (top_vars[:, 48] > 150) * (top_vars[:, 48] < 200)\n",
    "mask_qcd = (vars_test[labels == 0, 48] > 150) * \\\n",
    "    (vars_test[labels == 0, 48] < 200)\n",
    "mask_wz = (vars_test[labels == 1, 48] > 150) * \\\n",
    "    (vars_test[labels == 1, 48] < 200)\n",
    "bins = np.linspace(-2, 150, 100)\n",
    "_ = plt.hist(np.min(top_distance[mask], axis=1), bins=bins, alpha=0.5,\n",
    "             label='Top [OOD Sample]', density=True, histtype='step', lw=2,color=cols[2])\n",
    "_ = plt.hist(np.min(qcd_distance, axis=1), bins=bins, alpha=0.5,\n",
    "             label='QCD', density=True, histtype='step', lw=2,color=cols[0])\n",
    "_ = plt.hist(np.min(wz_distance, axis=1), bins=bins, alpha=0.5,\n",
    "             label='WZ', density=True, histtype='step', lw=2,color=cols[1])\n",
    "plt.yscale('log')\n",
    "plt.xlabel('Mahalanobis Distance')\n",
    "plt.legend()\n",
    "plt.savefig('max_mahalanobis_distance.pdf', bbox_inches='tight')\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred = np.concatenate([np.max(top_distance[mask],axis=1), np.max(qcd_distance,axis=1)])\n",
    "y = np.concatenate([np.ones(len(top_distance[mask])), np.zeros(len(qcd_distance))])\n",
    "fpr_, tpr_, thresholds = roc_curve(y, pred)\n",
    "auc_roc = auc(fpr_, tpr_)\n",
    "plt.plot(fpr_, tpr_, label='Mahalanobis Distance (AUC = %0.2f)' % auc_roc,color=cols[1])\n",
    "\n",
    "\n",
    "true_label_top = np.concatenate(( np.ones(top_loss.shape[0]), np.zeros(qcd_loss.shape[0]) ))\n",
    "predicted_label_top = np.concatenate(( top_loss, qcd_loss ))\n",
    "fpr_loss, tpr_loss, threshold_loss = roc_curve(true_label_top, predicted_label_top)\n",
    "auc_loss = auc(fpr_loss, tpr_loss)\n",
    "plt.plot(fpr_loss, tpr_loss, label='%s (AUC = %0.2f)' %\n",
    "         ('Variational Autoencoder', auc_loss), color=cols[0])\n",
    "\n",
    "\n",
    "qcd_max_lg = np.max(logits[labels==0],axis=1)\n",
    "wz_max_lg = np.max(logits[labels==1],axis=1)\n",
    "top_max_lg = np.max(logits_ood,axis=1)\n",
    "\n",
    "pred = np.concatenate([top_max_lg[mask], qcd_max_lg])\n",
    "y = np.concatenate([np.ones(len(top_max_lg[mask])), np.zeros(len(qcd_max_lg))])\n",
    "fpr, tpr, thresholds = roc_curve(y, pred)\n",
    "auc_roc = auc(fpr, tpr)\n",
    "plt.plot(fpr, tpr, label='Max Logits curve (area = %0.2f)' % auc_roc,color=cols[2])\n",
    "plt.plot([1e-10,1],[1e-10,1], '--', color='0.75')\n",
    "plt.legend(loc='lower right')\n",
    "plt.xlabel('False Positive Rate')\n",
    "plt.ylabel('True Positive Rate')\n",
    "plt.xlim(1e-5,1)\n",
    "plt.ylim(1e-5,1)\n",
    "# plt.yscale('log')\n",
    "# plt.xscale('log')\n",
    "plt.savefig('roc_curve_v2.pdf', bbox_inches='tight')\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(tpr_,tpr_/np.sqrt(fpr_),color=cols[1],label='Max Mahalanobis Distance')\n",
    "plt.plot(tpr,tpr/np.sqrt(fpr),color=cols[2],label='Max Logits')\n",
    "plt.plot(tpr_loss,tpr_loss/np.sqrt(fpr_loss),color=cols[0],label='Variational Autoencoder')\n",
    "plt.legend()\n",
    "plt.xlabel('True Positive Rate')\n",
    "plt.ylabel(r'$ S_{eff} / \\sqrt(B_{eff})$')\n",
    "plt.savefig('sb_curve.pdf', bbox_inches='tight')\n",
    "print(np.max(tpr_/np.sqrt(fpr_)),np.max(tpr/np.sqrt(fpr)),np.max(tpr_loss/np.sqrt(fpr_loss)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(qcd_max_lg, bins=100, alpha=0.5, label='QCD', density=True, histtype='step', lw=2, color=cols[0])\n",
    "plt.hist(wz_max_lg, bins=100, alpha=0.5, label='W/Z', density=True, histtype='step', lw=2, color=cols[1])\n",
    "plt.hist(top_max_lg[mask], bins=100, alpha=0.5, label='Top [OOD Sample]', density=True, histtype='step', lw=2,  color=cols[2])\n",
    "plt.xlabel('Max Logits')\n",
    "plt.ylabel('Density')\n",
    "plt.legend()\n",
    "plt.savefig('max_logits.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qcd_mass = vars_test[labels==0][:,48]\n",
    "qcd_distance_max = np.max(qcd_distance, axis=1)\n",
    "th = np.quantile(qcd_distance_max, 0.5)\n",
    "bins=np.linspace(0, 300, 60)\n",
    "hist_qcd ,_,_= plt.hist(qcd_mass, bins=bins, alpha=0.5, label='QCD Mass Dist.', density=True, histtype='step', lw=4)\n",
    "hist_md ,_,_= plt.hist(qcd_mass[qcd_distance_max>th], bins=bins, alpha=0.5, label='After NuRD selection', density=True, histtype='step', lw=2)\n",
    "\n",
    "qcd_mass_vae = vars_test_vae[:,48]\n",
    "th_vae = np.quantile(qcd_loss, 0.5)\n",
    "hist_vae, _, _ = plt.hist(qcd_mass_vae[qcd_loss > th_vae], bins=bins, alpha=0.5,\n",
    "                         label='After VAE selection', density=True, histtype='step', lw=2)\n",
    "plt.xlabel('Mass [GeV]')\n",
    "\n",
    "\n",
    "th_qcd_lg = np.quantile(qcd_max_lg, 0.5)\n",
    "hist_ml ,_,_= plt.hist(qcd_mass[qcd_max_lg>th_qcd_lg], bins=bins, alpha=0.5, label='After Logits selection', density=True, histtype='step', lw=2)\n",
    "plt.legend()\n",
    "plt.ylabel('Normalized Counts')\n",
    "plt.yscale('log')\n",
    "# plt.savefig('mass_dist_log_v2.pdf', bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import wasserstein_distance\n",
    "print('Wasserstein Distance between QCD and Top [OOD Sample] (NuRD):', wasserstein_distance(hist_qcd*10000, hist_md*10000))\n",
    "print('Wasserstein Distance between QCD and Top [OOD Sample] (VAE):', wasserstein_distance(hist_qcd*10000, hist_vae*10000))\n",
    "print('Wasserstein Distance between QCD and Top [OOD Sample] (Logits):', wasserstein_distance(hist_qcd*10000, hist_ml*10000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.spatial.distance import jensenshannon\n",
    "print('Jensen-Shannon Divergence between QCD and Top [OOD Sample] (NuRD-MD):', jensenshannon(hist_qcd, hist_md)**2)\n",
    "print('Jensen-Shannon Divergence between QCD and Top [OOD Sample] (VAE):', jensenshannon(hist_qcd, hist_vae)**2)\n",
    "print('Jensen-Shannon Divergence between QCD and Top [OOD Sample] (NuRD-ML):', jensenshannon(hist_qcd, hist_ml)**2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# bins = np.linspace(0, 1000, 50)\n",
    "_ = plt.hist(np.min(top_distance, axis=1), bins=bins, alpha=0.5,\n",
    "             label='Top [OOD Sample]', density=True, histtype='step', lw=2)\n",
    "_ = plt.hist(np.min(qcd_distance, axis=1), bins=bins, alpha=0.5,\n",
    "             label='QCD', density=True, histtype='step', lw=2)\n",
    "_ = plt.hist(np.min(wz_distance, axis=1), bins=bins, alpha=0.5,\n",
    "             label='WZ', density=True, histtype='step', lw=2)\n",
    "plt.yscale('log')\n",
    "plt.ylabel('Fraction')\n",
    "plt.xlabel('OOD Score')\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "# bins = np.linspace(-4, 10, 100)\n",
    "# _ = plt.hist(np.min(top_distance[mask], axis=1), bins=bins, alpha=0.5,\n",
    "#              label='Top [OOD Sample]', density=True, histtype='step', lw=2)\n",
    "# _ = plt.hist(np.min(qcd_distance, axis=1), bins=bins, alpha=0.5,\n",
    "#              label='QCD', density=True, histtype='step', lw=2)\n",
    "# _ = plt.hist(np.min(wz_distance, axis=1), bins=bins, alpha=0.5,\n",
    "#              label='WZ', density=True, histtype='step', lw=2)\n",
    "# plt.yscale('log')\n",
    "# plt.legend()\n",
    "# plt.show()\n",
    "\n",
    "pred = np.concatenate([np.max(top_distance, axis=1),\n",
    "                      np.max(qcd_distance, axis=1)])\n",
    "y = np.concatenate([np.ones(len(top_distance)), np.zeros(len(qcd_distance))])\n",
    "fpr, tpr, thresholds = roc_curve(y, pred)\n",
    "auc_roc = auc(fpr, tpr)\n",
    "plt.plot(fpr, tpr, label='ROC curve (area = %0.2f)' % auc_roc)\n",
    "plt.legend(loc='lower right')\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qcd_mass = vars_test[labels==0]\n",
    "qcd_min_dist = np.max(qcd_distance, axis=1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qcd_mass = vars_test[labels==0][:,48]\n",
    "qcd_distance_max = np.max(qcd_distance, axis=1)\n",
    "th = np.quantile(qcd_distance_max, 0.75)\n",
    "plt.hist(qcd_mass, bins=100, alpha=0.5, label='QCD Evetyhing', density=True, histtype='step', lw=2)\n",
    "plt.hist(qcd_mass[qcd_distance_max>th], bins=100, alpha=0.5, label='QCD after OOD selection', density=True, histtype='step', lw=2)\n",
    "plt.xlabel('Jet Mass[GeV]')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  }
 ],
 "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.16"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "ccad25ce197360c0f51c7d951f15faa3a823854e2a2bcc1c7d9e660ec9079ef2"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}