{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/tim/venvs/analysis-venv-2.7/lib/python2.7/site-packages/theano/tensor/signal/downsample.py:6: UserWarning: downsample module has been moved to the theano.tensor.signal.pool module.\n",
      "  \"downsample module has been moved to the theano.tensor.signal.pool module.\")\n",
      "/Users/tim/venvs/analysis-venv-2.7/lib/python2.7/site-packages/matplotlib/__init__.py:872: UserWarning: axes.color_cycle is deprecated and replaced with axes.prop_cycle; please use the latter.\n",
      "  warnings.warn(self.msg_depr % (key, alt_key))\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Couldn't import dot_parser, loading of dot files will not be possible.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/tim/venvs/analysis-venv-2.7/lib/python2.7/site-packages/IPython/html.py:14: ShimWarning: The `IPython.html` package has been deprecated. You should import from `notebook` instead. `IPython.html.widgets` has moved to `ipywidgets`.\n",
      "  \"`IPython.html.widgets` has moved to `ipywidgets`.\", ShimWarning)\n"
     ]
    }
   ],
   "source": [
    "import mhcflurry, seaborn, numpy, pandas, pickle, sklearn, collections, scipy\n",
    "import mhcflurry.data\n",
    "import sklearn.metrics\n",
    "import sklearn.cross_validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "data = mhcflurry.data.load_allele_datasets(\"/Users/tim/sinai/git/mhcflurry/bdata.2009.mhci.public.1.txt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 0,  0,  0,  0, 16,  1,  0,  9, 17])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a = data[\"HLA-A0201\"]\n",
    "a.X_index[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 672 ms, sys: 31.6 ms, total: 704 ms\n",
      "Wall time: 687 ms\n",
      "CPU times: user 876 ms, sys: 10.9 ms, total: 887 ms\n",
      "Wall time: 908 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([   0,    1,    2, ..., 6652, 6653, 6654], dtype=int32)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def projections(peptide, edit_distance=1):\n",
    "    if edit_distance == 0:\n",
    "        return set([peptide])\n",
    "    return set.union(*[\n",
    "            projections(p, edit_distance - 1)\n",
    "            for p in (peptide[0:i] + \".\" + peptide[(i+1):] for i in range(len(peptide)))\n",
    "    ])\n",
    "\n",
    "def edit_distance_subcomponent_labels(peptides, edit_distance=1):\n",
    "    adjacency_matrix = scipy.sparse.dok_matrix((len(peptides), len(peptides)), dtype=bool)\n",
    "    projection_to_peptide_indices = collections.defaultdict(set)\n",
    "    for (i, peptide) in enumerate(peptides):\n",
    "        for projection in projections(peptide, edit_distance=edit_distance):\n",
    "            projection_to_peptide_indices[projection].add(i)\n",
    "    for indices in projection_to_peptide_indices.values():\n",
    "        for index1 in indices:\n",
    "            for index2 in indices:\n",
    "                adjacency_matrix[index1, index2] = True\n",
    "    \n",
    "    (num_components, labels) = scipy.sparse.csgraph.connected_components(adjacency_matrix, directed=False)\n",
    "    return labels\n",
    "\n",
    "def balanced_folds(labels, num_folds=2):\n",
    "    # Simple greedy approximation to balanced partition problem\n",
    "    folds = [set() for i in range(num_folds)]\n",
    "    df = pandas.Series(labels).to_frame().reset_index()\n",
    "    df.columns = [\"point\", \"label\"]\n",
    "    grouped = df.groupby(\"label\").apply(lambda group: list(group.point)).to_frame()\n",
    "    grouped.columns = [\"indices\"]\n",
    "    grouped[\"count\"] = [len(x) for x in grouped.indices]\n",
    "    grouped.sort(\"count\", ascending=False, inplace=True)\n",
    "    for (label, row) in grouped.iterrows():\n",
    "        min_fold = folds[numpy.argmin([len(x) for x in folds])]\n",
    "        min_fold.update(row.indices)\n",
    "    \n",
    "    return [\n",
    "        numpy.array(sorted(fold))\n",
    "        for fold in folds\n",
    "    ]\n",
    "    \n",
    "def cv_indices(peptides, edit_distance=1, num_folds=3, shuffle=True):\n",
    "    if shuffle:\n",
    "        def maybe_shuffle(indices):\n",
    "            return numpy.random.permutation(indices)\n",
    "    else:\n",
    "        def maybe_shuffle(indices):\n",
    "            return indices\n",
    "    \n",
    "    labels = edit_distance_subcomponent_labels(peptides, edit_distance=edit_distance)\n",
    "    folds = balanced_folds(labels, num_folds=num_folds)\n",
    "    all_indices = set(numpy.arange(len(peptides)))\n",
    "    for test_fold in folds:\n",
    "        train_fold = numpy.array(list(all_indices.difference(test_fold)))\n",
    "        yield (maybe_shuffle(train_fold), maybe_shuffle(test_fold))\n",
    "        \n",
    "    \n",
    "%time labels = edit_distance_subcomponent_labels(a.peptides, 1)\n",
    "%time folds = balanced_folds(labels, 5)\n",
    "labels\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 95.5 ms, sys: 22.3 ms, total: 118 ms\n",
      "Wall time: 111 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(6961,)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%time labels = edit_distance_subcomponent_labels(a.peptides, 0)\n",
    "labels.max()\n",
    "labels.shape\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(array([3244, 4894, 6646, ..., 6832,  780,  692]), array([ 767, 3387, 1328, ..., 4490, 3659, 1911]))\n",
      "(array([4220, 5170,  874, ...,  454, 1144, 1778]), array([3446,  997, 3655, ..., 4569, 1844, 3468]))\n",
      "(array([3549, 1279, 2552, ..., 2142, 2311,  675]), array([2566, 6411, 4595, ..., 2731, 2938, 1601]))\n"
     ]
    }
   ],
   "source": [
    "for (train, test) in cv_indices(a.peptides, edit_distance=0):\n",
    "    print(train, test)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<module 'mhcflurry.class1_binding_predictor' from '/Users/tim/sinai/git/mhcflurry/mhcflurry/class1_binding_predictor.py'>"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "reload(mhcflurry)\n",
    "reload(mhcflurry.class1_binding_predictor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "model = mhcflurry.class1_binding_predictor.Class1BindingPredictor(\n",
    "    embedding_input_dim=21,\n",
    "    embedding_output_dim=32,\n",
    "    learning_rate=0.01,\n",
    "    activation=\"tanh\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 0.86992363,  0.84365301,  0.83955602])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def roc_auc_ic50_threshold(y, y_pred, threshold_nm=500):\n",
    "    return sklearn.metrics.roc_auc_score(\n",
    "        model.log_to_ic50(y) < threshold_nm,\n",
    "        model.log_to_ic50(y_pred) < threshold_nm)\n",
    "    \n",
    "scorer = sklearn.metrics.make_scorer(roc_auc_ic50_threshold)\n",
    "sklearn.cross_validation.cross_val_score(model, a.X_index, a.Y, scorer, cv_indices(a.peptides))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 0.84956871,  0.87736409,  0.86106563])"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sklearn.cross_validation.cross_val_score(model, a.X_index, a.Y, scorer, cv_indices(a.peptides, edit_distance=0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 0.8549138 ,  0.85409191,  0.87046823])"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sklearn.cross_validation.cross_val_score(model, a.X_index, a.Y, scorer, cv_indices(a.peptides, edit_distance=2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(14,)"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data[\"Gogo-B0101\"].Y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('Gogo-B0101', array([ 0.35,  0.25,  0.5 ]))\n",
      "('H-2-DB', array([ 0.80445675,  0.84123684,  0.89734347]))\n",
      "('H-2-DD', array([ 0.58387097,  0.60912698,  0.59325397]))\n",
      "('H-2-KB', array([ 0.73026071,  0.68359534,  0.7560893 ]))\n",
      "('H-2-KD', array([ 0.76336898,  0.70289855,  0.79196487]))\n",
      "('H-2-KK', array([ 0.68577075,  0.67935484,  0.78709677]))\n",
      "('H-2-LD', array([ 0.63095238,  0.86363636,  0.85151515]))\n",
      "('HLA-A0101', array([ 0.82295861,  0.79946528,  0.86610991]))\n",
      "('HLA-A0201', array([ 0.87802025,  0.86698864,  0.86856158]))\n",
      "('HLA-A0202', array([ 0.84085982,  0.85036179,  0.8273601 ]))\n",
      "('HLA-A0201', array([ 0.87802025,  0.86698864,  0.86856158]))\n",
      "('HLA-A0202', array([ 0.84085982,  0.85036179,  0.8273601 ]))\n",
      "('H-2-DB', array([ 0.80445675,  0.84123684,  0.89734347]))\n",
      "('H-2-DD', array([ 0.58387097,  0.60912698,  0.59325397]))\n",
      "('H-2-KK', array([ 0.68577075,  0.67935484,  0.78709677]))\n",
      "('Gogo-B0101', array([ 0.35,  0.25,  0.5 ]))\n",
      "('H-2-LD', array([ 0.63095238,  0.86363636,  0.85151515]))\n",
      "('H-2-KB', array([ 0.73026071,  0.68359534,  0.7560893 ]))\n",
      "('H-2-KD', array([ 0.76336898,  0.70289855,  0.79196487]))\n",
      "('HLA-A0101', array([ 0.82295861,  0.79946528,  0.86610991]))\n"
     ]
    }
   ],
   "source": [
    "result = {}\n",
    "for allele in list(sorted(data))[:10]:\n",
    "    values = data[allele]\n",
    "    result[allele] = sklearn.cross_validation.cross_val_score(\n",
    "        model, values.X_index, values.Y, scorer, cv_indices(values.peptides))\n",
    "    print(allele, result[allele])\n",
    "    \n",
    "for (key, values) in result.iteritems():\n",
    "    print(key, values)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('HLA-A0201', 6961)\n",
      "('HLA-A0201', array([ 0.88591598,  0.8921594 ,  0.8863821 ]))\n",
      "('HLA-A0301', 4601)\n",
      "('HLA-A0301', array([ 0.82211658,  0.84389135,  0.80404472]))\n",
      "('HLA-A0203', 3937)\n",
      "('HLA-A0203', array([ 0.87039776,  0.87301872,  0.85440652]))\n",
      "('HLA-A1101', 3862)\n",
      "('HLA-A1101', array([ 0.89471562,  0.81754575,  0.87795699]))\n",
      "('HLA-A3101', 3309)\n",
      "('HLA-A3101', array([ 0.84547391,  0.83338317,  0.80496805]))\n",
      "('HLA-A0206', 3223)\n",
      "('HLA-A0206', array([ 0.81704192,  0.82848796,  0.84186113]))\n",
      "('HLA-A6802', 3188)\n",
      "('HLA-A6802', array([ 0.82971764,  0.8122962 ,  0.79950669]))\n",
      "('HLA-A0101', 3169)\n",
      "('HLA-A0101', array([ 0.81362577,  0.78240669,  0.82866379]))\n",
      "('HLA-B1501', 3142)\n",
      "('HLA-B1501', array([ 0.81558246,  0.76136838,  0.80164387]))\n",
      "('HLA-B0702', 2974)\n",
      "('HLA-B0702', array([ 0.87536917,  0.85384887,  0.85821342]))\n",
      "('HLA-A0206', array([ 0.81704192,  0.82848796,  0.84186113]))\n",
      "('HLA-A1101', array([ 0.89471562,  0.81754575,  0.87795699]))\n",
      "('HLA-A0203', array([ 0.87039776,  0.87301872,  0.85440652]))\n",
      "('HLA-B0702', array([ 0.87536917,  0.85384887,  0.85821342]))\n",
      "('HLA-A3101', array([ 0.84547391,  0.83338317,  0.80496805]))\n",
      "('HLA-A0301', array([ 0.82211658,  0.84389135,  0.80404472]))\n",
      "('HLA-B1501', array([ 0.81558246,  0.76136838,  0.80164387]))\n",
      "('HLA-A6802', array([ 0.82971764,  0.8122962 ,  0.79950669]))\n",
      "('HLA-A0101', array([ 0.81362577,  0.78240669,  0.82866379]))\n",
      "('HLA-A0201', array([ 0.88591598,  0.8921594 ,  0.8863821 ]))\n"
     ]
    }
   ],
   "source": [
    "model = mhcflurry.class1_binding_predictor.Class1BindingPredictor(\n",
    "    embedding_input_dim=21,\n",
    "    embedding_output_dim=32,\n",
    "    learning_rate=0.01,\n",
    "    activation=\"tanh\")\n",
    "\n",
    "result = {}\n",
    "for allele in list(sorted(data, key=lambda x: -1 * data[x].Y.shape[0]))[:10]:\n",
    "    values = data[allele]\n",
    "    #print(allele, values.Y.shape[0])\n",
    "    result[allele] = sklearn.cross_validation.cross_val_score(\n",
    "        model, values.X_index, values.Y, scorer, cv_indices(values.peptides))\n",
    "    print(allele, result[allele])\n",
    "    \n",
    "for (key, values) in result.iteritems():\n",
    "    print(key, values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'Gogo-B0101': array([ 0.35,  0.  ,  0.5 ]),\n",
       " 'H-2-DB': array([ 0.80259192,  0.80323763,  0.88438937]),\n",
       " 'H-2-DD': array([ 0.6       ,  0.48412698,  0.60912698]),\n",
       " 'H-2-KB': array([ 0.74255326,  0.70298595,  0.71285589]),\n",
       " 'H-2-KD': array([ 0.73729947,  0.67687747,  0.79033832]),\n",
       " 'H-2-KK': array([ 0.63372859,  0.70322581,  0.72709677]),\n",
       " 'H-2-LD': array([ 0.70238095,  0.79186603,  0.8       ]),\n",
       " 'HLA-A0101': array([ 0.79496009,  0.77293741,  0.8203125 ]),\n",
       " 'HLA-A0201': array([ 0.86748346,  0.85718596,  0.84801616])}"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
