exatrkx_remove_noise.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import pickle\n",
    "import logging\n",
    "import argparse\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import functools\n",
    "\n",
    "import multiprocessing\n",
    "from multiprocessing.dummy import Pool as ThreadPool \n",
    "from multiprocessing import Pool as ProcessPool \n",
    "\n",
    "import trackml.dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def remove_some_noise(hits, cells, truth, keep=1.0):\n",
    "    if keep == 1.0:\n",
    "        return hits, cells, truth\n",
    "    unique_ids = truth.particle_id.unique()\n",
    "    track_ids_to_keep = unique_ids[np.where(unique_ids != 0)]\n",
    "\n",
    "    where_to_keep = truth['particle_id'].isin(track_ids_to_keep)\n",
    "    hit_idx  = hits[where_to_keep].hit_id.values\n",
    "    \n",
    "    noise_hit_idx = hits[~where_to_keep].hit_id.values\n",
    "    if keep <= 0.0:\n",
    "        noise_hit_idx = []\n",
    "    else:\n",
    "        noise_hit_idx = np.random.permutation(noise_hit_idx)[:int(keep * noise_hit_idx.shape[0])]\n",
    "    hit_idx = np.concatenate([hit_idx, noise_hit_idx])\n",
    "        \n",
    "    hits_reduced  = hits[hits['hit_id'].isin(hit_idx)]\n",
    "    truth_reduced = truth[truth['hit_id'].isin(hit_idx)]\n",
    "    cells_reduced = cells[cells['hit_id'].isin(hit_idx)]\n",
    "    return hits_reduced, cells_reduced, truth_reduced"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#dir_path = \"/global/cfs/cdirs/m3443/data/trackml-kaggle/train_all\"\n",
    "dir_path = \"/global/homes/a/aoka/exatrkx-ctd2020/trackml/train_100_events\"\n",
    "event_path = os.path.join(dir_path, \"event000001005\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hits, cells, particles, truth = trackml.dataset.load_event(event_path)\n",
    "print(hits.shape, cells.shape, truth.shape)\n",
    "hits, cells, truth = remove_some_noise(hits, cells, truth, 0.5)\n",
    "print(hits.shape, cells.shape, truth.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.  0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1. ]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x2aaadfbd0050>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAZEAAAEGCAYAAACkQqisAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy86wFpkAAAACXBIWXMAAAsTAAALEwEAmpwYAAAa+klEQVR4nO3de7xWVb3v8c9PUJblNaFOirWw8IJKXvCWZZZuvMQOU/OlJ/OyKfPSztq5E6uTtzyndr3ULMvt9oKaecPtEdNS83K8JCImXlATUkzMklBQS0rqd/54BviAC3iYa81nseDzfr2e15pzzDHnGHMJfplzzGfMyEwkSapitd7ugCSp7zJEJEmVGSKSpMoMEUlSZYaIJKmy/r3dgXYbOHBgdnZ29nY3JKnPePDBB/+UmYO62rbKhUhnZyeTJ0/u7W5IUp8REc8uaZu3syRJlRkikqTKDBFJUmWr3JiIpJXTG2+8wcyZM5k3b15vd6XP6ujoYPDgway++uot72OISFopzJw5k7XXXpvOzk4iore70+dkJrNnz2bmzJkMGTKk5f28nSVppTBv3jw22GADA6SiiGCDDTZY7is5Q0TSSsMA6Z4qvz9DRJJUmWMiklZKnWNv7NHjzfj2x3v0eADf/OY32W233dhzzz17/NjtYojU4ZR1W6w3t2+21xttrqjn6O90xW+z3X7/UGv1NtyW0047ra3t1cHbWZLUQ2bMmMEWH9mfz/376Wz50QMZecixvP76PKY89ht2HnUYw/c8iE+O+Qovz3kFgCOOOILx48cDMHbsWIYNG8bw4cM54YQTAJg1axYHHHAAO+ywAzvssAP33ntvr53bkhgiktSDpj3zHMcdfhBT7xjPeuuszbU33cZhX/pffOfrx/PIL69m683fz6ln/uci+8yePZvrrruOqVOn8sgjj/CNb3wDgOOPP54vf/nLPPDAA1x77bV89rOf7Y1TWipvZ0lSDxqy8YZss9VmAGw/fAt+++xM5sx9jY/ssj0Ah39qFJ/6/ImL7LPuuuvS0dHBmDFjGDVqFKNGjQLgl7/8JY8//vjCeq+88gqvvfYaa621VpvOZtkMEUnqQQMGrLFwuV+/1Zgz99Vl7tO/f38mTZrEbbfdxvjx4/nhD3/I7bffzj/+8Q8mTpxIR0dHnV3uFm9nSVKN1l1nLdZfd23uvv/XAFx27Y18ZOftFqnz2muvMXfuXPbdd1/OOussHn74YQBGjhzJD37wg4X1pkyZ0rZ+t8orEUkrpRnf/nivP7m0wCVnn8bRY8/gL/Pmscl7BnPxmacssv3VV19l9OjRzJs3j8zkzDPPBOCcc87huOOOY/jw4cyfP5/ddtuN8847r9a+Li9DRJJ6SGdnJ4/dfs3C9ROOPmzh8sSfXfqW+uPGjVu4PGnSpLdsHzhwIFdddVXPdrKHeTtLklSZISJJqswQkSRVZohIkiozRCRJlRkikqTKfMRX0sqp1RmDWz5ez88sPG7cOEaOHMmGG27Y7WNNfvhxLh3/M845/as90LPWGSKS1EvGjRvHVltt1SMhMuIDwxjxgWE90Kvl4+0sSeoBM2bMYIsttnjLNPBAl1PBjx8/nsmTJ/PpT3+abbbZhtdff32R4+2+++6ceOKJ7Ljjjmy66abcfffdQONd8kceeSRbb7012267LXfc+wAAd/5qMqMO+yIA/+++B9nmnw5mm386mG1HHsKrr/0ZgO9+97vssMMODB8+nJNPPrlHztsQkaQeMm3atLdMAw90ORX8gQceyIgRI7j88suZMmUKa6655luON3/+fCZNmsTZZ5/NqaeeCsC5555LRPDoo49yxRVXcPiXTmbevL8ust/3zruUc//3WKbceiV3X3cha3YM4JZbbmHatGlMmjSJKVOm8OCDD3LXXXd1+5wNEUnqIUOGDFlkGvgZz73A3FdefctU8Hfd39qcXvvvv3/jWNtvz4wZMwC45557OPTQQwHYfPPNee/g/8FTTz+7yH677rAN/3bqmZxz4RXMmfsq/fv355ZbbuGWW25h2223ZbvttuPJJ59k2rRp3T5nx0QkqYcMGDBg4XK/fqvx+rz5PXK8fv36MX9+68ca+4Uj+fgeH+Km2+9l1/2O5OafnktmctJJJ/H5z3++W31anFciklSjdddZe4lTwa+99tq8+uqy3zfS7MMf/jCXX345AE899RS/e/4PbPa+zkXq/HbGc2y9xVBOPO4IdvjAljw5fQZ77bUXF110Ea+99hoAzz//PC+++GI3z84rEUkrq1PmrvBTwR9xxBEcffTRrLnmmtx3331djoss7thjj+WYY45h6623pn///ow769RFXoQFcPYFP+WOX01mtdWCLTd9H/t8dFcGDNmJJ554gl122QWAtdZai5/85Ce8853v7Na5GSKS1AM6Ozt57LHHFgZX8zTw22y1WZdTwR9wwAEccMABXR7vzjvvXLg8cODAhWMiHR0dXHzxxW9WLO3t/sER7P7BEQD84FuLvn53geOPP57jjz++5XNqhbezJEmV1RoiEfHliJgaEY9FxBUR0RERQyLi/oiYHhFXRcQape6Asj69bO9sOs5Jpfw3EbFXU/nepWx6RIyt81wkSW9VW4hExEbAF4ERmbkV0A84GPgOcFZmvh94GRhTdhkDvFzKzyr1iIhhZb8tgb2BH0VEv4joB5wL7AMMAw4pdSWtojKzt7vQp1X5/dV9O6s/sGZE9AfeBrwAfAwYX7ZfAuxXlkeXdcr2PSIiSvmVmfnXzHwGmA7sWD7TM/PpzPwbcGWpK2kV1NHRwezZsw2SijKT2bNn09HRsVz71TawnpnPR8T3gN8BrwO3AA8CczJzwQPPM4GNyvJGwHNl3/kRMRfYoJRPbDp08z7PLVa+U1d9iYijgKMA3vOe93TvxCStkAYPHszMmTOZNWvWm4VzWnyEde4TPdeRdrfZg+11dHQwePDg5Wq+thCJiPVpXBkMAeYA19C4HdV2mXk+cD7AiBEj/GeKtBJaffXVGTJkyKKFp+zc2s49OUNvu9vsjXNsUuftrD2BZzJzVma+Afw3sCuwXrm9BTAYeL4sPw9sDFC2rwvMbi5fbJ8llUuS2qTOEPkdsHNEvK2MbewBPA7cARxY6hwOXF+WJ5R1yvbbs3FzcwJwcHl6awgwFJgEPAAMLU97rUFj8H1CjecjSVpMnWMi90fEeODXwHzgIRq3lG4EroyIb5WyC8suFwKXRcR04CUaoUBmTo2Iq2kE0HzguMz8O0BEfAG4mcaTXxdl5tS6zkeS9Fa1fmM9M08GFp+0/mkaT1YtXnce8KklHOcM4Iwuym8Cbup+TyVJVfiNdUlSZYaIJKkyQ0SSVJkhIkmqzBCRJFVmiEiSKjNEJEmVGSKSpMoMEUlSZYaIJKkyQ0SSVJkhIkmqzBCRJFVmiEiSKjNEJEmVGSKSpMoMEUlSZYaIJKkyQ0SSVJkhIkmqzBCRJFVmiEiSKjNEJEmVGSKSpMoMEUlSZYaIJKkyQ0SSVJkhIkmqzBCRJFVmiEiSKjNEJEmVGSKSpMoMEUlSZYaIJKkyQ0SSVJkhIkmqrNYQiYj1ImJ8RDwZEU9ExC4R8Y6IuDUippWf65e6ERHnRMT0iHgkIrZrOs7hpf60iDi8qXz7iHi07HNORESd5yNJWlTdVyLfB36RmZsDHwCeAMYCt2XmUOC2sg6wDzC0fI4CfgwQEe8ATgZ2AnYETl4QPKXO55r227vm85EkNaktRCJiXWA34EKAzPxbZs4BRgOXlGqXAPuV5dHApdkwEVgvIt4N7AXcmpkvZebLwK3A3mXbOpk5MTMTuLTpWJKkNqjzSmQIMAu4OCIeiogLIuLtwLsy84VS5w/Au8ryRsBzTfvPLGVLK5/ZRflbRMRRETE5IibPmjWrm6clSVqgzhDpD2wH/DgztwX+zJu3rgAoVxBZYx8WtHN+Zo7IzBGDBg2quzlJWmXUGSIzgZmZeX9ZH08jVP5YbkVRfr5Ytj8PbNy0/+BStrTywV2US5LapLYQycw/AM9FxGalaA/gcWACsOAJq8OB68vyBOCw8pTWzsDcctvrZmBkRKxfBtRHAjeXba9ExM7lqazDmo4lSWqD/jUf/1+ByyNiDeBp4EgawXV1RIwBngUOKnVvAvYFpgN/KXXJzJci4nTggVLvtMx8qSwfC4wD1gR+Xj6SpDapNUQycwowootNe3RRN4HjlnCci4CLuiifDGzVvV5KkqparttZEbFaRKxTV2ckSX3LMkMkIn4aEeuUx3MfAx6PiH+vv2uSpBVdK1ciwzLzFRpf5Ps5je9/fKbOTkmS+oZWQmT1iFidRohMyMw36u2SJKmvaCVE/hOYAbwduCsi3gvMrbNTkqS+oZUQuSEzN8rMfcsTVL8D/qXmfkmS+oBWQuTa5pUSJFfW0x1JUl+yxO+JRMTmwJbAuhGxf9OmdYCOujsmSVrxLe3LhpsBo4D1gH9uKn+Vxjs8JEmruCWGSGZeD1wfEbtk5n1t7JMkqY9Y2u2sr2bmfwD/MyIOWXx7Zn6x1p5JklZ4S7ud9UT5ObkdHZEk9T1Lu511Q/l5yZLqSJJWbcucxTciNgVOADqb62fmx+rrliSpL2hlKvhrgPOAC4C/19sdSVJf0kqIzM/MH9feE0lSn7O0p7PeURZviIhjgeuAvy7Y3vR2QUnSKmppVyIPAglEWW9+h0gCm9TVKUlS37C0p7OGtLMjkqS+Z7lejytJUjNDRJJU2RJDJCJ2LT8HtK87kqS+ZGlXIueUn06+KEnq0tKeznojIs4HNoqIcxbf6ASMkqSlhcgoYE9gLxqP+0qStIilPeL7J+DKiHgiMx9uY58kSX1EK09nzY6I6yLixfK5NiIG194zSdIKr5UQuRiYAGxYPjeUMknSKq6VEHlnZl6cmfPLZxwwqOZ+SZL6gFZC5E8RcWhE9CufQ4HZdXdMkrTiayVE/gU4CPgD8AJwIHBknZ2SJPUNy3yfSGY+C3yiDX2RJPUxzp0lSarMEJEkVWaISJIqW2aIRMQ3mpad0VeStNDSpoI/MSJ2ofE01gLLPaNveSz4oYj4WVkfEhH3R8T0iLgqItYo5QPK+vSyvbPpGCeV8t9ExF5N5XuXsukRMXZ5+yZJ6p6lXYk8CXwK2CQi7o6I/wI2iIjNlrON44Enmta/A5yVme8HXgbGlPIxwMul/KxSj4gYBhwMbAnsDfxowXdWgHOBfYBhwCGlriSpTZYWInOArwHTgd2B75fysRHxq1YOXubY+jhwQVkP4GPA+FLlEmC/sjy6rFO271HqjwauzMy/ZuYzpT87ls/0zHw6M/8GXFnqSpLaZGkhshdwI/A+4ExgJ+DPmXlkZn6wxeOfDXwV+EdZ3wCYk5nzy/pMYKOyvBHwHEDZPrfUX1i+2D5LKn+LiDgqIiZHxORZs2a12HVJ0rIsMUQy82uZuQcwA7gM6AcMioh7IuKGZR04IkYBL2Zmr7+LJDPPz8wRmTli0CCn/ZKknrLMb6wDN2fmZGByRByTmR+KiIEt7Lcr8ImI2BfoANahcUtsvYjoX642BgPPl/rPAxsDMyOiP7AujTm6FpQv0LzPksolSW2wzEd8M/OrTatHlLI/tbDfSZk5ODM7aQyM356Znwbu4M0nvg4Hri/LE8o6ZfvtmZml/ODy9NYQYCgwCXgAGFqe9lqjtDFhWf2SJPWcVq5EFuqhNxyeSOONid8CHgIuLOUXApdFxHTgJRqhQGZOjYirgceB+cBxmfl3gIj4AnAzjVttF2Xm1B7onySpRcsVIlVl5p3AnWX5aRpPVi1eZx6NR4q72v8M4Iwuym8CburBrkqSloPTnkiSKjNEJEmVGSKSpMoMEUlSZYaIJKkyQ0SSVJkhIkmqzBCRJFVmiEiSKjNEJEmVGSKSpMoMEUlSZYaIJKkyQ0SSVJkhIkmqzBCRJFVmiEiSKjNEJEmVGSKSpMoMEUlSZYaIJKkyQ0SSVJkhIkmqzBCRJFVmiEiSKjNEJEmVGSKSpMoMEUlSZYaIJKkyQ0SSVJkhIkmqzBCRJFVmiEiSKjNEJEmVGSKSpMoMEUlSZbWFSERsHBF3RMTjETE1Io4v5e+IiFsjYlr5uX4pj4g4JyKmR8QjEbFd07EOL/WnRcThTeXbR8SjZZ9zIiLqOh9J0lvVeSUyH/hKZg4DdgaOi4hhwFjgtswcCtxW1gH2AYaWz1HAj6EROsDJwE7AjsDJC4Kn1Plc035713g+kqTF1BYimflCZv66LL8KPAFsBIwGLinVLgH2K8ujgUuzYSKwXkS8G9gLuDUzX8rMl4Fbgb3LtnUyc2JmJnBp07EkSW3QljGRiOgEtgXuB96VmS+UTX8A3lWWNwKea9ptZilbWvnMLsq7av+oiJgcEZNnzZrVvZORJC1Ue4hExFrAtcCXMvOV5m3lCiLr7kNmnp+ZIzJzxKBBg+puTpJWGbWGSESsTiNALs/M/y7Ffyy3oig/XyzlzwMbN+0+uJQtrXxwF+WSpDap8+msAC4EnsjMM5s2TQAWPGF1OHB9U/lh5SmtnYG55bbXzcDIiFi/DKiPBG4u216JiJ1LW4c1HUuS1Ab9azz2rsBngEcjYkop+xrwbeDqiBgDPAscVLbdBOwLTAf+AhwJkJkvRcTpwAOl3mmZ+VJZPhYYB6wJ/Lx8JEltUluIZOY9wJK+t7FHF/UTOG4Jx7oIuKiL8snAVt3opiStkDrH3thSvRkdNXdkGfzGuiSpMkNEklSZISJJqqzOgXVJWmn0lTGKdvNKRJJUmSEiSarMEJEkVeaYiKQ+x/GJFYdXIpKkygwRSVJlhogkqTLHRCR1m2MUqy6vRCRJlRkikqTKDBFJUmWOiUgrGccn1E5eiUiSKjNEJEmVGSKSpMocE5Fq5hiFVmZeiUiSKjNEJEmVGSKSpMocE9EqxzEKqed4JSJJqswQkSRVZohIkipzTES9yvEJqW/zSkSSVJkhIkmqzBCRJFXmmIgW4RiFpOXhlYgkqTJDRJJUmSEiSarMMZEVmOMTklZ0ff5KJCL2jojfRMT0iBjb2/2RpFVJnw6RiOgHnAvsAwwDDomIYb3bK0ladfTpEAF2BKZn5tOZ+TfgSmB0L/dJklYZkZm93YfKIuJAYO/M/GxZ/wywU2Z+YbF6RwFHldXNgN/U3LWBwJ9qbqM32+uNNj3Hvt9eb7TpOfaM92bmoK42rBID65l5PnB+u9qLiMmZOWJlba832vQc+357vdGm51i/vn4763lg46b1waVMktQGfT1EHgCGRsSQiFgDOBiY0Mt9kqRVRp++nZWZ8yPiC8DNQD/gosyc2svdgjbeOuul9nqjTc+x77fXG216jjXr0wPrkqTe1ddvZ0mSepEhIkmqzBCpaFnTrUTEgIi4qmy/PyI629DmbhHx64iYX75DU3d7/xYRj0fEIxFxW0S8tw1tHh0Rj0bElIi4p7szFLQ6bU5EHBARGRHdfpSyhXM8IiJmlXOcEhGfrbO9Uueg8t9yakT8tDvttdJmRJzVdH5PRcScmtt7T0TcEREPlT+v+9bc3nvL34lHIuLOiBjczfYuiogXI+KxJWyPiDin9OeRiNiuO+0tl8z0s5wfGoP4vwU2AdYAHgaGLVbnWOC8snwwcFUb2uwEhgOXAge2ob2PAm8ry8e06RzXaVr+BPCLOtsr9dYG7gImAiPacI5HAD9s45/VocBDwPpl/Z11t7lY/X+l8VBMned4PnBMWR4GzKi5vWuAw8vyx4DLuvk73Q3YDnhsCdv3BX4OBLAzcH9P/Plp5eOVSDWtTLcyGrikLI8H9oiIqLPNzJyRmY8A/+hGO8vT3h2Z+ZeyOpHG93TqbvOVptW3A915MqTVaXNOB74DzOtGW8vbZk9ppb3PAedm5ssAmfliG9psdghwRc3tJbBOWV4X+H3N7Q0Dbi/Ld3Sxfblk5l3AS0upMhq4NBsmAutFxLu702arDJFqNgKea1qfWcq6rJOZ84G5wAY1t9mTlre9MTT+JVR7mxFxXET8FvgP4It1tlduC2ycma3Ny98DbRYHlNsS4yNi4y6292R7mwKbRsS9ETExIvbuRnuttgk0bvsAQ3jzf7h1tXcKcGhEzARuonH1U2d7DwP7l+VPAmtHRHf+/vdEn2phiKjbIuJQYATw3Xa0l5nnZub7gBOBb9TVTkSsBpwJfKWuNpbgBqAzM4cDt/LmFW1d+tO4pbU7jauC/4qI9Wpuc4GDgfGZ+fea2zkEGJeZg2nc+rms/PetywnARyLiIeAjNGbSqPsce4UhUk0r060srBMR/WlcQs+uuc2e1FJ7EbEn8HXgE5n513a02eRKYL8a21sb2Aq4MyJm0LjXPKGbg+vLPMfMnN30u7wA2L7O9mj8q3VCZr6Rmc8AT9EIlTrbXOBguncrq9X2xgBXA2TmfUAHjYkLa2kvM3+fmftn5rY0/n6QmXMqttcjfapNuwZfVqYPjX+5PU3jMnzBwNqWi9U5jkUH1q+uu82muuPo/sB6K+e4LY0BxqFt/L0ObVr+Z2ByO36npf6ddH9gvZVzfHfT8ieBiTW3tzdwSVkeSOO2yAZ1/16BzYEZlC8913yOPweOKMtb0BgTqdRui+0NBFYry2cAp3XnHMtxOlnywPrHWXRgfVJ322u5X+1qaGX70Lgkfqr8T/Trpew0Gv8ih8a/dK4BpgOTgE3a0OYONP5V+WcaVz1Ta27vl8AfgSnlM6EN5/h9YGpp746u/ufUk+0tVvdOuhkiLZ7j/ynn+HA5x81rbi9o3LZ7HHgUOLjucyzrpwDf7m5bLZ7jMODe8judAoysub0DgWmlzgXAgG62dwXwAvBG+Ts+BjgaOLrpv+G5pT+P9sSf01Y/TnsiSarMMRFJUmWGiCSpMkNEklSZISJJqswQkSRVZohINYqIziXNvCqtDAwRSVJlhojUJhGxSXmfxU4R8YuIeDAi7o6Izcv2QRFxbUQ8UD67lvJTIuKyiLgvIqZFxOd690ykN/Xv7Q5Iq4KI2IzGXF9H0Ph2+NGZOS0idgJ+ROOdE98HzsrMeyLiPcDNNKbogMZ7YnamMf39QxFxY2Z2ZzpzqUcYIlL9BgHX05ga/HfAB4Frml4vM6D83BMY1lS+TkSsVZavz8zXgdcj4g4a77T4v/V3XVo6Q0Sq31wa4fEhGlcjczJzmy7qrQbsnJmLvPyqhMri8xM5X5FWCI6JSPX7G43ZeA8DRgHPRMSnYOG7sT9Q6t1C08uSImKbpmOMjoiO8mKj3YEH2tBvaZkMEakNMvPPNALky8BVwJiIeJjGbL0LXp36RWBEeaPh4zRmaV3gERoz+k4ETnc8RCsKZ/GVVnARcQrwWmZ+r7f7Ii3OKxFJUmVeiUiSKvNKRJJUmSEiSarMEJEkVWaISJIqM0QkSZX9f/Zo20T17+GfAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "hits, cells, particles, truth = trackml.dataset.load_event(event_path)\n",
    "pt = np.sqrt(particles.px**2 + particles.py**2 + particles.pz**2)\n",
    "particles = particles.assign(pt=pt)\n",
    "keeps = np.arange(0.0, 1.1, 0.1)\n",
    "print(keeps)\n",
    "\n",
    "noise_hits = []\n",
    "not_noise_hits = []\n",
    "\n",
    "for i in keeps:\n",
    "    _, _, reduced_truth = remove_some_noise(hits, cells, truth, i)\n",
    "    n_noise = sum(reduced_truth['particle_id'] == 0)\n",
    "    noise_hits.append(n_noise)\n",
    "    not_noise_hits.append(len(reduced_truth) - n_noise)\n",
    "\n",
    "x = np.arange(len(keeps))  # the label locations\n",
    "width = 0.35  # the width of the bars\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "rects1 = ax.bar(x - width/2, noise_hits, width, label='noise')\n",
    "rects2 = ax.bar(x + width/2, not_noise_hits, width, label='not noise')\n",
    "\n",
    "# Add some text for labels, title and custom x-axis tick labels, etc.\n",
    "ax.set_ylabel('# of hits')\n",
    "ax.set_xlabel('keep')\n",
    "ax.set_xticks(x)\n",
    "labels = ['%.1f' % (i,) for i in keeps]\n",
    "ax.set_xticklabels(labels)\n",
    "ax.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from preprocess import get_one_event, load_detector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading detector...\n",
      "Detector loaded.\n"
     ]
    }
   ],
   "source": [
    "detector_path = \"/global/homes/a/aoka/exatrkx-ctd2020/trackml/detectors.csv\"\n",
    "detector_orig, detector_proc = load_detector(detector_path)\n",
    "event = get_one_event(event_path,detector_orig,detector_proc,True, 0.5, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[       hit_id           x          y       z  volume_id  layer_id  module_id  \\\n",
       " 0           1  -68.878403   2.077260 -1502.5          7         2          1   \n",
       " 1           2  -91.987801 -10.524700 -1502.5          7         2          1   \n",
       " 2           3  -54.743698  -6.891380 -1502.5          7         2          1   \n",
       " 3           4  -56.549900  -6.883070 -1502.5          7         2          1   \n",
       " 4           5  -55.004700 -11.325500 -1502.5          7         2          1   \n",
       " ...       ...         ...        ...     ...        ...       ...        ...   \n",
       " 92213  110019 -752.534973  65.378098  2944.5         18        12         97   \n",
       " 92214  110020 -783.512024  69.813599  2944.5         18        12         97   \n",
       " 92215  110021 -777.473999   2.915470  2944.5         18        12         97   \n",
       " 92216  110022 -974.481995  10.208400  2952.5         18        12         98   \n",
       " 92217  110023 -961.247009  54.945599  2952.5         18        12         98   \n",
       " \n",
       "        cell_count  cell_val      leta      lphi    lx        ly   lz  \\\n",
       " 0             1.0  0.257111  2.091356  0.844154  0.05   0.05625  0.3   \n",
       " 1             2.0  0.273439  1.623512  1.152572  0.05   0.11250  0.3   \n",
       " 2             1.0  0.284049  2.091356  0.844154  0.05   0.05625  0.3   \n",
       " 3             1.0  0.301796  2.091356  0.844154  0.05   0.05625  0.3   \n",
       " 4             2.0  0.358431  1.623512  1.152572  0.05   0.11250  0.3   \n",
       " ...           ...       ...       ...       ...   ...       ...  ...   \n",
       " 92213         5.0  5.000000  0.067145  1.513168  0.60  10.40000  0.7   \n",
       " 92214         1.0  1.000000  0.067253  1.559258  0.12  10.40000  0.7   \n",
       " 92215         1.0  1.000000  0.067253  1.559258  0.12  10.40000  0.7   \n",
       " 92216         2.0  2.000000  0.067239  1.547723  0.24  10.40000  0.7   \n",
       " 92217         2.0  2.000000  0.067239  1.547723  0.24  10.40000  0.7   \n",
       " \n",
       "            geta      gphi  \n",
       " 0     -2.091356 -2.336410  \n",
       " 1     -1.623512 -2.644828  \n",
       " 2     -2.091356 -2.336410  \n",
       " 3     -2.091356 -2.336410  \n",
       " 4     -1.623512 -2.644828  \n",
       " ...         ...       ...  \n",
       " 92213  0.067145  3.018514  \n",
       " 92214  0.067252  3.064605  \n",
       " 92215  0.067252  3.064605  \n",
       " 92216  0.067239  3.055688  \n",
       " 92217  0.067239  3.055688  \n",
       " \n",
       " [92218 rows x 16 columns],\n",
       "         hit_id         particle_id          tx         ty      tz       tpx  \\\n",
       " hit_id                                                                        \n",
       " 0            1  409833544685191168  -68.863098   2.085910 -1502.5 -0.222967   \n",
       " 1            2  450363536149839872  -92.002296 -10.515500 -1502.5 -0.498183   \n",
       " 2            3  256711500951977984  -54.723099  -6.892780 -1502.5 -0.242608   \n",
       " 3            4  256711432232501248  -56.564301  -6.884460 -1502.5 -2.295520   \n",
       " 4            5   45036339871088640  -54.989498 -11.338000 -1502.5 -0.414272   \n",
       " ...        ...                 ...         ...        ...     ...       ...   \n",
       " 110018  110019   90076528049659906 -756.646973  65.619202  2944.5 -0.013601   \n",
       " 110019  110020  373800555778146304 -780.107971  69.570503  2944.5 -0.332783   \n",
       " 110020  110021  517933061162139648 -780.028015   3.119130  2944.5 -0.520974   \n",
       " 110021  110022   36049412861984768 -971.122009   9.990380  2952.5 -0.930866   \n",
       " 110022  110023   40556379743715328 -963.507996  55.077202  2952.5 -0.425574   \n",
       " \n",
       "              tpy        tpz    weight         pt  \n",
       " hit_id                                            \n",
       " 0       0.027029  -4.871460  0.000007   4.880675  \n",
       " 1      -0.083997  -8.219370  0.000009   8.240234  \n",
       " 2      -0.044478  -6.729290  0.000009   6.737144  \n",
       " 3      -0.264053 -60.863300  0.000031  60.910053  \n",
       " 4      -0.066555 -11.277500  0.000008  11.289529  \n",
       " ...          ...        ...       ...        ...  \n",
       " 110018  0.150778   0.208964  0.000000   0.259504  \n",
       " 110019 -0.068478   1.411590  0.000010   1.462401  \n",
       " 110020 -0.090525   2.133090  0.000008   2.208025  \n",
       " 110021  0.123969   2.767040  0.000019   2.935543  \n",
       " 110022 -0.084756   1.405110  0.000010   1.479655  \n",
       " \n",
       " [92218 rows x 10 columns]]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "event"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:exatrkx]",
   "language": "python",
   "name": "conda-env-exatrkx-py"
  },
  "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}