apply_cuts.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "de6716b1-8361-4bad-9d35-019f3147c376",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-08-08T23:12:02.413096Z",
     "iopub.status.busy": "2022-08-08T23:12:02.412330Z",
     "iopub.status.idle": "2022-08-08T23:12:02.417374Z",
     "shell.execute_reply": "2022-08-08T23:12:02.416634Z",
     "shell.execute_reply.started": "2022-08-08T23:12:02.413055Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import uproot\n",
    "from glob import glob\n",
    "import json \n",
    "import os\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "31e0d8bc-71e7-4aaa-85c2-164dec2fcf9e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-08-08T23:12:02.521395Z",
     "iopub.status.busy": "2022-08-08T23:12:02.521027Z",
     "iopub.status.idle": "2022-08-08T23:12:02.531810Z",
     "shell.execute_reply": "2022-08-08T23:12:02.531101Z",
     "shell.execute_reply.started": "2022-08-08T23:12:02.521366Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "pi0_dir = '/clusterfs/ml4hep/mpettee/ml4pions/data/root_files/user.angerami.mc16_13TeV.900246.PG_singlepi0_logE0p2to2000.e8312_e7400_s3170_r12383.v01-45-gaa27bcb_OutputStream/'\n",
    "pion_dir = '/clusterfs/ml4hep/mpettee/ml4pions/data/root_files/user.angerami.mc16_13TeV.900247.PG_singlepion_logE0p2to2000.e8312_e7400_s3170_r12383.v01-45-gaa27bcb_OutputStream/'\n",
    "pi0_files = sorted(glob(pi0_dir+\"*.root\"))\n",
    "pion_files = sorted(glob(pion_dir+\"*.root\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b71a1293-2243-4d2d-b0fc-efcaab9e2276",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('pion_train.json', 'w') as f:\n",
    "    json.dump(pion_files[:350], f)\n",
    "with open('pion_val.json', 'w') as f:\n",
    "    json.dump(pion_files[350:400], f)\n",
    "with open('pion_test.json', 'w') as f:\n",
    "    json.dump(pion_files[400:], f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "cbca2fc3-f9ee-4dac-a3c7-d35b09d32d56",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-08-08T23:12:07.092493Z",
     "iopub.status.busy": "2022-08-08T23:12:07.092193Z",
     "iopub.status.idle": "2022-08-08T23:12:07.104023Z",
     "shell.execute_reply": "2022-08-08T23:12:07.103385Z",
     "shell.execute_reply.started": "2022-08-08T23:12:07.092462Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import math \n",
    "def delta_phi(phi1, phi2):\n",
    "    return (phi1 - phi2 + math.pi) % (2*math.pi) - math.pi\n",
    "\n",
    "def apply_cuts(arrays):\n",
    "    df = pd.DataFrame(arrays)\n",
    "\n",
    "    ### Single-track, multi-cluster \n",
    "    df = df[(df.nTrack == 1)]\n",
    "    dR_pass = []\n",
    "    deltaR_list = []\n",
    "    for row in df.index:\n",
    "        try:\n",
    "            deltaR = np.sqrt((df['cluster_Eta'][row].astype('float') - df['trackEta'][row].astype('float'))**2 + \n",
    "                               delta_phi(df['cluster_Phi'][row].astype('float'), df['trackPhi'][row].astype('float'))**2)\n",
    "        except:\n",
    "            deltaR = np.array(999)\n",
    "        deltaR_list.append(deltaR)\n",
    "        dR_pass.append(deltaR < 1.2)\n",
    "    df[\"deltaR\"] = deltaR_list\n",
    "    df[\"dR_pass\"] = dR_pass\n",
    "    df[\"event_number\"] = df.index\n",
    "    df.reset_index(inplace=True)\n",
    "    indices_pass = [] \n",
    "    for row in df.index: # kill all events with no clusters passing the delta R cut; deal with individual clusters at the training stage\n",
    "        if len(df['trackPt'][row]) > 1:\n",
    "            continue\n",
    "        elif df.dR_pass[row].sum() > 0:\n",
    "            indices_pass.append(row)\n",
    "    df = df.iloc[indices_pass]\n",
    "    df = df[(df.trackPt < 10**5)] # Track pT cut \n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "7b2b209c-44ba-442b-a2b4-8df3111cd403",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-08-08T23:12:52.255618Z",
     "iopub.status.busy": "2022-08-08T23:12:52.255326Z",
     "iopub.status.idle": "2022-08-08T23:12:52.259776Z",
     "shell.execute_reply": "2022-08-08T23:12:52.259331Z",
     "shell.execute_reply.started": "2022-08-08T23:12:52.255602Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# use a pared-down list of variables \n",
    "variables = ['cluster_cell_E', 'cluster_cell_ID',\n",
    "             'trackPt','trackD0','trackZ0',\n",
    "             'trackEta_EMB2','trackPhi_EMB2',\n",
    "             'trackEta_EME2','trackPhi_EME2',\n",
    "             'trackEta','trackPhi',\n",
    "             'nCluster','nTrack','truthPartE', 'truthPartPt',\n",
    "             'cluster_ENG_CALIB_TOT','cluster_E','cluster_Eta','cluster_Phi',\n",
    "             'cluster_EM_PROBABILITY','cluster_E_LCCalib','cluster_HAD_WEIGHT',             \n",
    "            ]\n",
    "\n",
    "### SANMAY'S VARIABLES\n",
    "cluster_var = ['cluster_EM_PROBABILITY', 'cluster_HAD_WEIGHT', 'cluster_OOC_WEIGHT',\n",
    "               'cluster_DM_WEIGHT', 'cluster_CENTER_MAG', 'cluster_FIRST_ENG_DENS', \n",
    "               'cluster_CENTER_LAMBDA', 'cluster_ISOLATION'\n",
    "              ]\n",
    "\n",
    "track_var = ['trackPt',\n",
    "             'trackP',\n",
    "             'trackMass',\n",
    "             'trackEta',\n",
    "             'trackPhi',\n",
    "             'trackNumberOfPixelHits',\n",
    "             'trackNumberOfSCTHits',\n",
    "             'trackNumberOfPixelDeadSensors',\n",
    "             'trackNumberOfSCTDeadSensors',\n",
    "             'trackNumberOfInnermostPixelLayerHits',\n",
    "             'trackNumberOfNextToInnermostPixelLayerHits',\n",
    "             'trackExpectInnermostPixelLayerHit',\n",
    "             'trackExpectNextToInnermostPixelLayerHit',\n",
    "             'trackNumberOfTRTHits',\n",
    "             'trackNumberOfTRTOutliers',\n",
    "             'trackChiSquared',\n",
    "             'trackNumberDOF',\n",
    "             'trackD0',\n",
    "             'trackZ0'\n",
    "            ]\n",
    "\n",
    "sanmay_vars = variables + cluster_var + track_var"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cab8e04e-7687-4a9f-bba7-a2ca19bef40c",
   "metadata": {},
   "outputs": [],
   "source": [
    "for file in tqdm(pi0_files):\n",
    "    prefix = file.split(\"/\")[:-2]\n",
    "    number = file.split(\"000\")[-1][:-5]\n",
    "    folder = os.path.join(\"/\".join(prefix), \n",
    "                          \"onetrack_multicluster\", \n",
    "                          \"pi0_files\")\n",
    "    os.makedirs(folder, exist_ok=True)\n",
    "    npy_filename = os.path.join(folder, str(number)+\".npy\")\n",
    "    a = uproot.open(file)[\"EventTree\"].arrays(variables, library = \"np\")\n",
    "    df = apply_cuts(a)\n",
    "    a_cuts = df.to_dict('list')\n",
    "    np.save(npy_filename, a_cuts) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "c5089240-89f9-4709-9cc3-981e4b28a450",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-08-08T23:33:21.499286Z",
     "iopub.status.busy": "2022-08-08T23:33:21.498753Z",
     "iopub.status.idle": "2022-08-08T23:42:56.714435Z",
     "shell.execute_reply": "2022-08-08T23:42:56.713892Z",
     "shell.execute_reply.started": "2022-08-08T23:33:21.499236Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 100/100 [09:35<00:00,  5.75s/it]\n"
     ]
    }
   ],
   "source": [
    "for file in tqdm(pion_files[400:]):\n",
    "    # prefix = file.split(\"/\")[:-2]\n",
    "    prefix = \"/clusterfs/ml4hep/mpettee/ml4pions/data/\"\n",
    "    number = file.split(\"000\")[-1][:-5]\n",
    "    folder = os.path.join(\n",
    "        prefix, # \"/\".join(prefix), \n",
    "                          \"onetrack_multicluster_sanmay\", \n",
    "                          \"pion_files\")\n",
    "    os.makedirs(folder, exist_ok=True)\n",
    "    npy_filename = os.path.join(folder, str(number)+\".npy\")\n",
    "    a = uproot.open(file)[\"EventTree\"].arrays(sanmay_vars, library = \"np\")\n",
    "    df = apply_cuts(a)\n",
    "    a_cuts = df.to_dict('list')\n",
    "    np.save(npy_filename, a_cuts) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "928f8ef6-a602-497e-a133-9cec85df81da",
   "metadata": {},
   "source": [
    "### Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "019b07bd-2d15-4960-bcc8-82539af5b836",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-08-08T22:09:00.426569Z",
     "iopub.status.busy": "2022-08-08T22:09:00.425992Z",
     "iopub.status.idle": "2022-08-08T22:09:06.603241Z",
     "shell.execute_reply": "2022-08-08T22:09:06.602668Z",
     "shell.execute_reply.started": "2022-08-08T22:09:00.426546Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "a = uproot.open(pion_files[400])[\"EventTree\"].arrays(\n",
    "    # variables, \n",
    "    library = \"np\")\n",
    "df = apply_cuts(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "d3a21460-ead6-420a-a32a-80ad3b3e7a8c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-08-01T21:32:57.212256Z",
     "iopub.status.busy": "2022-08-01T21:32:57.211957Z",
     "iopub.status.idle": "2022-08-01T21:32:57.216918Z",
     "shell.execute_reply": "2022-08-01T21:32:57.216576Z",
     "shell.execute_reply.started": "2022-08-01T21:32:57.212221Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0         [5.482727]\n",
       "2        [26.158798]\n",
       "3        [14.467227]\n",
       "4        [4.7492914]\n",
       "5        [4.5480275]\n",
       "            ...     \n",
       "12143     [17.39683]\n",
       "12145    [41.555126]\n",
       "12146    [505.93658]\n",
       "12148    [114.87165]\n",
       "12149    [19.100674]\n",
       "Name: truthPartE, Length: 10273, dtype: object"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.truthPartE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "f442204a-e85d-40f7-a231-cf530a6a5a61",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-08-01T21:28:39.362279Z",
     "iopub.status.busy": "2022-08-01T21:28:39.361889Z",
     "iopub.status.idle": "2022-08-01T21:28:39.446723Z",
     "shell.execute_reply": "2022-08-01T21:28:39.445901Z",
     "shell.execute_reply.started": "2022-08-01T21:28:39.362245Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(df[np.abs(df.truthPartE.explode() - 1.638227) < 0.001].truthPartE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "993348d0-b6ce-463a-965c-9bc8464e4acc",
   "metadata": {},
   "source": [
    "##### Look for the indices to compare with Sanmay:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "9a1ee82c-8478-4df7-9eee-387f43f4a79a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-08-01T21:42:12.307713Z",
     "iopub.status.busy": "2022-08-01T21:42:12.307327Z",
     "iopub.status.idle": "2022-08-01T21:42:12.311532Z",
     "shell.execute_reply": "2022-08-01T21:42:12.310938Z",
     "shell.execute_reply.started": "2022-08-01T21:42:12.307679Z"
    }
   },
   "outputs": [],
   "source": [
    "energies = [1.638227, 9.119415, 29.691723, 100.291687, 499.597717, 995.071655]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "0a44316c-104b-448f-b05b-d50c0f55702d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-08-01T21:43:24.257232Z",
     "iopub.status.busy": "2022-08-01T21:43:24.256966Z",
     "iopub.status.idle": "2022-08-01T21:49:08.794308Z",
     "shell.execute_reply": "2022-08-01T21:49:08.793755Z",
     "shell.execute_reply.started": "2022-08-01T21:43:24.257199Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  7%|▋         | 7/100 [00:24<05:14,  3.38s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Match found! File = /clusterfs/ml4hep/mpettee/ml4pions/data/root_files/user.angerami.mc16_13TeV.900247.PG_singlepion_logE0p2to2000.e8312_e7400_s3170_r12383.v01-45-gaa27bcb_OutputStream/user.angerami.24559744.OutputStream._000410.root\n",
      "4109    [1.638227]\n",
      "Name: truthPartE, dtype: object\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 28%|██▊       | 28/100 [01:36<04:11,  3.50s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Match found! File = /clusterfs/ml4hep/mpettee/ml4pions/data/root_files/user.angerami.mc16_13TeV.900247.PG_singlepion_logE0p2to2000.e8312_e7400_s3170_r12383.v01-45-gaa27bcb_OutputStream/user.angerami.24559744.OutputStream._000430.root\n",
      "8254    [499.59772]\n",
      "Name: truthPartE, dtype: object\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 30%|███       | 30/100 [01:43<03:57,  3.39s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Match found! File = /clusterfs/ml4hep/mpettee/ml4pions/data/root_files/user.angerami.mc16_13TeV.900247.PG_singlepion_logE0p2to2000.e8312_e7400_s3170_r12383.v01-45-gaa27bcb_OutputStream/user.angerami.24559744.OutputStream._000432.root\n",
      "5861    [9.11941]\n",
      "Name: truthPartE, dtype: object\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 38%|███▊      | 38/100 [02:11<03:40,  3.55s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Match found! File = /clusterfs/ml4hep/mpettee/ml4pions/data/root_files/user.angerami.mc16_13TeV.900247.PG_singlepion_logE0p2to2000.e8312_e7400_s3170_r12383.v01-45-gaa27bcb_OutputStream/user.angerami.24559744.OutputStream._000440.root\n",
      "9082    [9.119415]\n",
      "Name: truthPartE, dtype: object\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 59%|█████▉    | 59/100 [03:23<02:20,  3.44s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Match found! File = /clusterfs/ml4hep/mpettee/ml4pions/data/root_files/user.angerami.mc16_13TeV.900247.PG_singlepion_logE0p2to2000.e8312_e7400_s3170_r12383.v01-45-gaa27bcb_OutputStream/user.angerami.24559744.OutputStream._000461.root\n",
      "3985    [29.691723]\n",
      "Name: truthPartE, dtype: object\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 97%|█████████▋| 97/100 [05:34<00:10,  3.41s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Match found! File = /clusterfs/ml4hep/mpettee/ml4pions/data/root_files/user.angerami.mc16_13TeV.900247.PG_singlepion_logE0p2to2000.e8312_e7400_s3170_r12383.v01-45-gaa27bcb_OutputStream/user.angerami.24559744.OutputStream._000499.root\n",
      "3592    [100.29169]\n",
      "Name: truthPartE, dtype: object\n",
      "Match found! File = /clusterfs/ml4hep/mpettee/ml4pions/data/root_files/user.angerami.mc16_13TeV.900247.PG_singlepion_logE0p2to2000.e8312_e7400_s3170_r12383.v01-45-gaa27bcb_OutputStream/user.angerami.24559744.OutputStream._000499.root\n",
      "8425    [995.07166]\n",
      "Name: truthPartE, dtype: object\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 100/100 [05:44<00:00,  3.45s/it]\n"
     ]
    }
   ],
   "source": [
    "for file in tqdm(pion_files[400:]):\n",
    "    a = uproot.open(file)[\"EventTree\"].arrays(variables, library = \"np\")\n",
    "    df = apply_cuts(a)\n",
    "    for energy in energies:\n",
    "        if len(df[np.abs(df.truthPartE.explode() - energy) < 0.00001].truthPartE) != 0:\n",
    "            print(\"Match found! File = {}\".format(file))\n",
    "            print(df[np.abs(df.truthPartE.explode() - energy) < 0.00001].truthPartE)\n",
    "# print(df.truthPartE.min())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "11f912a6-7940-4de9-a589-0863ad919c06",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-08-01T21:54:36.802163Z",
     "iopub.status.busy": "2022-08-01T21:54:36.801950Z",
     "iopub.status.idle": "2022-08-01T21:54:36.805930Z",
     "shell.execute_reply": "2022-08-01T21:54:36.805364Z",
     "shell.execute_reply.started": "2022-08-01T21:54:36.802143Z"
    }
   },
   "outputs": [],
   "source": [
    "selected_pion_files = [\"/clusterfs/ml4hep/mpettee/ml4pions/data/root_files/user.angerami.mc16_13TeV.900247.PG_singlepion_logE0p2to2000.e8312_e7400_s3170_r12383.v01-45-gaa27bcb_OutputStream/user.angerami.24559744.OutputStream._000410.root\",\n",
    "\"/clusterfs/ml4hep/mpettee/ml4pions/data/root_files/user.angerami.mc16_13TeV.900247.PG_singlepion_logE0p2to2000.e8312_e7400_s3170_r12383.v01-45-gaa27bcb_OutputStream/user.angerami.24559744.OutputStream._000440.root\",\n",
    "\"/clusterfs/ml4hep/mpettee/ml4pions/data/root_files/user.angerami.mc16_13TeV.900247.PG_singlepion_logE0p2to2000.e8312_e7400_s3170_r12383.v01-45-gaa27bcb_OutputStream/user.angerami.24559744.OutputStream._000461.root\",\n",
    "\"/clusterfs/ml4hep/mpettee/ml4pions/data/root_files/user.angerami.mc16_13TeV.900247.PG_singlepion_logE0p2to2000.e8312_e7400_s3170_r12383.v01-45-gaa27bcb_OutputStream/user.angerami.24559744.OutputStream._000430.root\",\n",
    "\"/clusterfs/ml4hep/mpettee/ml4pions/data/root_files/user.angerami.mc16_13TeV.900247.PG_singlepion_logE0p2to2000.e8312_e7400_s3170_r12383.v01-45-gaa27bcb_OutputStream/user.angerami.24559744.OutputStream._000499.root\",]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "2de0ea57-50bc-44f3-a154-3198ea5fd2a4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-08-01T22:24:33.745713Z",
     "iopub.status.busy": "2022-08-01T22:24:33.745387Z",
     "iopub.status.idle": "2022-08-01T22:25:09.154410Z",
     "shell.execute_reply": "2022-08-01T22:25:09.153690Z",
     "shell.execute_reply.started": "2022-08-01T22:24:33.745679Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/5 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Match found! File = /clusterfs/ml4hep/mpettee/ml4pions/data/root_files/user.angerami.mc16_13TeV.900247.PG_singlepion_logE0p2to2000.e8312_e7400_s3170_r12383.v01-45-gaa27bcb_OutputStream/user.angerami.24559744.OutputStream._000410.root\n",
      "4109    [1.638227]\n",
      "Name: truthPartE, dtype: object\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 20%|██        | 1/5 [00:07<00:28,  7.25s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Match found! File = /clusterfs/ml4hep/mpettee/ml4pions/data/root_files/user.angerami.mc16_13TeV.900247.PG_singlepion_logE0p2to2000.e8312_e7400_s3170_r12383.v01-45-gaa27bcb_OutputStream/user.angerami.24559744.OutputStream._000440.root\n",
      "9082    [9.119415]\n",
      "Name: truthPartE, dtype: object\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 40%|████      | 2/5 [00:14<00:21,  7.20s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Match found! File = /clusterfs/ml4hep/mpettee/ml4pions/data/root_files/user.angerami.mc16_13TeV.900247.PG_singlepion_logE0p2to2000.e8312_e7400_s3170_r12383.v01-45-gaa27bcb_OutputStream/user.angerami.24559744.OutputStream._000461.root\n",
      "3985    [29.691723]\n",
      "Name: truthPartE, dtype: object\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 80%|████████  | 4/5 [00:28<00:07,  7.12s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Match found! File = /clusterfs/ml4hep/mpettee/ml4pions/data/root_files/user.angerami.mc16_13TeV.900247.PG_singlepion_logE0p2to2000.e8312_e7400_s3170_r12383.v01-45-gaa27bcb_OutputStream/user.angerami.24559744.OutputStream._000430.root\n",
      "8254    [499.59772]\n",
      "Name: truthPartE, dtype: object\n",
      "Match found! File = /clusterfs/ml4hep/mpettee/ml4pions/data/root_files/user.angerami.mc16_13TeV.900247.PG_singlepion_logE0p2to2000.e8312_e7400_s3170_r12383.v01-45-gaa27bcb_OutputStream/user.angerami.24559744.OutputStream._000499.root\n",
      "3592    [100.29169]\n",
      "Name: truthPartE, dtype: object\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5/5 [00:35<00:00,  7.08s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Match found! File = /clusterfs/ml4hep/mpettee/ml4pions/data/root_files/user.angerami.mc16_13TeV.900247.PG_singlepion_logE0p2to2000.e8312_e7400_s3170_r12383.v01-45-gaa27bcb_OutputStream/user.angerami.24559744.OutputStream._000499.root\n",
      "8425    [995.07166]\n",
      "Name: truthPartE, dtype: object\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "events = []\n",
    "for file in tqdm(selected_pion_files):\n",
    "    a = uproot.open(file)[\"EventTree\"].arrays(library = \"np\")\n",
    "    df = apply_cuts(a)\n",
    "    for energy in energies:\n",
    "        if len(df[np.abs(df.truthPartE.explode() - energy) < 0.00001].truthPartE) != 0:\n",
    "            print(\"Match found! File = {}\".format(file))\n",
    "            print(df[np.abs(df.truthPartE.explode() - energy) < 0.00001].truthPartE)\n",
    "            events.append(df[np.abs(df.truthPartE.explode() - energy) < 0.00001])\n",
    "# print(df.truthPartE.min())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "7cd71809-e2c3-4e6c-9bb7-3254fffc566b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-08-01T22:25:09.156245Z",
     "iopub.status.busy": "2022-08-01T22:25:09.155827Z",
     "iopub.status.idle": "2022-08-01T22:25:09.194065Z",
     "shell.execute_reply": "2022-08-01T22:25:09.193483Z",
     "shell.execute_reply.started": "2022-08-01T22:25:09.156213Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "df = pd.concat(events)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "4d0a4716-ccff-47ef-bb4b-83eec0ba5789",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-08-01T22:25:09.198844Z",
     "iopub.status.busy": "2022-08-01T22:25:09.198145Z",
     "iopub.status.idle": "2022-08-01T22:25:09.272190Z",
     "shell.execute_reply": "2022-08-01T22:25:09.271777Z",
     "shell.execute_reply.started": "2022-08-01T22:25:09.198809Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>index</th>\n",
       "      <th>runNumber</th>\n",
       "      <th>eventNumber</th>\n",
       "      <th>lumiBlock</th>\n",
       "      <th>coreFlags</th>\n",
       "      <th>mcEventNumber</th>\n",
       "      <th>mcChannelNumber</th>\n",
       "      <th>mcEventWeight</th>\n",
       "      <th>nTruthPart</th>\n",
       "      <th>G4PreCalo_n_EM</th>\n",
       "      <th>...</th>\n",
       "      <th>cluster_CENTER_LAMBDA</th>\n",
       "      <th>cluster_ISOLATION</th>\n",
       "      <th>cluster_ENERGY_DigiHSTruth</th>\n",
       "      <th>cluster_cell_ID</th>\n",
       "      <th>cluster_cell_E</th>\n",
       "      <th>cluster_hitsTruthIndex</th>\n",
       "      <th>cluster_hitsTruthE</th>\n",
       "      <th>deltaR</th>\n",
       "      <th>dR_pass</th>\n",
       "      <th>event_number</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>4109</th>\n",
       "      <td>6699</td>\n",
       "      <td>284500</td>\n",
       "      <td>8805468</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>8805468</td>\n",
       "      <td>900247</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>...</td>\n",
       "      <td>[16.759415]</td>\n",
       "      <td>[1.0]</td>\n",
       "      <td>[-999.0]</td>\n",
       "      <td>((754979952, 757113968, 757114480, 757115504, ...</td>\n",
       "      <td>((0.5277696, 0.058863826, 0.075782515, 0.01802...</td>\n",
       "      <td>((0))</td>\n",
       "      <td>((0.21558206))</td>\n",
       "      <td>[0.8515652624760468]</td>\n",
       "      <td>[True]</td>\n",
       "      <td>6699</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9082</th>\n",
       "      <td>15185</td>\n",
       "      <td>284500</td>\n",
       "      <td>1744200</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1744200</td>\n",
       "      <td>900247</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1</td>\n",
       "      <td>4</td>\n",
       "      <td>...</td>\n",
       "      <td>[274.6379]</td>\n",
       "      <td>[1.0]</td>\n",
       "      <td>[-999.0]</td>\n",
       "      <td>((759182878, 759182880, 759182882, 759182366, ...</td>\n",
       "      <td>((0.19118226, 0.33541742, 0.03791072, 0.040082...</td>\n",
       "      <td>((0))</td>\n",
       "      <td>((0.4949144))</td>\n",
       "      <td>[0.8796527826676044]</td>\n",
       "      <td>[True]</td>\n",
       "      <td>15185</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3985</th>\n",
       "      <td>6542</td>\n",
       "      <td>284500</td>\n",
       "      <td>6740038</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>6740038</td>\n",
       "      <td>900247</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1</td>\n",
       "      <td>7</td>\n",
       "      <td>...</td>\n",
       "      <td>[44.152363]</td>\n",
       "      <td>[1.0]</td>\n",
       "      <td>[-999.0]</td>\n",
       "      <td>((749207614, 749207612, 749207616, 748946494, ...</td>\n",
       "      <td>((0.12675342, 0.005030082, 0.015049332, 0.1322...</td>\n",
       "      <td>((0))</td>\n",
       "      <td>((0.3957773))</td>\n",
       "      <td>[0.4413612757326127]</td>\n",
       "      <td>[True]</td>\n",
       "      <td>6542</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8254</th>\n",
       "      <td>13421</td>\n",
       "      <td>284500</td>\n",
       "      <td>9453782</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>9453782</td>\n",
       "      <td>900247</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>[1655.8729, 280.56177, 951.5224, 886.79816, 17...</td>\n",
       "      <td>[0.5583132, 0.6243533, 0.51502, 0.67100334, 0....</td>\n",
       "      <td>[-999.0, -999.0, -999.0, -999.0, -999.0, -999....</td>\n",
       "      <td>((1150091280, 1150107664, 1141702672, 11500915...</td>\n",
       "      <td>((59.199345, 0.426856, 6.5276375, 0.8904916, 7...</td>\n",
       "      <td>((0), (0), (0), (0), (0), (0), (0), (0))</td>\n",
       "      <td>((129.5406), (47.10336), (6.7947793), (2.45943...</td>\n",
       "      <td>[0.028639940728819663, 0.012352409796186377, 0...</td>\n",
       "      <td>[True, True, True, True, True, True, True, True]</td>\n",
       "      <td>13421</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3592</th>\n",
       "      <td>5892</td>\n",
       "      <td>284500</td>\n",
       "      <td>6070536</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>6070536</td>\n",
       "      <td>900247</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1</td>\n",
       "      <td>9</td>\n",
       "      <td>...</td>\n",
       "      <td>[652.2834]</td>\n",
       "      <td>[1.0]</td>\n",
       "      <td>[-999.0]</td>\n",
       "      <td>((752879606, 752879608, 752880118, 752879096, ...</td>\n",
       "      <td>((0.5086221, 0.051815987, 0.32129198, 0.112986...</td>\n",
       "      <td>((0))</td>\n",
       "      <td>((1.3566374))</td>\n",
       "      <td>[0.9779932768951003]</td>\n",
       "      <td>[True]</td>\n",
       "      <td>5892</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8425</th>\n",
       "      <td>13863</td>\n",
       "      <td>284500</td>\n",
       "      <td>2149922</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>2149922</td>\n",
       "      <td>900247</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>[3134.7249, 1667.5095, 2183.9583, 2426.0288]</td>\n",
       "      <td>[0.78668576, 0.6643609, 0.52810866, 0.75504273]</td>\n",
       "      <td>[-999.0, -999.0, -999.0, -999.0]</td>\n",
       "      <td>((1208846880, 1208830496, 1208863264, 12088473...</td>\n",
       "      <td>((101.894714, 3.6700416, 1.9256063, 0.15658991...</td>\n",
       "      <td>((0), (0), (0), (0))</td>\n",
       "      <td>((134.67303), (5.35034), (4.1942773), (0.33746...</td>\n",
       "      <td>[0.028517369896244267, 0.07940313518030324, 0....</td>\n",
       "      <td>[True, True, True, True]</td>\n",
       "      <td>13863</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>6 rows × 129 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "      index  runNumber  eventNumber  lumiBlock  coreFlags  mcEventNumber  \\\n",
       "4109   6699     284500      8805468          1          0        8805468   \n",
       "9082  15185     284500      1744200          1          0        1744200   \n",
       "3985   6542     284500      6740038          1          0        6740038   \n",
       "8254  13421     284500      9453782          1          0        9453782   \n",
       "3592   5892     284500      6070536          1          0        6070536   \n",
       "8425  13863     284500      2149922          1          0        2149922   \n",
       "\n",
       "      mcChannelNumber  mcEventWeight  nTruthPart  G4PreCalo_n_EM  ...  \\\n",
       "4109           900247            1.0           1               1  ...   \n",
       "9082           900247            1.0           1               4  ...   \n",
       "3985           900247            1.0           1               7  ...   \n",
       "8254           900247            1.0           1               0  ...   \n",
       "3592           900247            1.0           1               9  ...   \n",
       "8425           900247            1.0           1               0  ...   \n",
       "\n",
       "                                  cluster_CENTER_LAMBDA  \\\n",
       "4109                                        [16.759415]   \n",
       "9082                                         [274.6379]   \n",
       "3985                                        [44.152363]   \n",
       "8254  [1655.8729, 280.56177, 951.5224, 886.79816, 17...   \n",
       "3592                                         [652.2834]   \n",
       "8425       [3134.7249, 1667.5095, 2183.9583, 2426.0288]   \n",
       "\n",
       "                                      cluster_ISOLATION  \\\n",
       "4109                                              [1.0]   \n",
       "9082                                              [1.0]   \n",
       "3985                                              [1.0]   \n",
       "8254  [0.5583132, 0.6243533, 0.51502, 0.67100334, 0....   \n",
       "3592                                              [1.0]   \n",
       "8425    [0.78668576, 0.6643609, 0.52810866, 0.75504273]   \n",
       "\n",
       "                             cluster_ENERGY_DigiHSTruth  \\\n",
       "4109                                           [-999.0]   \n",
       "9082                                           [-999.0]   \n",
       "3985                                           [-999.0]   \n",
       "8254  [-999.0, -999.0, -999.0, -999.0, -999.0, -999....   \n",
       "3592                                           [-999.0]   \n",
       "8425                   [-999.0, -999.0, -999.0, -999.0]   \n",
       "\n",
       "                                        cluster_cell_ID  \\\n",
       "4109  ((754979952, 757113968, 757114480, 757115504, ...   \n",
       "9082  ((759182878, 759182880, 759182882, 759182366, ...   \n",
       "3985  ((749207614, 749207612, 749207616, 748946494, ...   \n",
       "8254  ((1150091280, 1150107664, 1141702672, 11500915...   \n",
       "3592  ((752879606, 752879608, 752880118, 752879096, ...   \n",
       "8425  ((1208846880, 1208830496, 1208863264, 12088473...   \n",
       "\n",
       "                                         cluster_cell_E  \\\n",
       "4109  ((0.5277696, 0.058863826, 0.075782515, 0.01802...   \n",
       "9082  ((0.19118226, 0.33541742, 0.03791072, 0.040082...   \n",
       "3985  ((0.12675342, 0.005030082, 0.015049332, 0.1322...   \n",
       "8254  ((59.199345, 0.426856, 6.5276375, 0.8904916, 7...   \n",
       "3592  ((0.5086221, 0.051815987, 0.32129198, 0.112986...   \n",
       "8425  ((101.894714, 3.6700416, 1.9256063, 0.15658991...   \n",
       "\n",
       "                        cluster_hitsTruthIndex  \\\n",
       "4109                                     ((0))   \n",
       "9082                                     ((0))   \n",
       "3985                                     ((0))   \n",
       "8254  ((0), (0), (0), (0), (0), (0), (0), (0))   \n",
       "3592                                     ((0))   \n",
       "8425                      ((0), (0), (0), (0))   \n",
       "\n",
       "                                     cluster_hitsTruthE  \\\n",
       "4109                                     ((0.21558206))   \n",
       "9082                                      ((0.4949144))   \n",
       "3985                                      ((0.3957773))   \n",
       "8254  ((129.5406), (47.10336), (6.7947793), (2.45943...   \n",
       "3592                                      ((1.3566374))   \n",
       "8425  ((134.67303), (5.35034), (4.1942773), (0.33746...   \n",
       "\n",
       "                                                 deltaR  \\\n",
       "4109                               [0.8515652624760468]   \n",
       "9082                               [0.8796527826676044]   \n",
       "3985                               [0.4413612757326127]   \n",
       "8254  [0.028639940728819663, 0.012352409796186377, 0...   \n",
       "3592                               [0.9779932768951003]   \n",
       "8425  [0.028517369896244267, 0.07940313518030324, 0....   \n",
       "\n",
       "                                               dR_pass event_number  \n",
       "4109                                            [True]         6699  \n",
       "9082                                            [True]        15185  \n",
       "3985                                            [True]         6542  \n",
       "8254  [True, True, True, True, True, True, True, True]        13421  \n",
       "3592                                            [True]         5892  \n",
       "8425                          [True, True, True, True]        13863  \n",
       "\n",
       "[6 rows x 129 columns]"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "05407d2f-04db-4a59-9bd8-0a3b56f5d1ea",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-08-01T22:25:40.296091Z",
     "iopub.status.busy": "2022-08-01T22:25:40.295744Z",
     "iopub.status.idle": "2022-08-01T22:25:40.339162Z",
     "shell.execute_reply": "2022-08-01T22:25:40.338268Z",
     "shell.execute_reply.started": "2022-08-01T22:25:40.296059Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/clusterfs/ml4hep/mpettee/miniconda3/envs/nbdev/lib/python3.9/site-packages/pandas/core/generic.py:2703: PerformanceWarning: \n",
      "your performance may suffer as PyTables will pickle object types that it cannot\n",
      "map directly to c-types [inferred_type->mixed,key->block3_values] [items->Index(['truthPartPdgId', 'truthPartStatus', 'truthPartBarcode', 'truthPartPt',\n",
      "       'truthPartE', 'truthPartMass', 'truthPartEta', 'truthPartPhi',\n",
      "       'trackPt', 'trackP',\n",
      "       ...\n",
      "       'cluster_FIRST_ENG_DENS', 'cluster_CENTER_LAMBDA', 'cluster_ISOLATION',\n",
      "       'cluster_ENERGY_DigiHSTruth', 'cluster_cell_ID', 'cluster_cell_E',\n",
      "       'cluster_hitsTruthIndex', 'cluster_hitsTruthE', 'deltaR', 'dR_pass'],\n",
      "      dtype='object', length=110)]\n",
      "\n",
      "  pytables.to_hdf(\n"
     ]
    }
   ],
   "source": [
    "df.to_hdf(\"df_transformer_debug.h5\", \"df\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "fb3ec913-d5b5-4588-85eb-4bd421eb9a53",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-08-08T22:10:06.541982Z",
     "iopub.status.busy": "2022-08-08T22:10:06.541768Z",
     "iopub.status.idle": "2022-08-08T22:10:06.582202Z",
     "shell.execute_reply": "2022-08-08T22:10:06.581438Z",
     "shell.execute_reply.started": "2022-08-08T22:10:06.541962Z"
    }
   },
   "outputs": [],
   "source": [
    "df = pd.read_hdf(\"df_transformer_debug.h5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "bea76150-d44f-43ca-a6b1-f11223932413",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-08-08T22:10:08.814292Z",
     "iopub.status.busy": "2022-08-08T22:10:08.814110Z",
     "iopub.status.idle": "2022-08-08T22:10:08.819256Z",
     "shell.execute_reply": "2022-08-08T22:10:08.818685Z",
     "shell.execute_reply.started": "2022-08-08T22:10:08.814272Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Index(['index', 'runNumber', 'eventNumber', 'lumiBlock', 'coreFlags',\n",
       "       'mcEventNumber', 'mcChannelNumber', 'mcEventWeight', 'nTruthPart',\n",
       "       'G4PreCalo_n_EM',\n",
       "       ...\n",
       "       'cluster_CENTER_LAMBDA', 'cluster_ISOLATION',\n",
       "       'cluster_ENERGY_DigiHSTruth', 'cluster_cell_ID', 'cluster_cell_E',\n",
       "       'cluster_hitsTruthIndex', 'cluster_hitsTruthE', 'deltaR', 'dR_pass',\n",
       "       'event_number'],\n",
       "      dtype='object', length=129)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "3ae6ac76-36bf-4d7d-aee7-13c011d584dc",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-08-08T22:19:27.223228Z",
     "iopub.status.busy": "2022-08-08T22:19:27.223019Z",
     "iopub.status.idle": "2022-08-08T22:19:27.228772Z",
     "shell.execute_reply": "2022-08-08T22:19:27.227995Z",
     "shell.execute_reply.started": "2022-08-08T22:19:27.223209Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['truthPartEta',\n",
       " 'trackEta',\n",
       " 'trackEta_PreSamplerB',\n",
       " 'trackEta_PreSamplerE',\n",
       " 'trackEta_EMB1',\n",
       " 'trackEta_EMB2',\n",
       " 'trackEta_EMB3',\n",
       " 'trackEta_EME1',\n",
       " 'trackEta_EME2',\n",
       " 'trackEta_EME3',\n",
       " 'trackEta_HEC0',\n",
       " 'trackEta_HEC1',\n",
       " 'trackEta_HEC2',\n",
       " 'trackEta_HEC3',\n",
       " 'trackEta_TileBar0',\n",
       " 'trackEta_TileBar1',\n",
       " 'trackEta_TileBar2',\n",
       " 'trackEta_TileGap1',\n",
       " 'trackEta_TileGap2',\n",
       " 'trackEta_TileGap3',\n",
       " 'trackEta_TileExt0',\n",
       " 'trackEta_TileExt1',\n",
       " 'trackEta_TileExt2',\n",
       " 'AntiKt4EMTopoJetsEta',\n",
       " 'AntiKt4LCTopoJetsEta',\n",
       " 'AntiKt4TruthJetsEta',\n",
       " 'cluster_Eta']"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[key for key in df.keys() if \"Eta\" in key]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "1d678621-a2a0-459e-83b9-89bd1b4e3acd",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-08-08T22:20:12.054481Z",
     "iopub.status.busy": "2022-08-08T22:20:12.054177Z",
     "iopub.status.idle": "2022-08-08T22:20:12.102581Z",
     "shell.execute_reply": "2022-08-08T22:20:12.101899Z",
     "shell.execute_reply.started": "2022-08-08T22:20:12.054450Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>cluster_E</th>\n",
       "      <th>cluster_Eta</th>\n",
       "      <th>trackEta</th>\n",
       "      <th>truthPartEta</th>\n",
       "      <th>cluster_Phi</th>\n",
       "      <th>trackPhi</th>\n",
       "      <th>truthPartPhi</th>\n",
       "      <th>cluster_EM_PROBABILITY</th>\n",
       "      <th>cluster_HAD_WEIGHT</th>\n",
       "      <th>cluster_OOC_WEIGHT</th>\n",
       "      <th>cluster_DM_WEIGHT</th>\n",
       "      <th>cluster_CENTER_MAG</th>\n",
       "      <th>cluster_FIRST_ENG_DENS</th>\n",
       "      <th>cluster_CENTER_LAMBDA</th>\n",
       "      <th>cluster_ISOLATION</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>4109</th>\n",
       "      <td>[0.5934837]</td>\n",
       "      <td>[-0.25973687]</td>\n",
       "      <td>[-0.2805243]</td>\n",
       "      <td>[1.607805]</td>\n",
       "      <td>[-0.74060816]</td>\n",
       "      <td>[-1.5919197]</td>\n",
       "      <td>[-1.7966828]</td>\n",
       "      <td>[0.49211332]</td>\n",
       "      <td>[0.98009163]</td>\n",
       "      <td>[1.1358624]</td>\n",
       "      <td>[1.0]</td>\n",
       "      <td>[1536.88]</td>\n",
       "      <td>[4.742205e-07]</td>\n",
       "      <td>[16.759415]</td>\n",
       "      <td>[1.0]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9082</th>\n",
       "      <td>[0.82536364]</td>\n",
       "      <td>[-0.6838584]</td>\n",
       "      <td>[-0.628881]</td>\n",
       "      <td>[-2.25457]</td>\n",
       "      <td>[0.3976988]</td>\n",
       "      <td>[-0.4802343]</td>\n",
       "      <td>[-0.19429713]</td>\n",
       "      <td>[0.06517576]</td>\n",
       "      <td>[1.0400422]</td>\n",
       "      <td>[1.9183998]</td>\n",
       "      <td>[1.0502725]</td>\n",
       "      <td>[2111.5806]</td>\n",
       "      <td>[2.865492e-07]</td>\n",
       "      <td>[274.6379]</td>\n",
       "      <td>[1.0]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3985</th>\n",
       "      <td>[0.66128916]</td>\n",
       "      <td>[-1.5136156]</td>\n",
       "      <td>[-1.9086738]</td>\n",
       "      <td>[-1.9151504]</td>\n",
       "      <td>[3.1018682]</td>\n",
       "      <td>[2.9050717]</td>\n",
       "      <td>[2.8025804]</td>\n",
       "      <td>[0.5116583]</td>\n",
       "      <td>[0.97907114]</td>\n",
       "      <td>[2.1411436]</td>\n",
       "      <td>[1.288497]</td>\n",
       "      <td>[4174.968]</td>\n",
       "      <td>[8.01778e-07]</td>\n",
       "      <td>[44.152363]</td>\n",
       "      <td>[1.0]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8254</th>\n",
       "      <td>[131.32864, 42.44467, 5.466173, 3.823332, 2.44...</td>\n",
       "      <td>[0.02523974, 0.032032937, 0.013351178, 0.01600...</td>\n",
       "      <td>[0.019860815]</td>\n",
       "      <td>[0.019807491]</td>\n",
       "      <td>[-1.1544378, -1.1846708, -1.2848923, -0.943424...</td>\n",
       "      <td>[-1.1825681]</td>\n",
       "      <td>[-1.1825013]</td>\n",
       "      <td>[0.000458727, 0.05385432, 0.0013848847, 0.0008...</td>\n",
       "      <td>[1.1211182, 1.0571781, 1.1965067, 1.2111669, 1...</td>\n",
       "      <td>[1.005254, 1.0529661, 1.0632713, 1.132284, 1.0...</td>\n",
       "      <td>[1.1125239, 1.0070257, 1.5859705, 1.2677855, 1...</td>\n",
       "      <td>[3159.362, 1778.8763, 2434.2268, 2389.474, 323...</td>\n",
       "      <td>[4.771859e-07, 9.329076e-06, 6.883655e-08, 5.0...</td>\n",
       "      <td>[1655.8729, 280.56177, 951.5224, 886.79816, 17...</td>\n",
       "      <td>[0.5583132, 0.6243533, 0.51502, 0.67100334, 0....</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3592</th>\n",
       "      <td>[1.7634326]</td>\n",
       "      <td>[-1.6830969]</td>\n",
       "      <td>[-2.182635]</td>\n",
       "      <td>[-2.9042387]</td>\n",
       "      <td>[-0.09693346]</td>\n",
       "      <td>[0.7438593]</td>\n",
       "      <td>[1.1400995]</td>\n",
       "      <td>[0.001]</td>\n",
       "      <td>[1.2155149]</td>\n",
       "      <td>[1.8951625]</td>\n",
       "      <td>[1.3286014]</td>\n",
       "      <td>[4589.337]</td>\n",
       "      <td>[5.8728745e-07]</td>\n",
       "      <td>[652.2834]</td>\n",
       "      <td>[1.0]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8425</th>\n",
       "      <td>[119.363686, 7.9105496, 5.7105308, 3.256291]</td>\n",
       "      <td>[-1.0072887, -1.0556288, -1.0082572, -1.1504852]</td>\n",
       "      <td>[-0.97878176]</td>\n",
       "      <td>[-0.9791123]</td>\n",
       "      <td>[-0.9378377, -0.9186243, -0.77520394, -1.1300832]</td>\n",
       "      <td>[-0.9386091]</td>\n",
       "      <td>[-0.9387261]</td>\n",
       "      <td>[0.001, 0.0003290493, 0.0006894411, 0.02137577]</td>\n",
       "      <td>[1.1482576, 1.0881703, 1.2168978, 1.0106353]</td>\n",
       "      <td>[1.0178477, 1.1170146, 1.0841136, 1.0658295]</td>\n",
       "      <td>[1.0379997, 1.0457845, 1.065008, 1.1431478]</td>\n",
       "      <td>[5359.1133, 4091.796, 4526.9287, 4245.3477]</td>\n",
       "      <td>[3.2673833e-07, 7.6536173e-07, 1.10334085e-07,...</td>\n",
       "      <td>[3134.7249, 1667.5095, 2183.9583, 2426.0288]</td>\n",
       "      <td>[0.78668576, 0.6643609, 0.52810866, 0.75504273]</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                              cluster_E  \\\n",
       "4109                                        [0.5934837]   \n",
       "9082                                       [0.82536364]   \n",
       "3985                                       [0.66128916]   \n",
       "8254  [131.32864, 42.44467, 5.466173, 3.823332, 2.44...   \n",
       "3592                                        [1.7634326]   \n",
       "8425       [119.363686, 7.9105496, 5.7105308, 3.256291]   \n",
       "\n",
       "                                            cluster_Eta       trackEta  \\\n",
       "4109                                      [-0.25973687]   [-0.2805243]   \n",
       "9082                                       [-0.6838584]    [-0.628881]   \n",
       "3985                                       [-1.5136156]   [-1.9086738]   \n",
       "8254  [0.02523974, 0.032032937, 0.013351178, 0.01600...  [0.019860815]   \n",
       "3592                                       [-1.6830969]    [-2.182635]   \n",
       "8425   [-1.0072887, -1.0556288, -1.0082572, -1.1504852]  [-0.97878176]   \n",
       "\n",
       "       truthPartEta                                        cluster_Phi  \\\n",
       "4109     [1.607805]                                      [-0.74060816]   \n",
       "9082     [-2.25457]                                        [0.3976988]   \n",
       "3985   [-1.9151504]                                        [3.1018682]   \n",
       "8254  [0.019807491]  [-1.1544378, -1.1846708, -1.2848923, -0.943424...   \n",
       "3592   [-2.9042387]                                      [-0.09693346]   \n",
       "8425   [-0.9791123]  [-0.9378377, -0.9186243, -0.77520394, -1.1300832]   \n",
       "\n",
       "          trackPhi   truthPartPhi  \\\n",
       "4109  [-1.5919197]   [-1.7966828]   \n",
       "9082  [-0.4802343]  [-0.19429713]   \n",
       "3985   [2.9050717]    [2.8025804]   \n",
       "8254  [-1.1825681]   [-1.1825013]   \n",
       "3592   [0.7438593]    [1.1400995]   \n",
       "8425  [-0.9386091]   [-0.9387261]   \n",
       "\n",
       "                                 cluster_EM_PROBABILITY  \\\n",
       "4109                                       [0.49211332]   \n",
       "9082                                       [0.06517576]   \n",
       "3985                                        [0.5116583]   \n",
       "8254  [0.000458727, 0.05385432, 0.0013848847, 0.0008...   \n",
       "3592                                            [0.001]   \n",
       "8425    [0.001, 0.0003290493, 0.0006894411, 0.02137577]   \n",
       "\n",
       "                                     cluster_HAD_WEIGHT  \\\n",
       "4109                                       [0.98009163]   \n",
       "9082                                        [1.0400422]   \n",
       "3985                                       [0.97907114]   \n",
       "8254  [1.1211182, 1.0571781, 1.1965067, 1.2111669, 1...   \n",
       "3592                                        [1.2155149]   \n",
       "8425       [1.1482576, 1.0881703, 1.2168978, 1.0106353]   \n",
       "\n",
       "                                     cluster_OOC_WEIGHT  \\\n",
       "4109                                        [1.1358624]   \n",
       "9082                                        [1.9183998]   \n",
       "3985                                        [2.1411436]   \n",
       "8254  [1.005254, 1.0529661, 1.0632713, 1.132284, 1.0...   \n",
       "3592                                        [1.8951625]   \n",
       "8425       [1.0178477, 1.1170146, 1.0841136, 1.0658295]   \n",
       "\n",
       "                                      cluster_DM_WEIGHT  \\\n",
       "4109                                              [1.0]   \n",
       "9082                                        [1.0502725]   \n",
       "3985                                         [1.288497]   \n",
       "8254  [1.1125239, 1.0070257, 1.5859705, 1.2677855, 1...   \n",
       "3592                                        [1.3286014]   \n",
       "8425        [1.0379997, 1.0457845, 1.065008, 1.1431478]   \n",
       "\n",
       "                                     cluster_CENTER_MAG  \\\n",
       "4109                                          [1536.88]   \n",
       "9082                                        [2111.5806]   \n",
       "3985                                         [4174.968]   \n",
       "8254  [3159.362, 1778.8763, 2434.2268, 2389.474, 323...   \n",
       "3592                                         [4589.337]   \n",
       "8425        [5359.1133, 4091.796, 4526.9287, 4245.3477]   \n",
       "\n",
       "                                 cluster_FIRST_ENG_DENS  \\\n",
       "4109                                     [4.742205e-07]   \n",
       "9082                                     [2.865492e-07]   \n",
       "3985                                      [8.01778e-07]   \n",
       "8254  [4.771859e-07, 9.329076e-06, 6.883655e-08, 5.0...   \n",
       "3592                                    [5.8728745e-07]   \n",
       "8425  [3.2673833e-07, 7.6536173e-07, 1.10334085e-07,...   \n",
       "\n",
       "                                  cluster_CENTER_LAMBDA  \\\n",
       "4109                                        [16.759415]   \n",
       "9082                                         [274.6379]   \n",
       "3985                                        [44.152363]   \n",
       "8254  [1655.8729, 280.56177, 951.5224, 886.79816, 17...   \n",
       "3592                                         [652.2834]   \n",
       "8425       [3134.7249, 1667.5095, 2183.9583, 2426.0288]   \n",
       "\n",
       "                                      cluster_ISOLATION  \n",
       "4109                                              [1.0]  \n",
       "9082                                              [1.0]  \n",
       "3985                                              [1.0]  \n",
       "8254  [0.5583132, 0.6243533, 0.51502, 0.67100334, 0....  \n",
       "3592                                              [1.0]  \n",
       "8425    [0.78668576, 0.6643609, 0.52810866, 0.75504273]  "
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df[['cluster_E','cluster_Eta','trackEta', 'truthPartEta', 'cluster_Phi','trackPhi', 'truthPartPhi','cluster_EM_PROBABILITY', 'cluster_HAD_WEIGHT', 'cluster_OOC_WEIGHT',\n",
    "               'cluster_DM_WEIGHT', 'cluster_CENTER_MAG', 'cluster_FIRST_ENG_DENS', \n",
    "               'cluster_CENTER_LAMBDA', 'cluster_ISOLATION'\n",
    "              ]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "907ad84b-b20c-4e13-831f-b05ce2df3d9d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-08-08T22:10:19.142968Z",
     "iopub.status.busy": "2022-08-08T22:10:19.142760Z",
     "iopub.status.idle": "2022-08-08T22:10:19.182347Z",
     "shell.execute_reply": "2022-08-08T22:10:19.181859Z",
     "shell.execute_reply.started": "2022-08-08T22:10:19.142949Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>truthPartE</th>\n",
       "      <th>trackPt</th>\n",
       "      <th>trackP</th>\n",
       "      <th>trackMass</th>\n",
       "      <th>trackEta</th>\n",
       "      <th>trackPhi</th>\n",
       "      <th>trackNumberOfPixelHits</th>\n",
       "      <th>trackNumberOfSCTHits</th>\n",
       "      <th>trackNumberOfPixelDeadSensors</th>\n",
       "      <th>trackNumberOfSCTDeadSensors</th>\n",
       "      <th>trackNumberOfInnermostPixelLayerHits</th>\n",
       "      <th>trackNumberOfNextToInnermostPixelLayerHits</th>\n",
       "      <th>trackExpectInnermostPixelLayerHit</th>\n",
       "      <th>trackExpectNextToInnermostPixelLayerHit</th>\n",
       "      <th>trackNumberOfTRTHits</th>\n",
       "      <th>trackNumberOfTRTOutliers</th>\n",
       "      <th>trackChiSquared</th>\n",
       "      <th>trackNumberDOF</th>\n",
       "      <th>trackD0</th>\n",
       "      <th>trackZ0</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>4109</th>\n",
       "      <td>[1.638227]</td>\n",
       "      <td>[0.56499124]</td>\n",
       "      <td>[0.5873681]</td>\n",
       "      <td>[0.13957007]</td>\n",
       "      <td>[-0.2805243]</td>\n",
       "      <td>[-1.5919197]</td>\n",
       "      <td>[3]</td>\n",
       "      <td>[7]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[1]</td>\n",
       "      <td>[1]</td>\n",
       "      <td>[1]</td>\n",
       "      <td>[25]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[45.287254]</td>\n",
       "      <td>[33]</td>\n",
       "      <td>[-7.4432144]</td>\n",
       "      <td>[87.2737]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9082</th>\n",
       "      <td>[9.119415]</td>\n",
       "      <td>[0.59222865]</td>\n",
       "      <td>[0.7132502]</td>\n",
       "      <td>[0.13957018]</td>\n",
       "      <td>[-0.628881]</td>\n",
       "      <td>[-0.4802343]</td>\n",
       "      <td>[5]</td>\n",
       "      <td>[8]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[2]</td>\n",
       "      <td>[1]</td>\n",
       "      <td>[1]</td>\n",
       "      <td>[1]</td>\n",
       "      <td>[28]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[48.939526]</td>\n",
       "      <td>[41]</td>\n",
       "      <td>[6.6680737]</td>\n",
       "      <td>[27.934269]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3985</th>\n",
       "      <td>[29.691723]</td>\n",
       "      <td>[0.7130422]</td>\n",
       "      <td>[2.4572914]</td>\n",
       "      <td>[0.13957112]</td>\n",
       "      <td>[-1.9086738]</td>\n",
       "      <td>[2.9050717]</td>\n",
       "      <td>[3]</td>\n",
       "      <td>[10]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[1]</td>\n",
       "      <td>[1]</td>\n",
       "      <td>[1]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[17.507769]</td>\n",
       "      <td>[11]</td>\n",
       "      <td>[-3.952645]</td>\n",
       "      <td>[22.175701]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8254</th>\n",
       "      <td>[499.59772]</td>\n",
       "      <td>[436.06116]</td>\n",
       "      <td>[436.14716]</td>\n",
       "      <td>[0.16326419]</td>\n",
       "      <td>[0.019860815]</td>\n",
       "      <td>[-1.1825681]</td>\n",
       "      <td>[5]</td>\n",
       "      <td>[8]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[2]</td>\n",
       "      <td>[1]</td>\n",
       "      <td>[1]</td>\n",
       "      <td>[1]</td>\n",
       "      <td>[31]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[40.224117]</td>\n",
       "      <td>[44]</td>\n",
       "      <td>[0.0017354734]</td>\n",
       "      <td>[25.76988]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3592</th>\n",
       "      <td>[100.29169]</td>\n",
       "      <td>[1.1556271]</td>\n",
       "      <td>[5.1901484]</td>\n",
       "      <td>[0.13957529]</td>\n",
       "      <td>[-2.182635]</td>\n",
       "      <td>[0.7438593]</td>\n",
       "      <td>[5]</td>\n",
       "      <td>[10]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[2]</td>\n",
       "      <td>[1]</td>\n",
       "      <td>[1]</td>\n",
       "      <td>[1]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[12.896353]</td>\n",
       "      <td>[15]</td>\n",
       "      <td>[9.772532]</td>\n",
       "      <td>[-44.83223]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8425</th>\n",
       "      <td>[995.07166]</td>\n",
       "      <td>[1118.2161]</td>\n",
       "      <td>[1698.0002]</td>\n",
       "      <td>[0.36632022]</td>\n",
       "      <td>[-0.97878176]</td>\n",
       "      <td>[-0.9386091]</td>\n",
       "      <td>[4]</td>\n",
       "      <td>[9]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[1]</td>\n",
       "      <td>[1]</td>\n",
       "      <td>[1]</td>\n",
       "      <td>[1]</td>\n",
       "      <td>[20]</td>\n",
       "      <td>[0]</td>\n",
       "      <td>[23.839802]</td>\n",
       "      <td>[32]</td>\n",
       "      <td>[-0.004356428]</td>\n",
       "      <td>[-1.2269732]</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       truthPartE       trackPt       trackP     trackMass       trackEta  \\\n",
       "4109   [1.638227]  [0.56499124]  [0.5873681]  [0.13957007]   [-0.2805243]   \n",
       "9082   [9.119415]  [0.59222865]  [0.7132502]  [0.13957018]    [-0.628881]   \n",
       "3985  [29.691723]   [0.7130422]  [2.4572914]  [0.13957112]   [-1.9086738]   \n",
       "8254  [499.59772]   [436.06116]  [436.14716]  [0.16326419]  [0.019860815]   \n",
       "3592  [100.29169]   [1.1556271]  [5.1901484]  [0.13957529]    [-2.182635]   \n",
       "8425  [995.07166]   [1118.2161]  [1698.0002]  [0.36632022]  [-0.97878176]   \n",
       "\n",
       "          trackPhi trackNumberOfPixelHits trackNumberOfSCTHits  \\\n",
       "4109  [-1.5919197]                    [3]                  [7]   \n",
       "9082  [-0.4802343]                    [5]                  [8]   \n",
       "3985   [2.9050717]                    [3]                 [10]   \n",
       "8254  [-1.1825681]                    [5]                  [8]   \n",
       "3592   [0.7438593]                    [5]                 [10]   \n",
       "8425  [-0.9386091]                    [4]                  [9]   \n",
       "\n",
       "     trackNumberOfPixelDeadSensors trackNumberOfSCTDeadSensors  \\\n",
       "4109                           [0]                         [0]   \n",
       "9082                           [0]                         [0]   \n",
       "3985                           [0]                         [0]   \n",
       "8254                           [0]                         [0]   \n",
       "3592                           [0]                         [0]   \n",
       "8425                           [0]                         [0]   \n",
       "\n",
       "     trackNumberOfInnermostPixelLayerHits  \\\n",
       "4109                                  [0]   \n",
       "9082                                  [2]   \n",
       "3985                                  [0]   \n",
       "8254                                  [2]   \n",
       "3592                                  [2]   \n",
       "8425                                  [1]   \n",
       "\n",
       "     trackNumberOfNextToInnermostPixelLayerHits  \\\n",
       "4109                                        [1]   \n",
       "9082                                        [1]   \n",
       "3985                                        [1]   \n",
       "8254                                        [1]   \n",
       "3592                                        [1]   \n",
       "8425                                        [1]   \n",
       "\n",
       "     trackExpectInnermostPixelLayerHit  \\\n",
       "4109                               [1]   \n",
       "9082                               [1]   \n",
       "3985                               [1]   \n",
       "8254                               [1]   \n",
       "3592                               [1]   \n",
       "8425                               [1]   \n",
       "\n",
       "     trackExpectNextToInnermostPixelLayerHit trackNumberOfTRTHits  \\\n",
       "4109                                     [1]                 [25]   \n",
       "9082                                     [1]                 [28]   \n",
       "3985                                     [1]                  [0]   \n",
       "8254                                     [1]                 [31]   \n",
       "3592                                     [1]                  [0]   \n",
       "8425                                     [1]                 [20]   \n",
       "\n",
       "     trackNumberOfTRTOutliers trackChiSquared trackNumberDOF         trackD0  \\\n",
       "4109                      [0]     [45.287254]           [33]    [-7.4432144]   \n",
       "9082                      [0]     [48.939526]           [41]     [6.6680737]   \n",
       "3985                      [0]     [17.507769]           [11]     [-3.952645]   \n",
       "8254                      [0]     [40.224117]           [44]  [0.0017354734]   \n",
       "3592                      [0]     [12.896353]           [15]      [9.772532]   \n",
       "8425                      [0]     [23.839802]           [32]  [-0.004356428]   \n",
       "\n",
       "           trackZ0  \n",
       "4109     [87.2737]  \n",
       "9082   [27.934269]  \n",
       "3985   [22.175701]  \n",
       "8254    [25.76988]  \n",
       "3592   [-44.83223]  \n",
       "8425  [-1.2269732]  "
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df[[\n",
    "    'truthPartE',\n",
    "    'trackPt',\n",
    "             'trackP',\n",
    "             'trackMass',\n",
    "             'trackEta',\n",
    "             'trackPhi',\n",
    "             'trackNumberOfPixelHits',\n",
    "             'trackNumberOfSCTHits',\n",
    "             'trackNumberOfPixelDeadSensors',\n",
    "             'trackNumberOfSCTDeadSensors',\n",
    "#              'trackNumberOfPixelSharedHits',\n",
    "#              'trackNumberOfSCTSharedHits',\n",
    "#              'trackNumberOfPixelHoles',\n",
    "#              'trackNumberOfSCTHoles',\n",
    "             'trackNumberOfInnermostPixelLayerHits',\n",
    "             'trackNumberOfNextToInnermostPixelLayerHits',\n",
    "             'trackExpectInnermostPixelLayerHit',\n",
    "             'trackExpectNextToInnermostPixelLayerHit',\n",
    "             'trackNumberOfTRTHits',\n",
    "             'trackNumberOfTRTOutliers',\n",
    "             'trackChiSquared',\n",
    "             'trackNumberDOF',\n",
    "             'trackD0',\n",
    "             'trackZ0'\n",
    "            ]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1788da2d-6de4-43a2-9823-45ddf2d3aed8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df = pd.DataFrame(a)\n",
    "# df = df[(df.nTrack == 1)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "136e97f3-59f9-43f9-b838-bb58d3799da1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "np.mean(df.cluster_nCells.explode())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39304907-6f7c-48da-aa84-6293248270f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "for row in df.index:\n",
    "    np.sqrt((df['cluster_Eta'][row].astype('float') - \n",
    "             df['trackEta'][row].astype('float'))**2 + \n",
    "                           delta_phi(df['cluster_Phi'][row].astype('float'), \n",
    "                                     df['trackPhi'][row].astype('float'))**2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "670f2634-fc1a-4240-80b5-d8262c8654e3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df[\"truthPartE\"] = df.truthPartE.explode()\n",
    "df[\"trackPt\"] = df.trackPt.explode()\n",
    "df[\"trackChiSquared\"] = df.trackChiSquared.explode()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d36b02ba-17cb-4be4-9e2b-5b114f6a1bed",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dfs = []\n",
    "for file in tqdm(pion_files[:50]):\n",
    "    a = uproot.open(file)[\"EventTree\"].arrays(library = \"np\")\n",
    "    dff = pd.DataFrame(a)\n",
    "    dff[\"nCluster\"] = dff.nCluster.explode()\n",
    "    dff[\"pion_eta\"] = dff.truthPartEta.explode()\n",
    "    dfs.append(dff)\n",
    "df_pion = pd.concat(dfs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ee50304-6ef6-4dd7-8257-bde0a8d09d0e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dfs = []\n",
    "for file in tqdm(pi0_files[:50]):\n",
    "    a = uproot.open(file)[\"EventTree\"].arrays(library = \"np\")\n",
    "    dff = pd.DataFrame(a)\n",
    "    dff[\"nCluster\"] = dff.nCluster.explode()\n",
    "    pion_eta = []\n",
    "    for row in dff.index:\n",
    "        pion_eta.append(dff.truthPartEta[row][0])\n",
    "    dff[\"pion_eta\"] = np.array(pion_eta)\n",
    "    dfs.append(dff)\n",
    "df_pi0 = pd.concat(dfs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4eb86584-874b-40a6-9823-1d70467d5442",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from scipy import stats \n",
    "binned_stats_pion = stats.binned_statistic(np.array(np.abs(df_pion.pion_eta), dtype=float), \n",
    "                                           np.array(df_pion.nCluster, dtype=float), \n",
    "                                           statistic='mean', bins=5)\n",
    "\n",
    "binned_stats_pi0 = stats.binned_statistic(np.array(np.abs(df_pi0.pion_eta), dtype=float), \n",
    "                                           np.array(df_pi0.nCluster, dtype=float), \n",
    "                                           statistic='mean', bins=5)\n",
    "\n",
    "plt.figure(dpi=200)\n",
    "plt.plot(binned_stats_pion.bin_edges[:-1], binned_stats_pion.statistic, label=r\"$\\pi^{\\pm}$\")\n",
    "plt.plot(binned_stats_pi0.bin_edges[:-1], binned_stats_pi0.statistic, label=r\"$\\pi^{0}$\")\n",
    "plt.xlabel(r\"True Pion $|\\eta|$\");\n",
    "plt.ylabel(\"Median Number of Clusters per Pion\");\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86e3d5c0-0b30-4ceb-80c6-8840254f7e8a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df = pd.DataFrame(a)\n",
    "\n",
    "#     ### Single-track, multi-cluster \n",
    "df = df[(df.nTrack == 1)]\n",
    "#     dR_pass = []\n",
    "#     deltaR_list = []\n",
    "#     for row in df.index:\n",
    "#         try:\n",
    "#             deltaR = np.sqrt((df['cluster_Eta'][row].astype('float') - df['trackEta'][row].astype('float'))**2 + \n",
    "#                                delta_phi(df['cluster_Phi'][row].astype('float'), df['trackPhi'][row].astype('float'))**2)\n",
    "#         except:\n",
    "#             deltaR = np.array(999)\n",
    "#         deltaR_list.append(deltaR)\n",
    "#         dR_pass.append(deltaR < 1.2)\n",
    "#     df[\"deltaR\"] = deltaR_list\n",
    "#     df[\"dR_pass\"] = dR_pass\n",
    "#     df[\"event_number\"] = df.index\n",
    "#     df.reset_index(inplace=True)\n",
    "#     indices_pass = [] \n",
    "#     for row in df.index: # kill all events with no clusters passing the delta R cut; deal with individual clusters at the training stage\n",
    "#         if len(df['trackPt'][row]) > 1:\n",
    "#             continue\n",
    "#         elif df.dR_pass[row].sum() > 0:\n",
    "#             indices_pass.append(row)\n",
    "#     df = df.iloc[indices_pass]\n",
    "#     df = df[(df.trackPt < 10**5)] # Track pT cut \n",
    "#     return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "352c0381-9d4f-41fe-b968-af50cb5dfdb4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for file in tqdm(pion_files[:10]):\n",
    "    a = uproot.open(file)[\"EventTree\"].arrays(variables, library = \"np\")\n",
    "    df = pd.DataFrame(a)\n",
    "    if len(df[(df.truthPartE < 0.3) & (df.nTrack == 1)]) != 0:\n",
    "        print(\"Match found!\")\n",
    "# print(df.truthPartE.min())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09a988be-4a20-4731-9012-8e62d78675ec",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df = pd.concat([pd.DataFrame(uproot.open(file)[\"EventTree\"].arrays(variables, library = \"np\")) \n",
    "                for file in tqdm(pion_files[400:])])\n",
    "df.reset_index(inplace=True)\n",
    "print(len(df[(df.truthPartE < 0.3) & (df.nTrack == 1)]))\n",
    "print(df[(df.truthPartE < 0.3) & (df.nTrack == 1)].truthPartE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07075a4c-d8bf-415c-90e9-7006251f8a4d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df[(df.truthPartE < 0.3) & (df.nTrack == 1)].truthPartE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06d2de4b-99e5-42d3-95d8-abab92e4f108",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df = pd.concat([pd.DataFrame(uproot.open(file)[\"EventTree\"].arrays(variables, library = \"np\")) \n",
    "                for file in tqdm(pion_files[300:400])])\n",
    "df.reset_index(inplace=True)\n",
    "print(len(df[(df.truthPartE < 0.3) & (df.nTrack == 1)]))\n",
    "print(df[(df.truthPartE < 0.3) & (df.nTrack == 1)].truthPartE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "230de6f1-0d41-4436-8903-b4c9071cc0f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.concat([pd.DataFrame(uproot.open(file)[\"EventTree\"].arrays(variables, library = \"np\")) \n",
    "                for file in tqdm(pion_files[:300])])\n",
    "df.reset_index(inplace=True)\n",
    "print(len(df[(df.truthPartE < 0.3) & (df.nTrack == 1)]))\n",
    "print(df[(df.truthPartE < 0.3) & (df.nTrack == 1)].truthPartE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b086bbe1-88b7-4475-8e3d-e6bbe5156863",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df[(df.truthPartE < 0.25)].nTrack.value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ceceb0b8-805b-4fc7-8555-cd33d3688969",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df[(df.truthPartE < 0.25) & (df.nTrack == 1)].truthPartE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa605dcb-8a10-4c85-bd7f-741d7de9f38c",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_test = pd.concat([pd.DataFrame(uproot.open(file)[\"EventTree\"].arrays(variables, library = \"np\")) \n",
    "                for file in tqdm(pion_files[400:])])\n",
    "df_test.reset_index(inplace=True)\n",
    "print(len(df_test[(df_test.truthPartE < 0.25) & (df_test.nTrack == 1)]))\n",
    "print(df_test[(df_test.truthPartE < 0.25) & (df_test.nTrack == 1)].truthPartE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bd900bc-ac50-4a93-a106-085cea962a9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_all = pd.concat([pd.DataFrame(uproot.open(file)[\"EventTree\"].arrays(variables, library = \"np\")) \n",
    "                for file in tqdm(pion_files)])\n",
    "df_all.reset_index(inplace=True)\n",
    "print(len(df_all[(df_all.truthPartE < 0.25) & (df_all.nTrack == 1)]))\n",
    "print(df_all[(df_all.truthPartE < 0.25) & (df_all.nTrack == 1)].truthPartE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "221b1367-e716-45e4-bf5d-97b7b3d66b4e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "len(df[(df.truthPartE < 0.3) & (df.nTrack == 1)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bfbc00e-61dc-4c55-8059-bcfbda6907c3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "100*6/len(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a0c230e-e384-40ff-b9ae-2acefc7c4429",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df[(df.truthPartE < 0.3) & (df.nTrack == 1)].truthPartE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8473a57c-3d33-4dc4-bf13-5401610537d1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_2 = pd.concat([pd.DataFrame(uproot.open(file)[\"EventTree\"].arrays(variables, library = \"np\")) \n",
    "                for file in tqdm(pion_files)])\n",
    "df_2.reset_index(inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56761eee-e4f9-441f-a5a3-cd335e13dcb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_2[(df_2.truthPartE < 0.3) & (df_2.nTrack == 1)].truthPartE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c8198e5-3357-4489-8909-a1e08536cd36",
   "metadata": {},
   "outputs": [],
   "source": [
    "a = uproot.open(pion_files[0])[\"EventTree\"].arrays(variables, library = \"np\")\n",
    "df = pd.DataFrame(a)\n",
    "df = df[(df.nTrack == 1)]\n",
    "print(df.truthPartE.min())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c8dd50d-a6b0-4721-ab62-15a6bbb6d16b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df.truthPartE.min()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "182f95f9-5b7b-4bc3-9bc7-e9e794c68ad3",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = apply_cuts(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df26211f-def5-46a6-a0eb-b50e52caa8aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_drs = []\n",
    "dfs = []\n",
    "for file in tqdm(pion_files[:10]):\n",
    "    prefix = file.split(\"/\")[:-2]\n",
    "    number = file.split(\"000\")[-1][:-5]\n",
    "    folder = os.path.join(\"/\".join(prefix), \"onetrack_multicluster\", \"pion_files\")\n",
    "    os.makedirs(folder, exist_ok=True)\n",
    "    npy_filename = os.path.join(folder, str(number)+\".npy\")\n",
    "    a = uproot.open(file)[\"EventTree\"].arrays(variables, library = \"np\")\n",
    "    df = apply_cuts(a)\n",
    "    list = df.deltaR.to_numpy().flatten()\n",
    "    drs = np.concatenate(list, axis=0)\n",
    "    all_drs.append(drs)\n",
    "    dfs.append(df)\n",
    "    \n",
    "import matplotlib.pyplot as plt\n",
    "all_drs = np.concatenate(all_drs, axis=0)\n",
    "\n",
    "plt.figure(dpi=200)\n",
    "plt.hist(all_drs, bins=np.linspace(0.5,1.5,30));\n",
    "plt.xlabel(r\"$\\Delta R$(cluster,track)\")\n",
    "\n",
    "plt.figure(dpi=200)\n",
    "plt.hist(all_drs, bins=np.linspace(0.5,1.5,30));\n",
    "plt.yscale(\"log\")\n",
    "plt.xlabel(r\"$\\Delta R$(cluster,track)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e013880f-82b5-4a9b-9eb7-fb188e63abf3",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.concat([df for df in dfs])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2629f9b9-1cc5-41d7-b9e8-58fb00534aa3",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(dpi=150)\n",
    "plt.hist(np.concatenate(df.trackPt.to_numpy().flatten(), axis=0), bins=10);\n",
    "plt.xlabel(r\"Track $p_T$\")\n",
    "plt.yscale(\"log\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a02f3399-8b5e-49ac-a103-7665d3bb5f36",
   "metadata": {},
   "source": [
    "### Scaling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9da55693-c71d-47c8-87b9-b6bd716213b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster_cell_e = []\n",
    "cluster_e = []\n",
    "cluster_eta = []\n",
    "cluster_phi = []\n",
    "track_pt = []\n",
    "track_z0 = []\n",
    "track_eta = []\n",
    "track_phi = []\n",
    "truth_part_e = []\n",
    "\n",
    "n_files = 10\n",
    "\n",
    "for file in tqdm(pion_files[:n_files]):\n",
    "    prefix = file.split(\"/\")[:-2]\n",
    "    number = file.split(\"000\")[-1][:-5]\n",
    "    a = uproot.open(file)[\"EventTree\"].arrays(variables, library = \"np\")\n",
    "    df = apply_cuts(a)\n",
    "    track_pt.append(df.trackPt.explode())\n",
    "    track_z0.append(df.trackZ0)\n",
    "    track_eta.append(df.trackEta)\n",
    "    track_phi.append(df.trackPhi)\n",
    "    truth_part_e.append(df.truthPartE.explode())\n",
    "    for i in range(len(df)): \n",
    "        for cluster in range(df.nCluster.iloc[i]):\n",
    "            cluster_cell_e.append(np.array(np.log10(df.cluster_cell_E.iloc[i][cluster])))\n",
    "            cluster_e.append(np.array(np.log10(df.cluster_E.iloc[i][cluster])))\n",
    "            cluster_eta.append(np.array(df.cluster_Eta.iloc[i][cluster]))\n",
    "            cluster_phi.append(np.array(df.cluster_Phi.iloc[i][cluster]))\n",
    "\n",
    "print(\"Track pT | mean: {} | std: {}\".format(np.mean([np.log10(x) for x in np.concatenate(track_pt)]), np.std([np.log10(x) for x in np.concatenate(track_pt)])))\n",
    "print(\"Track z0 | mean: {} | std: {}\".format(np.mean(np.concatenate(track_z0)), np.std(np.concatenate(track_z0))))\n",
    "print(\"Track eta | mean: {} | std: {}\".format(np.mean(np.concatenate(track_eta)), np.std(np.concatenate(track_eta))))\n",
    "print(\"Track phi | mean: {} | std: {}\".format(np.mean(np.concatenate(track_phi)), np.std(np.concatenate(track_phi))))\n",
    "\n",
    "print(\"Truth particle E | mean: {} | std: {}\".format(np.mean([np.log10(x) for x in np.concatenate(truth_part_e)]), np.std([np.log10(x) for x in np.concatenate(truth_part_e)])))\n",
    "\n",
    "print(\"Cluster cell e\")\n",
    "print(np.mean(np.concatenate(cluster_cell_e)))\n",
    "print(np.std(np.concatenate(cluster_cell_e)))\n",
    "        \n",
    "print(\"Cluster e\")\n",
    "print(np.mean(cluster_e))\n",
    "print(np.std(cluster_e))\n",
    "\n",
    "print(\"Cluster eta\")\n",
    "print(np.mean(cluster_eta))\n",
    "print(np.std(cluster_eta))\n",
    "\n",
    "print(\"Cluster phi\")\n",
    "print(np.mean(cluster_phi))\n",
    "print(np.std(cluster_phi))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56560bd1-6b80-479d-a2c1-0ee7a9de4253",
   "metadata": {},
   "outputs": [],
   "source": [
    "scales = {\n",
    "    'track_pt_mean': 1.633278727,\n",
    "    'track_pt_std': 0.8481947183,\n",
    "    'track_z0_mean': 0.08022017,\n",
    "    'track_z0_std': 42.53320004,    \n",
    "    'track_eta_mean': -0.00563187,\n",
    "    'track_eta_std': 1.35242735,    \n",
    "    'track_phi_mean': 0.00206431,\n",
    "    'track_phi_std': 1.81240248,   \n",
    "    'truth_part_e_mean': 1.92469358, \n",
    "    'truth_part_e_std': 0.8289864, \n",
    "    'cluster_cell_e_mean': -1.0121697,\n",
    "    'cluster_cell_e_std': 0.818357, \n",
    "    'cluster_e_mean': 0.89923394,\n",
    "    'cluster_e_std': 1.0585934,\n",
    "    'cluster_eta_mean': 0.016195267,\n",
    "    'cluster_eta_std': 1.3400925,\n",
    "    'cluster_phi_mean': 0.0050816955,\n",
    "    'cluster_phi_std': 1.8100655,\n",
    "         }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0edd91ff-d25b-45d7-b404-f2b4367aeac6",
   "metadata": {},
   "outputs": [],
   "source": [
    "cell_geo_file = \"/clusterfs/ml4hep/mpettee/ml4pions/data/cell_geo.root\"\n",
    "df = uproot.open(cell_geo_file)['CellGeo'].arrays(library='pd')\n",
    "df.reset_index(inplace=True)\n",
    "vars = ['cell_geo_sampling', 'cell_geo_eta', 'cell_geo_phi', 'cell_geo_rPerp', 'cell_geo_deta', 'cell_geo_dphi']\n",
    "for var in vars: \n",
    "    print(var)\n",
    "    print(np.mean(df[var]))\n",
    "    print(np.std(df[var]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "141f8437-9fa3-407b-9b5c-8db3081fcc1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "scales.update({\n",
    "    'cell_geo_sampling_mean': 3.8827391420197177,\n",
    "    'cell_geo_sampling_std': 3.9588233603576204,\n",
    "    'cell_geo_eta_mean': 0.0005979097,\n",
    "    'cell_geo_eta_std': 1.4709069,\n",
    "    'cell_geo_phi_mean': -2.8938382e-05,\n",
    "    'cell_geo_phi_std': 1.813651,\n",
    "    'cell_geo_rPerp_mean': 1478.9285,\n",
    "    'cell_geo_rPerp_std': 434.60815,\n",
    "    'cell_geo_deta_mean': 0.026611786,\n",
    "    'cell_geo_deta_std': 0.03396141,\n",
    "    'cell_geo_dphi_mean': 0.068693615,\n",
    "    'cell_geo_dphi_std': 0.038586758,  \n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a27cc559-eb45-47ee-a703-756189a68105",
   "metadata": {},
   "outputs": [],
   "source": [
    "scales"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd17841c-70c7-4427-b2c7-a5fac9841bf2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Node features: \n",
    "# {\n",
    "# np.log10(cluster_cell_E), \n",
    "# cell_geo_sampling, \n",
    "# cell_geo_eta, \n",
    "# cell_geo_phi, \n",
    "# cell_geo_rPerp, \n",
    "# cell_geo_deta, \n",
    "# cell_geo_dphi\n",
    "#}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "469b20f0-058b-4b00-be89-4bdc59333155",
   "metadata": {},
   "outputs": [],
   "source": [
    "node_means = [\n",
    "    scales[\"cluster_cell_e_mean\"], \n",
    "    scales[\"cell_geo_sampling_mean\"],\n",
    "    scales[\"cell_geo_eta_mean\"],\n",
    "    scales[\"cell_geo_phi_mean\"],\n",
    "    scales[\"cell_geo_rPerp_mean\"],\n",
    "    scales[\"cell_geo_deta_mean\"],\n",
    "    scales[\"cell_geo_dphi_mean\"],\n",
    "] \n",
    "\n",
    "node_stds = [\n",
    "    scales[\"cluster_cell_e_std\"], \n",
    "    scales[\"cell_geo_sampling_std\"],\n",
    "    scales[\"cell_geo_eta_std\"],\n",
    "    scales[\"cell_geo_phi_std\"],\n",
    "    scales[\"cell_geo_rPerp_std\"],\n",
    "    scales[\"cell_geo_deta_std\"],\n",
    "    scales[\"cell_geo_dphi_std\"],\n",
    "] "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfa0b5e1-0828-4863-ae5c-c37c59e18a77",
   "metadata": {},
   "outputs": [],
   "source": [
    "node_means"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7db48f4-41e8-4127-942d-29d42f6c94f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "node_stds"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nbdev",
   "language": "python",
   "name": "nbdev"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}