{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Machine learning - Features extraction\n",
    "\n",
    "Demo to create a feature vector for protein fold classification. \n",
    "In this demo we try to classify a protein chain as either an all alpha or all beta protein based on protein sequence. We use n-grams and a Word2Vec representation of the protein sequence as a feature vector.\n",
    "\n",
    "[Word2Vec model](https://spark.apache.org/docs/latest/mllib-feature-extraction.html#word2vec)\n",
    "\n",
    "[Word2Vec example](https://spark.apache.org/docs/latest/ml-features.html#word2vec)\n",
    "\n",
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql import SparkSession\n",
    "from pyspark.sql.functions import *\n",
    "from pyspark.sql.types import *\n",
    "from mmtfPyspark.io import mmtfReader\n",
    "from mmtfPyspark.webfilters import Pisces\n",
    "from mmtfPyspark.filters import ContainsLProteinChain\n",
    "from mmtfPyspark.mappers import StructureToPolymerChains\n",
    "from mmtfPyspark.datasets import secondaryStructureExtractor\n",
    "from mmtfPyspark.ml import ProteinSequenceEncoder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configure Spark Context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "spark = SparkSession.builder.appName(\"1-Features\").getOrCreate()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read MMTF File and get a set of L-protein chains"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "pdb = mmtfReader.read_sequence_file('../resources/mmtf_reduced_sample/') \\\n",
    "                .flatMap(StructureToPolymerChains()) \\\n",
    "                .filter(ContainsLProteinChain())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get secondary structure content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = secondaryStructureExtractor.get_dataset(pdb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----------------+--------------------+----------+----------+----------+--------------------+--------------------+\n",
      "|structureChainId|            sequence|     alpha|      beta|      coil|          dsspQ8Code|          dsspQ3Code|\n",
      "+----------------+--------------------+----------+----------+----------+--------------------+--------------------+\n",
      "|          4WMY.A|TDWSHPQFEKSTDEANT...|0.19081272|0.26855123|0.54063606|XXXXXXXXXXXXXXXXX...|XXXXXXXXXXXXXXXXX...|\n",
      "|          4WMY.B|TDWSHPQFEKSTDEANT...|0.17081851|0.26334518| 0.5658363|XXXXXXXXXXXXXXXXX...|XXXXXXXXXXXXXXXXX...|\n",
      "|          4WN5.A|GSHMGRGAFLSRHSLDM...| 0.2962963|0.37962964|0.32407406|XXCCCCCCEEEEECTTC...|XXCCCCCCEEEEECCCC...|\n",
      "|          4WN5.B|GSHMGRGAFLSRHSLDM...|0.33333334|0.37142858| 0.2952381|XXXXXCCCEEEEECTTC...|XXXXXCCCEEEEECCCC...|\n",
      "|          4WND.A|GPGSMEASCLELALEGE...| 0.8358663|       0.0|0.16413374|XXXXCCSCHHHHHHHHH...|XXXXCCCCHHHHHHHHH...|\n",
      "+----------------+--------------------+----------+----------+----------+--------------------+--------------------+\n",
      "only showing top 5 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "data.show(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define add_protein_fold_type function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_protein_fold_type(data, minThreshold, maxThreshold):\n",
    "    '''\n",
    "    Adds a column \"foldType\" with three major secondary structure class:\n",
    "    \"alpha\", \"beta\", \"alpha+beta\", and \"other\" based upon the fraction of alpha/beta content.\n",
    "\n",
    "    The simplified syntax used in this method relies on two imports:\n",
    "        from pyspark.sql.functions import when\n",
    "        from pyspark.sql.functions import col\n",
    "\n",
    "    Attributes:\n",
    "        data (Dataset<Row>): input dataset with alpha, beta composition\n",
    "        minThreshold (float): below this threshold, the secondary structure is ignored\n",
    "        maxThreshold (float): above this threshold, the secondary structure is ignored\n",
    "    '''\n",
    "\n",
    "    return data.withColumn(\"foldType\", \\\n",
    "                           when((col(\"alpha\") > maxThreshold) & (col(\"beta\") < minThreshold), \"alpha\"). \\\n",
    "                           when((col(\"beta\") > maxThreshold) & (col(\"alpha\") < minThreshold), \"beta\"). \\\n",
    "                           when((col(\"alpha\") > maxThreshold) & (col(\"beta\") > maxThreshold), \"alpha+beta\"). \\\n",
    "                           otherwise(\"other\")\\\n",
    "                           )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Classify chains by secondary structure type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = add_protein_fold_type(data, minThreshold=0.05, maxThreshold=0.15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----------------+--------------------+-----------+-----------+----------+--------------------+--------------------+----------+\n",
      "|structureChainId|            sequence|      alpha|       beta|      coil|          dsspQ8Code|          dsspQ3Code|  foldType|\n",
      "+----------------+--------------------+-----------+-----------+----------+--------------------+--------------------+----------+\n",
      "|          4WMY.A|TDWSHPQFEKSTDEANT...| 0.19081272| 0.26855123|0.54063606|XXXXXXXXXXXXXXXXX...|XXXXXXXXXXXXXXXXX...|alpha+beta|\n",
      "|          4WMY.B|TDWSHPQFEKSTDEANT...| 0.17081851| 0.26334518| 0.5658363|XXXXXXXXXXXXXXXXX...|XXXXXXXXXXXXXXXXX...|alpha+beta|\n",
      "|          4WN5.A|GSHMGRGAFLSRHSLDM...|  0.2962963| 0.37962964|0.32407406|XXCCCCCCEEEEECTTC...|XXCCCCCCEEEEECCCC...|alpha+beta|\n",
      "|          4WN5.B|GSHMGRGAFLSRHSLDM...| 0.33333334| 0.37142858| 0.2952381|XXXXXCCCEEEEECTTC...|XXXXXCCCEEEEECCCC...|alpha+beta|\n",
      "|          4WND.A|GPGSMEASCLELALEGE...|  0.8358663|        0.0|0.16413374|XXXXCCSCHHHHHHHHH...|XXXXCCCCHHHHHHHHH...|     alpha|\n",
      "|          4WND.B|GPLGSDLPPKVVPSKQL...|0.115384616|        0.0|0.88461536|XXXXXXXXXXXXXXXCC...|XXXXXXXXXXXXXXXCC...|     other|\n",
      "|          4WP6.A|GSHHHHHHSQDPMQAAQ...| 0.45695364|0.119205296|0.42384106|XXXXXXXXXXXXXXXXX...|XXXXXXXXXXXXXXXXX...|     other|\n",
      "|          4WP9.A|FQGAMGSRVVILFTDIE...|  0.3939394|  0.3151515|0.29090908|XXCCSSEEEEEEEEEET...|XXCCCCEEEEEEEEEEC...|alpha+beta|\n",
      "|          4WP9.B|FQGAMGSRVVILFTDIE...|        0.4|     0.3125|    0.2875|XXXCCSEEEEEEEEEET...|XXXCCCEEEEEEEEEEC...|alpha+beta|\n",
      "|          4WPG.A|GPLLEMILITGSNGQLG...| 0.39372823| 0.17073171|0.43554008|XCCSCCEEEESTTSHHH...|XCCCCCEEEECCCCHHH...|alpha+beta|\n",
      "|          4WPK.A|MHHHHHHGMASMTARPL...|  0.4122807|0.114035085|0.47368422|XXXXXXXXXXCTTTSCH...|XXXXXXXXXXCCCCCCH...|     other|\n",
      "|          4WQD.A|MEPPTVALTVPAAALLP...|  0.3991228|0.057017542|0.54385966|XXXXCBCCCCCCGGGCC...|XXXXCECCCCCCHHHCC...|     other|\n",
      "|          4WRI.A|GILANLKEPSAHWCRKM...| 0.62032086|0.053475935| 0.3262032|XXXXXCCCCCHHHHHHH...|XXXXXCCCCCHHHHHHH...|     other|\n",
      "|          4WSF.A|TTDTRRRVKLYALNAER...| 0.16216215|  0.4774775|0.36036035|XXCCTTEEEEEEECTTS...|XXCCCCEEEEEEECCCC...|alpha+beta|\n",
      "|          4WSF.B| PDESSADVVFKKPLAPAPR|        0.0|        0.0|       1.0| XXXXXXXCCSCCCSSCCCX| XXXXXXXCCCCCCCCCCCX|     other|\n",
      "|          1GWM.A|MNVRATYTVIFKNASGL...|0.039215688|   0.503268|0.45751634|CCCSCCEEEEESSCSSS...|CCCCCCEEEEECCCCCC...|      beta|\n",
      "|          1GXM.A|GLVPRGSHMTGRMLTLD...| 0.42901236| 0.13580246| 0.4351852|XXXXXXXXCBTTBCCCT...|XXXXXXXXCECCECCCC...|     other|\n",
      "|          1GXM.B|GLVPRGSHMTGRMLTLD...|  0.4186747| 0.12951808|0.45180723|CCCCTTTTCBTTBCCCT...|CCCCCCCCCECCECCCC...|     other|\n",
      "|          1GXR.A|DYFQGAMGSKPAYSFHV...|        0.0|  0.5432836|0.45671642|CCEEEEEEEEECCEEEE...|CCEEEEEEEEECCEEEE...|      beta|\n",
      "|          1GXR.B|DYFQGAMGSKPAYSFHV...|        0.0|  0.5555556|0.44444445|CCEEEEEEEEECCEEET...|CCEEEEEEEEECCEEEC...|      beta|\n",
      "+----------------+--------------------+-----------+-----------+----------+--------------------+--------------------+----------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "data.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a Word2Vec representation of the protein sequences\n",
    "\n",
    "**n = 2**     # create 2-grams \n",
    "\n",
    "**windowSize = 25**    # 25-amino residue window size for Word2Vector\n",
    "\n",
    "**vectorSize = 50**    # dimension of feature vector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "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>structureChainId</th>\n",
       "      <th>sequence</th>\n",
       "      <th>alpha</th>\n",
       "      <th>beta</th>\n",
       "      <th>coil</th>\n",
       "      <th>dsspQ8Code</th>\n",
       "      <th>dsspQ3Code</th>\n",
       "      <th>foldType</th>\n",
       "      <th>ngram</th>\n",
       "      <th>features</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>4WMY.A</td>\n",
       "      <td>TDWSHPQFEKSTDEANTYFKEWTCSSSPSLPRSCKEIKDECPSAFD...</td>\n",
       "      <td>0.190813</td>\n",
       "      <td>0.268551</td>\n",
       "      <td>0.540636</td>\n",
       "      <td>XXXXXXXXXXXXXXXXXXXXXXXCCCCCCCCSSHHHHHHHCTTCCS...</td>\n",
       "      <td>XXXXXXXXXXXXXXXXXXXXXXXCCCCCCCCCCHHHHHHHCCCCCC...</td>\n",
       "      <td>alpha+beta</td>\n",
       "      <td>[TD, DW, WS, SH, HP, PQ, QF, FE, EK, KS, ST, T...</td>\n",
       "      <td>[0.028354697964596942, 0.06656068684991266, 0....</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>4WMY.B</td>\n",
       "      <td>TDWSHPQFEKSTDEANTYFKEWTCSSSPSLPRSCKEIKDECPSAFD...</td>\n",
       "      <td>0.170819</td>\n",
       "      <td>0.263345</td>\n",
       "      <td>0.565836</td>\n",
       "      <td>XXXXXXXXXXXXXXXXXXXXXCCCXXXXCCCSSHHHHHHHCTTCCS...</td>\n",
       "      <td>XXXXXXXXXXXXXXXXXXXXXCCCXXXXCCCCCHHHHHHHCCCCCC...</td>\n",
       "      <td>alpha+beta</td>\n",
       "      <td>[TD, DW, WS, SH, HP, PQ, QF, FE, EK, KS, ST, T...</td>\n",
       "      <td>[0.028354697964596942, 0.06656068684991266, 0....</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>4WN5.A</td>\n",
       "      <td>GSHMGRGAFLSRHSLDMKFTYCDDRIAEVAGYSPDDLIGCSAYEYI...</td>\n",
       "      <td>0.296296</td>\n",
       "      <td>0.379630</td>\n",
       "      <td>0.324074</td>\n",
       "      <td>XXCCCCCCEEEEECTTCBEEEECGGHHHHHSCCHHHHBTSBGGGGB...</td>\n",
       "      <td>XXCCCCCCEEEEECCCCEEEEECHHHHHHHCCCHHHHECCEHHHHE...</td>\n",
       "      <td>alpha+beta</td>\n",
       "      <td>[GS, SH, HM, MG, GR, RG, GA, AF, FL, LS, SR, R...</td>\n",
       "      <td>[-0.04048257577641491, 0.1233881547426184, 0.3...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4WN5.B</td>\n",
       "      <td>GSHMGRGAFLSRHSLDMKFTYCDDRIAEVAGYSPDDLIGCSAYEYI...</td>\n",
       "      <td>0.333333</td>\n",
       "      <td>0.371429</td>\n",
       "      <td>0.295238</td>\n",
       "      <td>XXXXXCCCEEEEECTTCBEEEECGGHHHHHSCCHHHHBTSBGGGGB...</td>\n",
       "      <td>XXXXXCCCEEEEECCCCEEEEECHHHHHHHCCCHHHHECCEHHHHE...</td>\n",
       "      <td>alpha+beta</td>\n",
       "      <td>[GS, SH, HM, MG, GR, RG, GA, AF, FL, LS, SR, R...</td>\n",
       "      <td>[-0.04048257577641491, 0.1233881547426184, 0.3...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>4WND.A</td>\n",
       "      <td>GPGSMEASCLELALEGERLCKSGDCRAGVSFFEAAVQVGTEDLKTL...</td>\n",
       "      <td>0.835866</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.164134</td>\n",
       "      <td>XXXXCCSCHHHHHHHHHHHHHTTCHHHHHHHHHHHHHHCCSCHHHH...</td>\n",
       "      <td>XXXXCCCCHHHHHHHHHHHHHCCCHHHHHHHHHHHHHHCCCCHHHH...</td>\n",
       "      <td>alpha</td>\n",
       "      <td>[GP, PG, GS, SM, ME, EA, AS, SC, CL, LE, EL, L...</td>\n",
       "      <td>[-0.009619595496742813, 0.03677304709491171, 0...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  structureChainId                                           sequence  \\\n",
       "0           4WMY.A  TDWSHPQFEKSTDEANTYFKEWTCSSSPSLPRSCKEIKDECPSAFD...   \n",
       "1           4WMY.B  TDWSHPQFEKSTDEANTYFKEWTCSSSPSLPRSCKEIKDECPSAFD...   \n",
       "2           4WN5.A  GSHMGRGAFLSRHSLDMKFTYCDDRIAEVAGYSPDDLIGCSAYEYI...   \n",
       "3           4WN5.B  GSHMGRGAFLSRHSLDMKFTYCDDRIAEVAGYSPDDLIGCSAYEYI...   \n",
       "4           4WND.A  GPGSMEASCLELALEGERLCKSGDCRAGVSFFEAAVQVGTEDLKTL...   \n",
       "\n",
       "      alpha      beta      coil  \\\n",
       "0  0.190813  0.268551  0.540636   \n",
       "1  0.170819  0.263345  0.565836   \n",
       "2  0.296296  0.379630  0.324074   \n",
       "3  0.333333  0.371429  0.295238   \n",
       "4  0.835866  0.000000  0.164134   \n",
       "\n",
       "                                          dsspQ8Code  \\\n",
       "0  XXXXXXXXXXXXXXXXXXXXXXXCCCCCCCCSSHHHHHHHCTTCCS...   \n",
       "1  XXXXXXXXXXXXXXXXXXXXXCCCXXXXCCCSSHHHHHHHCTTCCS...   \n",
       "2  XXCCCCCCEEEEECTTCBEEEECGGHHHHHSCCHHHHBTSBGGGGB...   \n",
       "3  XXXXXCCCEEEEECTTCBEEEECGGHHHHHSCCHHHHBTSBGGGGB...   \n",
       "4  XXXXCCSCHHHHHHHHHHHHHTTCHHHHHHHHHHHHHHCCSCHHHH...   \n",
       "\n",
       "                                          dsspQ3Code    foldType  \\\n",
       "0  XXXXXXXXXXXXXXXXXXXXXXXCCCCCCCCCCHHHHHHHCCCCCC...  alpha+beta   \n",
       "1  XXXXXXXXXXXXXXXXXXXXXCCCXXXXCCCCCHHHHHHHCCCCCC...  alpha+beta   \n",
       "2  XXCCCCCCEEEEECCCCEEEEECHHHHHHHCCCHHHHECCEHHHHE...  alpha+beta   \n",
       "3  XXXXXCCCEEEEECCCCEEEEECHHHHHHHCCCHHHHECCEHHHHE...  alpha+beta   \n",
       "4  XXXXCCCCHHHHHHHHHHHHHCCCHHHHHHHHHHHHHHCCCCHHHH...       alpha   \n",
       "\n",
       "                                               ngram  \\\n",
       "0  [TD, DW, WS, SH, HP, PQ, QF, FE, EK, KS, ST, T...   \n",
       "1  [TD, DW, WS, SH, HP, PQ, QF, FE, EK, KS, ST, T...   \n",
       "2  [GS, SH, HM, MG, GR, RG, GA, AF, FL, LS, SR, R...   \n",
       "3  [GS, SH, HM, MG, GR, RG, GA, AF, FL, LS, SR, R...   \n",
       "4  [GP, PG, GS, SM, ME, EA, AS, SC, CL, LE, EL, L...   \n",
       "\n",
       "                                            features  \n",
       "0  [0.028354697964596942, 0.06656068684991266, 0....  \n",
       "1  [0.028354697964596942, 0.06656068684991266, 0....  \n",
       "2  [-0.04048257577641491, 0.1233881547426184, 0.3...  \n",
       "3  [-0.04048257577641491, 0.1233881547426184, 0.3...  \n",
       "4  [-0.009619595496742813, 0.03677304709491171, 0...  "
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "encoder = ProteinSequenceEncoder(data)\n",
    "data = encoder.overlapping_ngram_word2vec_encode(n=2, windowSize=25, vectorSize=50).cache()\n",
    "\n",
    "data.toPandas().head(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Keep only a subset of relevant fields for further processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = data.select(['structureChainId','alpha','beta','coil','foldType','features'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Write to parquet file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "data.write.mode('overwrite').format('parquet').save('./input_features')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Terminate Spark"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "spark.stop()"
   ]
  }
 ],
 "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}