nn.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/anaconda3/lib/python3.6/site-packages/h5py/__init__.py:36: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.\n",
      "  from ._conv import register_converters as _register_converters\n",
      "Using TensorFlow backend.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['train_3.zip', 'train_2.zip', 'detectors.zip', '.DS_Store', 'train_4.zip', 'test', 'train_1', 'code', 'trackml-participant-document-particle-v1.0.pdf', 'Trackml_accuracy_submissions_public', 'train_sample.zip', 'test.zip', '.ipynb_checkpoints', 'train_100_events', 'sample_submission.csv.zip', 'nohup.out', 'sample_submission.csv']\n",
      "['train_3.zip', 'train_2.zip', 'detectors.zip', '.DS_Store', 'train_4.zip', 'test', 'train_1', 'code', 'trackml-participant-document-particle-v1.0.pdf', 'Trackml_accuracy_submissions_public', 'train_sample.zip', 'test.zip', '.ipynb_checkpoints', 'train_100_events', 'sample_submission.csv.zip', 'nohup.out', 'sample_submission.csv']\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import os\n",
    "from keras.models import Model, Sequential, load_model\n",
    "from keras.layers import Dense, Input\n",
    "from keras.optimizers import Adam\n",
    "from tqdm import tqdm_notebook\n",
    "print(os.listdir(\"../\"))\n",
    "print(os.listdir(\"../\"))\n",
    "prefix='../'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "_uuid": "6cfcc4d7929c3744e906026c56c9f8d03a0a46df"
   },
   "outputs": [],
   "source": [
    "def init_model(fs = 10):\n",
    "    model = Sequential()\n",
    "    model.add(Dense(400, activation='selu', input_shape=(fs,)))\n",
    "    model.add(Dense(200, activation='selu'))\n",
    "    model.add(Dense(200, activation='selu'))\n",
    "    model.add(Dense(100, activation='selu'))\n",
    "    model.add(Dense(100, activation='selu'))\n",
    "    model.add(Dense(1, activation='sigmoid'))\n",
    "    return model\n",
    "\n",
    "def get_event(event):\n",
    "    hits= pd.read_csv(prefix+'train_1/%s-hits.csv'%event)\n",
    "    cells= pd.read_csv(prefix+'train_1/%s-cells.csv'%event)\n",
    "    truth= pd.read_csv(prefix+'train_1/%s-truth.csv'%event)\n",
    "    particles= pd.read_csv(prefix+'train_1/%s-particles.csv'%event)\n",
    "    return hits, cells, truth, particles"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_uuid": "22ca78487ca459c532750dd133c7ed401e592e3d"
   },
   "source": [
    "# Step 1 - Prepare training data\n",
    "* use 10 events for training\n",
    "* input: hit pair\n",
    "* output: 1 if two hits are the same particle_id, 0 otherwise.\n",
    "* feature size: 10 (5 per hit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "_uuid": "8a8aa70428d754645a5ab2c74f7bc6c3c99505b7"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d2be77216ec04db3bf1630aa8eecfd8b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=10), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "event000001010 (905342, 11)\n",
      "event000001011 (1049940, 11)\n",
      "event000001012 (975962, 11)\n",
      "event000001013 (937140, 11)\n",
      "event000001014 (1138386, 11)\n",
      "event000001015 (1096946, 11)\n",
      "event000001016 (1054472, 11)\n",
      "event000001017 (1125976, 11)\n",
      "event000001018 (787588, 11)\n",
      "event000001019 (1099946, 11)\n",
      "\n",
      "(40684242, 11)\n"
     ]
    }
   ],
   "source": [
    "# you can jump to step4 for test only.\n",
    "train = True\n",
    "if train:\n",
    "    Train = []\n",
    "    for i in tqdm_notebook(range(10,20)):\n",
    "        event = 'event0000010%02d'%i\n",
    "        hits, cells, truth, particles = get_event(event)\n",
    "        hit_cells = cells.groupby(['hit_id']).value.count().values\n",
    "        hit_value = cells.groupby(['hit_id']).value.sum().values\n",
    "        features = np.hstack((hits[['x','y','z']]/1000, hit_cells.reshape(len(hit_cells),1)/10,hit_value.reshape(len(hit_cells),1)))\n",
    "        particle_ids = truth.particle_id.unique()\n",
    "        particle_ids = particle_ids[np.where(particle_ids!=0)[0]]\n",
    "\n",
    "        pair = []\n",
    "        for particle_id in particle_ids:\n",
    "            hit_ids = truth[truth.particle_id == particle_id].hit_id.values-1\n",
    "            for i in hit_ids:\n",
    "                for j in hit_ids:\n",
    "                    if i != j:\n",
    "                        pair.append([i,j])\n",
    "        pair = np.array(pair)   \n",
    "        Train1 = np.hstack((features[pair[:,0]], features[pair[:,1]], np.ones((len(pair),1))))\n",
    "\n",
    "        if len(Train) == 0:\n",
    "            Train = Train1\n",
    "        else:\n",
    "            Train = np.vstack((Train,Train1))\n",
    "\n",
    "        n = len(hits)\n",
    "        size = len(Train1)*3\n",
    "        p_id = truth.particle_id.values\n",
    "        i =np.random.randint(n, size=size)\n",
    "        j =np.random.randint(n, size=size)\n",
    "        pair = np.hstack((i.reshape(size,1),j.reshape(size,1)))\n",
    "        pair = pair[((p_id[i]==0) | (p_id[i]!=p_id[j]))]\n",
    "\n",
    "        Train0 = np.hstack((features[pair[:,0]], features[pair[:,1]], np.zeros((len(pair),1))))\n",
    "\n",
    "        print(event, Train1.shape)\n",
    "\n",
    "        Train = np.vstack((Train,Train0))\n",
    "    del Train0, Train1\n",
    "\n",
    "    np.random.shuffle(Train)\n",
    "    print(Train.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_uuid": "9f38fbd32829f9ca2a18f8b00bf3db26f57f4297"
   },
   "source": [
    "# Step 2 - Train model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "_uuid": "3bab2db55c8d0abf38119e4f6780269adaaf906e"
   },
   "outputs": [],
   "source": [
    "if train:\n",
    "    model = init_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "_uuid": "9e96d105f22f0fb6d605332c43994b7608744bcc"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train on 38650032 samples, validate on 2034213 samples\n",
      "Epoch 1/1\n",
      " - 1137s - loss: 0.3495 - acc: 0.8344 - val_loss: 0.2634 - val_acc: 0.8789\n"
     ]
    }
   ],
   "source": [
    "#Train[:,:-1]\n",
    "if train:\n",
    "    lr=-5\n",
    "    model.compile(loss=['binary_crossentropy'], optimizer=Adam(lr=10**(lr)), metrics=['accuracy'])\n",
    "    History = model.fit(x=Train[:,:-1], y=Train[:,-1], batch_size=8000, epochs=1, verbose=2, validation_split=0.05, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "29eb9b3b2db4909cf9993435075bf0561b2cc468",
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "if train:\n",
    "    lr=-4\n",
    "    model.compile(loss=['binary_crossentropy'], optimizer=Adam(lr=10**(lr)), metrics=['accuracy'])\n",
    "    History = model.fit(x=Train[:,:-1], y=Train[:,-1], batch_size=8000, epochs=20, verbose=2, validation_split=0.05, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "ab41ee83a17a170f907b0789379420cee8da037f",
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "if train:\n",
    "    lr=-5\n",
    "    model.compile(loss=['binary_crossentropy'], optimizer=Adam(lr=10**(lr)), metrics=['accuracy'])\n",
    "    History = model.fit(x=Train[:,:-1], y=Train[:,-1], batch_size=8000, epochs=3, verbose=2, validation_split=0.05, shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_uuid": "185872dc6a27b54f7b810c66a5076796c8c8ec01"
   },
   "source": [
    "# Step 3 - Hard Negative Mining"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "20a591651544390c1940aedd8c48a4f1526e1f8f",
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# if you skip step2, you still need to run step1 to get training data.\n",
    "if train:\n",
    "    try:\n",
    "        model\n",
    "    except NameError:\n",
    "        print('load model')\n",
    "        model = load_model('../input/trackml/my_model.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "64ce296b3f18f0b4dd368ee232bb9aaa7655365f",
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "if train:\n",
    "    Train_hard = []\n",
    "\n",
    "    for i in tqdm_notebook(range(10,20)):\n",
    "\n",
    "        event = 'event0000010%02d'%i\n",
    "        hits, cells, truth, particles = get_event(event)\n",
    "        hit_cells = cells.groupby(['hit_id']).value.count().values\n",
    "        hit_value = cells.groupby(['hit_id']).value.sum().values\n",
    "        features = np.hstack((hits[['x','y','z']]/1000, hit_cells.reshape(len(hit_cells),1)/10,hit_value.reshape(len(hit_cells),1)))\n",
    "\n",
    "        size=30000000\n",
    "        n = len(truth)\n",
    "        i =np.random.randint(n, size=size)\n",
    "        j =np.random.randint(n, size=size)\n",
    "        p_id = truth.particle_id.values\n",
    "        pair = np.hstack((i.reshape(size,1),j.reshape(size,1)))\n",
    "        pair = pair[((p_id[i]==0) | (p_id[i]!=p_id[j]))]\n",
    "\n",
    "        Train0 = np.hstack((features[pair[:,0]], features[pair[:,1]], np.zeros((len(pair),1))))\n",
    "\n",
    "        pred = model.predict(Train0[:,:-1], batch_size=20000)\n",
    "        s = np.where(pred>0.5)[0]\n",
    "\n",
    "        print(event, len(Train0), len(s))\n",
    "\n",
    "        if len(Train_hard) == 0:\n",
    "            Train_hard = Train0[s]\n",
    "        else:\n",
    "            Train_hard = np.vstack((Train_hard,Train0[s]))\n",
    "    del Train0\n",
    "    print(Train_hard.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "72e5d14bd8617a4e3c5175cea7705cdabcaa0f5e",
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "if train:\n",
    "    Train = np.vstack((Train,Train_hard))\n",
    "    np.random.shuffle(Train)\n",
    "    print(Train.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "ba695b07a4d35af53127564ba11b7884d56fd1a9",
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "if train:\n",
    "    lr=-4\n",
    "    model.compile(loss=['binary_crossentropy'], optimizer=Adam(lr=10**(lr)), metrics=['accuracy'])\n",
    "    History = model.fit(x=Train[:,:-1], y=Train[:,-1], batch_size=8000, epochs=30, verbose=2, validation_split=0.05, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "ffecdec926f9645fe4ddac507787980fbf6142a6",
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "if train:\n",
    "    lr=-5\n",
    "    model.compile(loss=['binary_crossentropy'], optimizer=Adam(lr=10**(lr)), metrics=['accuracy'])\n",
    "    History = model.fit(x=Train[:,:-1], y=Train[:,-1], batch_size=8000, epochs=10, verbose=2, validation_split=0.05, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "b81f5b8a3e882fae96bc923f79720ef3f130ca5d",
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "if train:\n",
    "    lr=-6\n",
    "    model.compile(loss=['binary_crossentropy'], optimizer=Adam(lr=10**(lr)), metrics=['accuracy'])\n",
    "    History = model.fit(x=Train[:,:-1], y=Train[:,-1], batch_size=8000, epochs=2, verbose=2, validation_split=0.05, shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_uuid": "1e5db9b7e283f01e8016e919c0e4d4c4731bf247"
   },
   "source": [
    "# Step 4 - Test event 1001"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "ba7f35bbd13b606908207e4dbea86190a681f8cc"
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    model\n",
    "except NameError:\n",
    "    print('load model')\n",
    "    model = load_model('../input/trackml/my_model_h.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "b87b9218f8426fc7b86d09156d5d461713993e8b",
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "event = 'event000001001'\n",
    "hits, cells, truth, particles = get_event(event)\n",
    "hit_cells = cells.groupby(['hit_id']).value.count().values\n",
    "hit_value = cells.groupby(['hit_id']).value.sum().values\n",
    "features = np.hstack((hits[['x','y','z']]/1000, hit_cells.reshape(len(hit_cells),1)/10,hit_value.reshape(len(hit_cells),1)))\n",
    "count = hits.groupby(['volume_id','layer_id','module_id'])['hit_id'].count().values\n",
    "module_id = np.zeros(len(hits), dtype='int32')\n",
    "\n",
    "for i in range(len(count)):\n",
    "    si = np.sum(count[:i])\n",
    "    module_id[si:si+count[i]] = i"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "f882457cf0e251f4e8e8adce0629d381063a3ca2",
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_path(hit, mask, thr):\n",
    "    path = [hit]\n",
    "    a = 0\n",
    "    while True:\n",
    "        c = get_predict(path[-1], thr/2)\n",
    "        mask = (c > thr)*mask\n",
    "        mask[path[-1]] = 0\n",
    "        \n",
    "        if 1:\n",
    "            cand = np.where(c>thr)[0]\n",
    "            if len(cand)>0:\n",
    "                mask[cand[np.isin(module_id[cand], module_id[path])]]=0\n",
    "                \n",
    "        a = (c + a)*mask\n",
    "        if a.max() < thr*len(path):\n",
    "            break\n",
    "        path.append(a.argmax())\n",
    "    return path\n",
    "\n",
    "def get_predict(hit, thr=0.5):\n",
    "    Tx = np.zeros((len(truth),10))\n",
    "    Tx[:,5:] = features\n",
    "    Tx[:,:5] = np.tile(features[hit], (len(Tx), 1))\n",
    "    pred = model.predict(Tx, batch_size=len(Tx))[:,0]\n",
    "    # TTA\n",
    "    idx = np.where(pred > thr)[0]\n",
    "    Tx2 = np.zeros((len(idx),10))\n",
    "    Tx2[:,5:] = Tx[idx,:5]\n",
    "    Tx2[:,:5] = Tx[idx,5:]    \n",
    "    pred1 = model.predict(Tx2, batch_size=len(idx))[:,0]\n",
    "    pred[idx] = (pred[idx] + pred1)/2\n",
    "    return pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "98400d020f1339b0a467047d3acff43925f39080"
   },
   "outputs": [],
   "source": [
    "# select one hit to construct a track\n",
    "for hit in range(3):\n",
    "    path = get_path(hit, np.ones(len(truth)), 0.95)\n",
    "    gt = np.where(truth.particle_id==truth.particle_id[hit])[0]\n",
    "    print('hit_id = ', hit+1)\n",
    "    print('reconstruct :', path)\n",
    "    print('ground truth:', gt.tolist())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_uuid": "b218782577435528ac3e131eac33e7d4b8f64d75"
   },
   "source": [
    "# Step 5 - Predict and Score\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "6dd697c803477e88b7d0bd98b24e5c3c649898ae"
   },
   "outputs": [],
   "source": [
    "# Predict all pairs for reconstruct by all hits. (takes 2.5hr but can skip)\n",
    "skip_predict = True\n",
    "\n",
    "if skip_predict == False:\n",
    "    TestX = np.zeros((len(features), 10))\n",
    "    TestX[:,5:] = features\n",
    "\n",
    "    # for TTA\n",
    "    TestX1 = np.zeros((len(features), 10))\n",
    "    TestX1[:,:5] = features\n",
    "\n",
    "    preds = []\n",
    "\n",
    "    for i in tqdm_notebook(range(len(features)-1)):\n",
    "        TestX[i+1:,:5] = np.tile(features[i], (len(TestX)-i-1, 1))\n",
    "\n",
    "        pred = model.predict(TestX[i+1:], batch_size=20000)[:,0]                \n",
    "        idx = np.where(pred>0.2)[0]\n",
    "\n",
    "        if len(idx) > 0:\n",
    "            TestX1[idx+i+1,5:] = TestX[idx+i+1,:5]\n",
    "            pred1 = model.predict(TestX1[idx+i+1], batch_size=20000)[:,0]\n",
    "            pred[idx] = (pred[idx]+pred1)/2\n",
    "\n",
    "        idx = np.where(pred>0.5)[0]\n",
    "\n",
    "        preds.append([idx+i+1, pred[idx]])\n",
    "\n",
    "        #if i==0: print(preds[-1])\n",
    "\n",
    "    preds.append([np.array([], dtype='int64'), np.array([], dtype='float32')])\n",
    "\n",
    "    # rebuild to NxN\n",
    "    for i in range(len(preds)):\n",
    "        ii = len(preds)-i-1\n",
    "        for j in range(len(preds[ii][0])):\n",
    "            jj = preds[ii][0][j]\n",
    "            preds[jj][0] = np.insert(preds[jj][0], 0 ,ii)\n",
    "            preds[jj][1] = np.insert(preds[jj][1], 0 ,preds[ii][1][j])\n",
    "\n",
    "    #np.save('my_%s.npy'%event, preds)\n",
    "else:\n",
    "    print('load predicts')\n",
    "    preds = np.load('../input/trackml/my_%s.npy'%event)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "0d34e0bcdc6ba971b4bc7bab1e8c185e17523956",
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_path2(hit, mask, thr):\n",
    "    path = [hit]\n",
    "    a = 0\n",
    "    while True:\n",
    "        c = get_predict2(path[-1])\n",
    "        mask = (c > thr)*mask\n",
    "        mask[path[-1]] = 0\n",
    "        \n",
    "        if 1:\n",
    "            cand = np.where(c>thr)[0]\n",
    "            if len(cand)>0:\n",
    "                mask[cand[np.isin(module_id[cand], module_id[path])]]=0\n",
    "                \n",
    "        a = (c + a)*mask\n",
    "        if a.max() < thr*len(path):\n",
    "            break\n",
    "        path.append(a.argmax())\n",
    "    return path\n",
    "\n",
    "def get_predict2(p):\n",
    "    c = np.zeros(len(preds))\n",
    "    c[preds[p, 0]] = preds[p, 1]          \n",
    "    return c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "0a1af92b508896ec7fc6da0d5f89f153246f73b1"
   },
   "outputs": [],
   "source": [
    "# reconstruct by all hits. (takes 0.6hr but can skip)\n",
    "skip_reconstruct = True\n",
    "\n",
    "if skip_reconstruct == False:\n",
    "    tracks_all = []\n",
    "    thr = 0.85\n",
    "    x4 = True\n",
    "    for hit in tqdm_notebook(range(len(preds))):\n",
    "        m = np.ones(len(truth))\n",
    "        path  = get_path2(hit, m, thr)\n",
    "        if x4 and len(path) > 1:\n",
    "            m[path[1]]=0\n",
    "            path2  = get_path2(hit, m, thr)\n",
    "            if len(path) < len(path2):\n",
    "                path = path2\n",
    "                m[path[1]]=0\n",
    "                path2  = get_path2(hit, m, thr)\n",
    "                if len(path) < len(path2):\n",
    "                    path = path2\n",
    "            elif len(path2) > 1:\n",
    "                m[path[1]]=1\n",
    "                m[path2[1]]=0\n",
    "                path2  = get_path2(hit, m, thr)\n",
    "                if len(path) < len(path2):\n",
    "                    path = path2\n",
    "        tracks_all.append(path)\n",
    "    #np.save('my_tracks_all', tracks_all)\n",
    "else:\n",
    "    print('load tracks')\n",
    "    tracks_all = np.load('../input/trackml/my_tracks_all.npy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "ce9865fb97fd16066e062bc375027cb405cf1f51",
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_track_score(tracks_all, n=4):\n",
    "    scores = np.zeros(len(tracks_all))\n",
    "    for i, path in enumerate(tracks_all):\n",
    "        count = len(path)\n",
    "\n",
    "        if count > 1:\n",
    "            tp=0\n",
    "            fp=0\n",
    "            for p in path:\n",
    "                tp = tp + np.sum(np.isin(tracks_all[p], path, assume_unique=True))\n",
    "                fp = fp + np.sum(np.isin(tracks_all[p], path, assume_unique=True, invert=True))\n",
    "            scores[i] = (tp-fp*n-count)/count/(count-1)\n",
    "        else:\n",
    "            scores[i] = -np.inf\n",
    "    return scores\n",
    "\n",
    "def score_event_fast(truth, submission):\n",
    "    truth = truth[['hit_id', 'particle_id', 'weight']].merge(submission, how='left', on='hit_id')\n",
    "    df = truth.groupby(['track_id', 'particle_id']).hit_id.count().to_frame('count_both').reset_index()\n",
    "    truth = truth.merge(df, how='left', on=['track_id', 'particle_id'])\n",
    "    \n",
    "    df1 = df.groupby(['particle_id']).count_both.sum().to_frame('count_particle').reset_index()\n",
    "    truth = truth.merge(df1, how='left', on='particle_id')\n",
    "    df1 = df.groupby(['track_id']).count_both.sum().to_frame('count_track').reset_index()\n",
    "    truth = truth.merge(df1, how='left', on='track_id')\n",
    "    truth.count_both *= 2\n",
    "    score = truth[(truth.count_both > truth.count_particle) & (truth.count_both > truth.count_track)].weight.sum()\n",
    "    particles = truth[(truth.count_both > truth.count_particle) & (truth.count_both > truth.count_track)].particle_id.unique()\n",
    "\n",
    "    return score, truth[truth.particle_id.isin(particles)].weight.sum(), 1-truth[truth.track_id>0].weight.sum()\n",
    "\n",
    "def evaluate_tracks(tracks, truth):\n",
    "    submission = pd.DataFrame({'hit_id': truth.hit_id, 'track_id': tracks})\n",
    "    score = score_event_fast(truth, submission)[0]\n",
    "    track_id = tracks.max()\n",
    "    print('%.4f %2.2f %4d %5d %.4f %.4f'%(score, np.sum(tracks>0)/track_id, track_id, np.sum(tracks==0), 1-score-np.sum(truth.weight.values[tracks==0]), np.sum(truth.weight.values[tracks==0])))\n",
    "\n",
    "def extend_path(path, mask, thr, last = False):\n",
    "    a = 0\n",
    "    for p in path[:-1]:\n",
    "        c = get_predict2(p)\n",
    "        if last == False:\n",
    "            mask = (c > thr)*mask\n",
    "        mask[p] = 0\n",
    "        cand = np.where(c>thr)[0]\n",
    "        mask[cand[np.isin(module_id[cand], module_id[path])]]=0\n",
    "        a = (c + a)*mask\n",
    "\n",
    "    while True:\n",
    "        c = get_predict2(path[-1])\n",
    "        if last == False:\n",
    "            mask = (c > thr)*mask\n",
    "        mask[path[-1]] = 0\n",
    "        cand = np.where(c>thr)[0]\n",
    "        mask[cand[np.isin(module_id[cand], module_id[path])]]=0\n",
    "        a = (c + a)*mask\n",
    "            \n",
    "        if a.max() < thr*len(path):\n",
    "            break\n",
    "\n",
    "        path.append(a.argmax())\n",
    "        if last: break\n",
    "    \n",
    "    return path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "d4670b7668b14f70b2c64727ae48ff319e2c4552",
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# calculate track's confidence (about 2 mins)\n",
    "scores = get_track_score(tracks_all, 8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "6d5bde61fd95aba5e09536d81360003742857062",
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# merge tracks by confidence and get score\n",
    "idx = np.argsort(scores)[::-1]\n",
    "tracks = np.zeros(len(hits))\n",
    "track_id = 0\n",
    "\n",
    "for hit in idx:\n",
    "\n",
    "    path = np.array(tracks_all[hit])\n",
    "    path = path[np.where(tracks[path]==0)[0]]\n",
    "\n",
    "    if len(path)>3:\n",
    "        track_id = track_id + 1  \n",
    "        tracks[path] = track_id\n",
    "\n",
    "evaluate_tracks(tracks, truth)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "caa0cbbee67c1b4a6eb749d495f865436f3b25eb",
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# multistage\n",
    "idx = np.argsort(scores)[::-1]\n",
    "tracks = np.zeros(len(hits))\n",
    "track_id = 0\n",
    "\n",
    "for hit in idx:\n",
    "    path = np.array(tracks_all[hit])\n",
    "    path = path[np.where(tracks[path]==0)[0]]\n",
    "\n",
    "    if len(path)>6:\n",
    "        track_id = track_id + 1  \n",
    "        tracks[path] = track_id\n",
    "\n",
    "evaluate_tracks(tracks, truth)\n",
    "\n",
    "for track_id in range(1, int(tracks.max())+1):\n",
    "    path = np.where(tracks == track_id)[0]\n",
    "    path = extend_path(path.tolist(), 1*(tracks==0), 0.6)\n",
    "    tracks[path] = track_id\n",
    "        \n",
    "evaluate_tracks(tracks, truth)\n",
    "        \n",
    "for hit in idx:\n",
    "    path = np.array(tracks_all[hit])\n",
    "    path = path[np.where(tracks[path]==0)[0]]\n",
    "\n",
    "    if len(path)>3:\n",
    "        path = extend_path(path.tolist(), 1*(tracks==0), 0.6)\n",
    "        track_id = track_id + 1  \n",
    "        tracks[path] = track_id\n",
    "        \n",
    "evaluate_tracks(tracks, truth)\n",
    "\n",
    "for track_id in range(1, int(tracks.max())+1):\n",
    "    path = np.where(tracks == track_id)[0]\n",
    "    path = extend_path(path.tolist(), 1*(tracks==0), 0.5)\n",
    "    tracks[path] = track_id\n",
    "        \n",
    "evaluate_tracks(tracks, truth)\n",
    "\n",
    "for hit in idx:\n",
    "    path = np.array(tracks_all[hit])\n",
    "    path = path[np.where(tracks[path]==0)[0]]\n",
    "\n",
    "    if len(path)>1:\n",
    "        path = extend_path(path.tolist(), 1*(tracks==0), 0.5)\n",
    "    if len(path)>2:\n",
    "        track_id = track_id + 1\n",
    "        tracks[path] = track_id\n",
    "        \n",
    "evaluate_tracks(tracks, truth)\n",
    "\n",
    "for track_id in range(1, int(tracks.max())+1):\n",
    "    path = np.where(tracks== track_id)[0]\n",
    "    if len(path)%2 == 0:\n",
    "        path = extend_path(path.tolist(), 1*(tracks==0), 0.5, True)\n",
    "        tracks[path] = track_id\n",
    "        \n",
    "evaluate_tracks(tracks, truth)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "64f8a293a76cac80636c80a4adb8fe994fc8cef4"
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "from sklearn.preprocessing import MaxAbsScaler\n",
    "\n",
    "scaler = MaxAbsScaler()\n",
    "\n",
    "train_X = scaler.fit_transform(Train[:-1, :-1])\n",
    "train_X = \n",
    "test_X = scaler.transform(Train[-1, :-1].reshape(1, -1))\n",
    "\n",
    "clf = RandomForestClassifier(n_estimators=500,\n",
    "                             random_state=0, verbose = 1)\n",
    "clf.fit(train_X, Train[:-1, -1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'clf' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<timed eval>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'clf' is not defined"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "clf.score(test_X, Train[-1, -1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}