kMer Sequence Model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9c1d77e7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-07-07 23:57:54.004524: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2023-07-07 23:57:54.443736: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Num GPUs Available:  1\n",
      "/mnt/e/ML/cafa-5-protein-function-prediction\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-07-07 23:57:55.365123: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:982] could not open file to read NUMA node: /sys/bus/pci/devices/0000:2b:00.0/numa_node\n",
      "Your kernel may have been built without NUMA support.\n",
      "2023-07-07 23:57:55.447095: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:982] could not open file to read NUMA node: /sys/bus/pci/devices/0000:2b:00.0/numa_node\n",
      "Your kernel may have been built without NUMA support.\n",
      "2023-07-07 23:57:55.447408: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:982] could not open file to read NUMA node: /sys/bus/pci/devices/0000:2b:00.0/numa_node\n",
      "Your kernel may have been built without NUMA support.\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import pandas as pd\n",
    "import datetime\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "physical_devices = tf.config.list_physical_devices('GPU')\n",
    "print(\"Num GPUs Available: \", len(physical_devices))\n",
    "# try:\n",
    "#   tf.config.experimental.set_memory_growth(physical_devices[0], True)\n",
    "# except:\n",
    "#   # Invalid device or cannot modify virtual devices once initialized.\n",
    "#   pass\n",
    "\n",
    "load_dotenv(override=True)\n",
    "\n",
    "DATA_PATH = os.getenv('DATA_PATH')\n",
    "print(DATA_PATH)\n",
    "\n",
    "# Choose subontology (CCO, MFO or BPO)\n",
    "SO = 'BPO'"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e3544f8a",
   "metadata": {},
   "source": [
    "## Reading fasta, obo and tsv files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c8b515f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from Bio import SeqIO\n",
    "\n",
    "sequences = [rec.seq for rec in SeqIO.parse(os.path.join(DATA_PATH, \"Train/train_sequences.fasta\"),\"fasta\")]\n",
    "ids = [rec.id for rec in SeqIO.parse(os.path.join(DATA_PATH, \"Train/train_sequences.fasta\"),\"fasta\")]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2898414e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import networkx\n",
    "import obonet\n",
    "\n",
    "# Read the taxrank ontology\n",
    "url = os.path.join(DATA_PATH, \"Train/go-basic.obo\")\n",
    "graph = obonet.read_obo(url)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9f4bf949",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(3497732, 3)\n"
     ]
    }
   ],
   "source": [
    "df = pd.read_csv(os.path.join(DATA_PATH, \"Train/train_terms.tsv\"), sep='\\t')\n",
    "\n",
    "dfSO = df.loc[df[\"aspect\"]==SO]\n",
    "uniqueTerms = dfSO[\"term\"].unique()\n",
    "termsArr = list(dfSO[\"term\"].to_numpy())\n",
    "\n",
    "uniqueTermsDict={}\n",
    "for i,el in enumerate(uniqueTerms):\n",
    "    uniqueTermsDict[el] = i\n",
    "    \n",
    "print(dfSO.shape)\n",
    "df=dfSO\n",
    "\n",
    "df.set_index(\"EntryID\", inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e4dec705",
   "metadata": {},
   "outputs": [],
   "source": [
    "testID = df.index.to_list()[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b4e4bac6",
   "metadata": {},
   "outputs": [],
   "source": [
    "dfGo = pd.read_csv(os.path.join(DATA_PATH, \"Train/train_terms.tsv\"), sep='\\t')\n",
    "\n",
    "dfGo = dfGo.loc[dfGo[\"aspect\"]==SO]\n",
    "\n",
    "dfGo.set_index(\"term\", inplace=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5338b7f8",
   "metadata": {},
   "source": [
    "## GO analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1058ef1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "item_counts = df[\"term\"].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "eb35584f",
   "metadata": {},
   "outputs": [],
   "source": [
    "id_to_name = {id_: data.get('name') for id_, data in graph.nodes(data=True)}\n",
    "name_to_id = {data['name']: id_ for id_, data in graph.nodes(data=True) if 'name' in data}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "4d226631",
   "metadata": {},
   "outputs": [],
   "source": [
    "relevantNodes=[]\n",
    "for layer, nodes in enumerate(networkx.topological_generations(graph)):\n",
    "    if layer>11:\n",
    "        relevantNodes = relevantNodes + nodes"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "3c099fc2",
   "metadata": {},
   "source": [
    "## Label encoding"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "cae9a3b0",
   "metadata": {},
   "source": [
    "The task is a multilabel classification: The output has several possible targets (Gene Ontologies) but each can only be 1 (existing) or 0 (non existing)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c8e63ed9",
   "metadata": {},
   "source": [
    "Extract label weights from IA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "c7e3c316",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Not found GOs: 0 (set to 0)\n"
     ]
    }
   ],
   "source": [
    "dfIa = pd.read_csv(os.path.join(DATA_PATH, \"IA.txt\"), sep='\\t', header=None)\n",
    "\n",
    "dfIa.set_index(0, inplace=True)\n",
    "\n",
    "labelWeights=[]\n",
    "allIndices = dfIa.index.tolist()\n",
    "\n",
    "# if SO == \"BPO\":\n",
    "#     item_counts=item_counts[0:10]\n",
    "#     # item_counts=item_counts[0:len(item_counts)//2]\n",
    "\n",
    "notFound=0\n",
    "for go in item_counts.index.to_list():\n",
    "    if go in allIndices:\n",
    "        labelWeights.append(dfIa.loc[go].to_numpy()[0])\n",
    "    else:\n",
    "        notFound += 1\n",
    "        labelWeights.append(0)\n",
    "\n",
    "print(\"Not found GOs: {} (set to 0)\".format(notFound))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "353d72fb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading presaved data\n"
     ]
    }
   ],
   "source": [
    "topGOs=item_counts.index.to_list()\n",
    "\n",
    "threshold=0\n",
    "labelWeights=np.array(labelWeights)\n",
    "selection = labelWeights>threshold\n",
    "topGOs=np.array(topGOs)[selection]\n",
    "\n",
    "if os.path.exists(os.path.join(DATA_PATH, \"GODataSizes_\"+SO+\".npy\")):\n",
    "    print(\"Loading presaved data\")\n",
    "    GODataSizes = np.load(os.path.join(DATA_PATH, \"GODataSizes_\"+SO+\".npy\"))\n",
    "else:\n",
    "    GODataSizes= [dfGo.loc[g].size for g in topGOs]\n",
    "    np.save(os.path.join(DATA_PATH, \"GODataSizes_\"+SO), GODataSizes)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "92d6a9e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "#At least 10 samples\n",
    "np.count_nonzero(np.array(GODataSizes)>20)\n",
    "GOsWithSufficientData = topGOs[np.array(GODataSizes)>20]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "c18c075f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "101\n"
     ]
    }
   ],
   "source": [
    "#Reduce to relevant GOs (see Data Exploration.ipynb)\n",
    "\n",
    "finalGOSet =[]\n",
    "for g in GOsWithSufficientData:\n",
    "    if g in relevantNodes:\n",
    "        finalGOSet.append(g)\n",
    "\n",
    "print(len(finalGOSet))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "99f6a6ca",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['GO:0008152' 'GO:0034655' 'GO:0072523' 'GO:0044270' 'GO:0006753'\n",
      " 'GO:1901292' 'GO:0044237' 'GO:1901360' 'GO:0008150' 'GO:1901564'\n",
      " 'GO:1901565' 'GO:0009117' 'GO:0006139' 'GO:0044281' 'GO:0046496'\n",
      " 'GO:0019362' 'GO:0046483' 'GO:0055086' 'GO:0044248' 'GO:0019439'\n",
      " 'GO:0019637' 'GO:0006807' 'GO:0019677' 'GO:1901361' 'GO:0006163'\n",
      " 'GO:0046700' 'GO:0009987' 'GO:0006725' 'GO:0006796' 'GO:0034641'\n",
      " 'GO:0072521' 'GO:0071704' 'GO:0019364' 'GO:1901575' 'GO:0072526'\n",
      " 'GO:0046434' 'GO:0009166' 'GO:0072524' 'GO:0006195' 'GO:0009056'\n",
      " 'GO:0044238' 'GO:0006793' 'GO:0019674']\n",
      "[[0 0 0 0 0 0 0 0 0 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 1 0 0 0 0 0 0\n",
      "  1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 0 0 0 0 1 1 0 1 0 1 0 0 0 0 0 0 0 0 0\n",
      "  0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 0 0 0 0 0 1 0 0]]\n",
      "101\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/manuel/miniconda3/envs/tf/lib/python3.9/site-packages/sklearn/preprocessing/_label.py:895: UserWarning: unknown class(es) ['GO:0006163', 'GO:0006195', 'GO:0006753', 'GO:0006796', 'GO:0008150', 'GO:0009117', 'GO:0009166', 'GO:0019362', 'GO:0019364', 'GO:0019439', 'GO:0019637', 'GO:0019674', 'GO:0019677', 'GO:0034655', 'GO:0044248', 'GO:0044270', 'GO:0046434', 'GO:0046496', 'GO:0046700', 'GO:0055086', 'GO:0072521', 'GO:0072523', 'GO:0072524', 'GO:0072526', 'GO:1901292', 'GO:1901361', 'GO:1901564', 'GO:1901565', 'GO:1901575'] will be ignored\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "from sklearn.preprocessing import MultiLabelBinarizer\n",
    "import pickle\n",
    "\n",
    "\n",
    "mlb = MultiLabelBinarizer()\n",
    "mlb.fit([finalGOSet])\n",
    "\n",
    "dftest=df.loc[testID]\n",
    "indices = dftest[\"term\"].to_numpy()\n",
    "print(indices)\n",
    "print(mlb.transform([indices]))\n",
    "print(len(mlb.classes_))\n",
    "\n",
    "with open(os.path.join(DATA_PATH,'MLB_'+SO+'.pkl'), 'wb') as f:\n",
    "    pickle.dump(mlb, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "afe3c8b2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Not found GOs: 0 (set to 0)\n"
     ]
    }
   ],
   "source": [
    "labelWeightsCorr=[]\n",
    "\n",
    "notFound=0\n",
    "for go in mlb.classes_:\n",
    "    if go in allIndices:\n",
    "        labelWeightsCorr.append(dfIa.loc[go].to_numpy()[0])\n",
    "    else:\n",
    "        notFound += 1\n",
    "        labelWeightsCorr.append(0)\n",
    "\n",
    "print(\"Not found GOs: {} (set to 0)\".format(notFound))\n",
    "labelWeightsCorr=np.array(labelWeightsCorr)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "92019f8c",
   "metadata": {},
   "source": [
    "## Amino acids encoding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "b747477d",
   "metadata": {},
   "outputs": [],
   "source": [
    "aminos_list = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'Y', 'X']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "6e42462e",
   "metadata": {},
   "outputs": [],
   "source": [
    "aa_dict = {'A': 1, 'B':24, 'C': 2, 'D': 3, 'E': 4, 'F': 5, 'G': 6, 'H': 7, 'I': 8, 'K': 9, 'L': 10, 'M': 11, 'N': 12, 'O': 21, 'P': 13, 'Q': 14, 'R': 15, 'S': 16, 'T': 17, 'U': 22, 'V': 18, 'W': 19, 'Y': 20, 'X':30, 'Z':23}"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "bfd823ba",
   "metadata": {},
   "source": [
    "## Build Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "431f9df8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The max. length of the sequences is 35375\n"
     ]
    }
   ],
   "source": [
    "seqLengths = [len(seq) for seq in sequences]\n",
    "maxLen = max(seqLengths)\n",
    "print(\"The max. length of the sequences is {}\".format(maxLen))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "e581e06c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['A0A009IHW8' 'A0A021WW32' 'A0A023FFD0' ... 'X5L1L5' 'X5L565' 'X5M5N0']\n"
     ]
    }
   ],
   "source": [
    "\n",
    "dfAll=pd.read_csv(os.path.join(DATA_PATH, \"Train/train_terms.tsv\"), sep='\\t')\n",
    "\n",
    "soEntries = dfAll.loc[dfAll[\"aspect\"]==SO]\n",
    "soEntryIds = soEntries[\"EntryID\"].unique()\n",
    "\n",
    "# print(len(seqEntries))\n",
    "print(soEntryIds)\n",
    "\n",
    "# SoSequences = []\n",
    "# for entry in soEntryIds:\n",
    "#     SoSequences.append(sequences[ids.index(entry)])\n",
    "\n",
    "# print(len(SoSequences))\n",
    "dfAll.set_index(\"EntryID\", inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "44f722e8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "99572\n"
     ]
    }
   ],
   "source": [
    "import warnings\n",
    "from itertools import product\n",
    "from tqdm import tqdm\n",
    "import os\n",
    "\n",
    "TRAIN_VAL_SPLIT = 0.7\n",
    "k = 3\n",
    "n=100\n",
    "maxFreqVectorLen = 5000//n\n",
    "\n",
    "allAA = list(aa_dict.keys())\n",
    "allAA.sort()\n",
    "allCombinations= list(product(*(allAA for i in range(k))))\n",
    "allCombinations=np.array([''.join(el) for el in allCombinations])\n",
    "\n",
    "positionDict = dict(zip(allCombinations, np.arange(0,allCombinations.size).T))\n",
    "\n",
    "xSize = allCombinations.shape[0]\n",
    "\n",
    "#Use numpy vectorize to speed up the mapping (hopefully)\n",
    "mapping = lambda x: aa_dict[x]\n",
    "vectMapping = np.vectorize(mapping)\n",
    "\n",
    "# Shuffle the data\n",
    "import random\n",
    "random.seed(516213)\n",
    "c = list(zip(sequences, ids))\n",
    "random.shuffle(c)\n",
    "sequencesShuffle, idsShuffle = zip(*c)\n",
    "\n",
    "\n",
    "#Train Validation Split\n",
    "split = int(np.floor(len(sequencesShuffle)*TRAIN_VAL_SPLIT))\n",
    "print(split)\n",
    "trainSeq = sequencesShuffle[0:split]\n",
    "valSeq = sequencesShuffle[split+1:]\n",
    "trainIds = idsShuffle[0:split]\n",
    "valIds = idsShuffle[split+1:]\n",
    "\n",
    "\n",
    "def generator():\n",
    "  for i,seq in enumerate(trainSeq):\n",
    "      entryId = trainIds[i]\n",
    "      \n",
    "      if entryId in soEntryIds:\n",
    "        labelData = df.loc[entryId]\n",
    "        # indices = labelData[\"termToken\"].to_numpy()\n",
    "        indices = labelData[\"term\"].to_numpy()\n",
    "      else:\n",
    "        indices=[]\n",
    "\n",
    "      with warnings.catch_warnings():\n",
    "          #supress the warnings for unknown classes\n",
    "          warnings.simplefilter(\"ignore\")\n",
    "          y = mlb.transform([indices])\n",
    "      \n",
    "      seqLen= len(seq)\n",
    "      freqVectorSeq=[]\n",
    "      for subIdx in np.arange(0, seqLen//n+1):\n",
    "        subSeq = seq[subIdx*n:(subIdx+1)*n]\n",
    "        if len(subSeq)<3:\n",
    "          continue\n",
    "        kmers = [subSeq[i:i+k] if i < len(subSeq)-(k-1) else 0 for i,el in enumerate(subSeq)]\n",
    "        kmers = kmers[0:-(k-1)]\n",
    "        kmers = [str(el) for el in kmers]\n",
    "        values, counts = np.unique(kmers, return_counts=True)\n",
    "        freqVector=np.zeros(allCombinations.shape)\n",
    "        for j,v in enumerate(values):\n",
    "          freqVector[positionDict[v]] = counts[j]\n",
    "        freqVectorSeq.append(freqVector)\n",
    "\n",
    "      #Pad or truncate the frequency vector sequence\n",
    "      if len(freqVectorSeq)>maxFreqVectorLen:\n",
    "          freqVectorSeq=np.array(freqVectorSeq[0:maxFreqVectorLen])\n",
    "      elif len(freqVectorSeq)<maxFreqVectorLen:\n",
    "          diffLen = maxFreqVectorLen - len(freqVectorSeq)\n",
    "          padding = np.zeros((diffLen,xSize))\n",
    "          freqVectorSeq = np.append(np.array(freqVectorSeq), padding, axis=0)\n",
    "      else:\n",
    "         freqVectorSeq= np.array(freqVectorSeq)\n",
    "      yield (freqVectorSeq,y[0])\n",
    "\n",
    "\n",
    "def generatorVal():\n",
    "  for i,seq in enumerate(valSeq):\n",
    "      entryId = valIds[i]\n",
    "      if entryId in soEntryIds:\n",
    "        labelData = df.loc[entryId]\n",
    "        # indices = labelData[\"termToken\"].to_numpy()\n",
    "        indices = labelData[\"term\"].to_numpy()\n",
    "      else:\n",
    "        indices=[]\n",
    "\n",
    "      with warnings.catch_warnings():\n",
    "          #supress the warnings for unknown classes\n",
    "          warnings.simplefilter(\"ignore\")\n",
    "          y = mlb.transform([indices])\n",
    "      \n",
    "      seqLen= len(seq)\n",
    "      freqVectorSeq=[]\n",
    "      for subIdx in np.arange(0, seqLen//n+1):\n",
    "        subSeq = seq[subIdx*n:(subIdx+1)*n]\n",
    "        if len(subSeq)<3:\n",
    "          continue\n",
    "        kmers = [subSeq[i:i+k] if i < len(subSeq)-(k-1) else 0 for i,el in enumerate(subSeq)]\n",
    "        kmers = kmers[0:-(k-1)]\n",
    "        kmers = [str(el) for el in kmers]\n",
    "        values, counts = np.unique(kmers, return_counts=True)\n",
    "        freqVector=np.zeros(allCombinations.shape)\n",
    "        for j,v in enumerate(values):\n",
    "          freqVector[positionDict[v]] = counts[j]\n",
    "        freqVectorSeq.append(freqVector)\n",
    "        \n",
    "      #Pad or truncate the frequency vector sequence\n",
    "      if len(freqVectorSeq)>maxFreqVectorLen:\n",
    "          freqVectorSeq=np.array(freqVectorSeq[0:maxFreqVectorLen])\n",
    "      elif len(freqVectorSeq)<maxFreqVectorLen:\n",
    "          diffLen = maxFreqVectorLen - len(freqVectorSeq)\n",
    "          padding = np.zeros((diffLen,xSize))\n",
    "          freqVectorSeq = np.append(np.array(freqVectorSeq), padding, axis=0)\n",
    "      else:\n",
    "         freqVectorSeq= np.array(freqVectorSeq)\n",
    "      yield (freqVectorSeq,y[0])\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "2727338a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The first sample: \n",
      "(50, 15625)\n",
      "[[0. 0. 0. ... 0. 0. 0.]\n",
      " [0. 0. 0. ... 0. 0. 0.]\n",
      " [0. 0. 0. ... 0. 0. 0.]\n",
      " ...\n",
      " [0. 0. 0. ... 0. 0. 0.]\n",
      " [0. 0. 0. ... 0. 0. 0.]\n",
      " [0. 0. 0. ... 0. 0. 0.]]\n",
      "The first sample has 0 classes\n",
      "519\n"
     ]
    }
   ],
   "source": [
    "g = generator()\n",
    "test = next(g)\n",
    "print(\"The first sample: \\n{}\\n{}\".format(test[0].shape, test[0]))\n",
    "print(\"The first sample has {} classes\".format(np.count_nonzero(test[1])))\n",
    "print(np.count_nonzero(test[0]))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "10c4a51a",
   "metadata": {},
   "source": [
    "## Tensorflow Classification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63f0d2e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "dataset = tf.data.Dataset.from_generator(generator, output_signature=(\n",
    "         tf.TensorSpec(shape=(maxFreqVectorLen,xSize), dtype=tf.int32),\n",
    "         tf.TensorSpec(shape=(len(mlb.classes_),), dtype=tf.int32)))\n",
    "print(list(dataset.take(1)))\n",
    "\n",
    "datasetVal = tf.data.Dataset.from_generator(generatorVal, output_signature=(\n",
    "         tf.TensorSpec(shape=(maxFreqVectorLen,xSize), dtype=tf.int32),\n",
    "         tf.TensorSpec(shape=(len(mlb.classes_),), dtype=tf.int32)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0a554f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext tensorboard\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73140569",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from tensorflow.keras import layers\n",
    "\n",
    "VOCAB_SIZE=len(aa_dict)\n",
    "EMBED_DIM=500\n",
    "OUT_SEQ_LENGTH=10\n",
    "\n",
    "def createRnnModel():\n",
    "    inputs = tf.keras.Input(shape=(maxFreqVectorLen, xSize,))\n",
    "    # x = tf.keras.layers.Masking(0)(inputs)\n",
    "    # x=layers.Embedding(VOCAB_SIZE, EMBED_DIM, name=\"embedding\")(inputs)\n",
    "\n",
    "    # x = layers.RepeatVector(OUT_SEQ_LENGTH)(x)\n",
    "    x=layers.BatchNormalization()(inputs)\n",
    "    x = layers.Bidirectional(layers.LSTM(10, return_sequences=True))(x)\n",
    "    x=layers.BatchNormalization()(x)\n",
    "    x = layers.Bidirectional(layers.LSTM(10))(x)\n",
    "    x=layers.BatchNormalization()(x)\n",
    "    # x = layers.LSTM(64)(x)\n",
    "    x = layers.Dense(128)(x)\n",
    "    x = layers.LeakyReLU()(x)\n",
    "    x = layers.Dense(128)(x)\n",
    "    outputs=layers.Dense(len(mlb.classes_), activation=tf.keras.activations.sigmoid)(x)\n",
    "    # outputs=layers.Softmax()(x)\n",
    "\n",
    "    return tf.keras.Model(inputs=inputs, outputs=outputs, name=\"embedRnnKmerSeqModel\")\n",
    "\n",
    "# model = createRnnModel()\n",
    "\n",
    "# model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23b1feaa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras import layers\n",
    "\n",
    "VOCAB_SIZE=len(aa_dict)\n",
    "EMBED_DIM=10\n",
    "\n",
    "def createModel():\n",
    "    inputs = tf.keras.Input(shape=(maxFreqVectorLen, xSize,))\n",
    "    # x=layers.Embedding(VOCAB_SIZE, EMBED_DIM, name=\"embedding\")(inputs)\n",
    "\n",
    "    x=layers.Conv2D(8, 5, strides=(1,2))(tf.expand_dims(inputs,3))\n",
    "    x = layers.BatchNormalization()(x)\n",
    "    x=layers.LeakyReLU()(x)\n",
    "\n",
    "    x=layers.Conv2D(8, 5, strides=(1,2))(x)\n",
    "    x = layers.BatchNormalization()(x)\n",
    "    x=layers.LeakyReLU()(x)\n",
    "\n",
    "    x=layers.Conv2D(8, 5, strides=(1,2))(x)\n",
    "    x = layers.BatchNormalization()(x)\n",
    "    x=layers.LeakyReLU()(x)\n",
    "\n",
    "    x=layers.Conv2D(8, 5, strides=(1,2))(x)\n",
    "    x = layers.BatchNormalization()(x)\n",
    "    x=layers.LeakyReLU()(x)\n",
    "\n",
    "    x=layers.Conv2D(8, 5, strides=(1,2))(x)\n",
    "    x = layers.BatchNormalization()(x)\n",
    "    x=layers.LeakyReLU()(x)\n",
    "\n",
    "    x=layers.Conv2D(8, 5)(x)\n",
    "    x = layers.BatchNormalization()(x)\n",
    "    x=layers.LeakyReLU()(x)\n",
    "    \n",
    "    x=layers.Flatten()(x)\n",
    "    x=layers.Dense(64)(x)\n",
    "    x=layers.LeakyReLU()(x)\n",
    "    x=layers.Dense(64)(x)\n",
    "    x=layers.LeakyReLU()(x)\n",
    "    x=layers.Dense(64)(x)\n",
    "    x=layers.LeakyReLU()(x)\n",
    "    outputs=layers.Dense(len(mlb.classes_), activation=tf.keras.activations.sigmoid)(x)\n",
    "    # outputs=layers.Softmax()(x)\n",
    "\n",
    "    return tf.keras.Model(inputs=inputs, outputs=outputs, name=\"embedConv2DKmerSeqModel\")\n",
    "\n",
    "model = createModel()\n",
    "\n",
    "model.summary()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a3ddc20",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "#Learning rate schedule\n",
    "initial_learning_rate = 0.001\n",
    "decaySteps=5000\n",
    "lr_schedule = tf.keras.optimizers.schedules.CosineDecayRestarts(initial_learning_rate, first_decay_steps=decaySteps,\n",
    "                                                                t_mul=2.0, m_mul=0.7)\n",
    "# lr_schedule = tf.keras.optimizers.schedules.CosineDecay(\n",
    "#     initial_learning_rate, decay_steps=decaySteps, alpha=0.01)\n",
    "# lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(\n",
    "#     initial_learning_rate,decay_steps=decaySteps,decay_rate=0.9,staircase=False)\n",
    "step = np.linspace(0,decaySteps*3)\n",
    "lr = lr_schedule(step)\n",
    "# plt.figure(figsize = (8,6))\n",
    "# plt.yscale(\"log\")\n",
    "# plt.plot(step, lr)\n",
    "# plt.ylim([0,max(plt.ylim())])\n",
    "# plt.xlabel('step')\n",
    "# _ = plt.ylabel('Learning Rate')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4760b66",
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9113c4a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import tensorflow_addons as tfa\n",
    "\n",
    "BATCH_SIZE=32\n",
    "LOG_INTERVAL=20\n",
    "epochs = 15\n",
    "saveModel=False\n",
    "\n",
    "\n",
    "log_dir = \"./logs/\"+model.name+\"/\" + datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\")+\"_\"+SO\n",
    "tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=1,\n",
    "                                                      write_graph=True, update_freq=5)\n",
    "\n",
    "\n",
    "summary_writer = tf.summary.create_file_writer(log_dir)\n",
    "\n",
    "# Instantiate an optimizer .\n",
    "# optimizer = tf.keras.optimizers.SGD(learning_rate=lr_schedule)\n",
    "# optimizer = tf.keras.optimizers.SGD(learning_rate=1e-3)\n",
    "optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)\n",
    "\n",
    "# Instantiate a loss function.\n",
    "# loss_fn = tf.keras.losses.BinaryCrossentropy()\n",
    "# loss_fn = WeightedBinaryCE(np.ones(len(mlb.classes_)))\n",
    "# loss_fn = WeightedBinaryCE(labelWeightsCorr)\n",
    "loss_fn = WeightedComboLoss(labelWeightsCorr, alpha=0.5, beta=0.5)\n",
    "\n",
    "train_acc_metric = WeightedAccuracy(classWeights=labelWeightsCorr)\n",
    "train_f1_metric = WeightedF1(classWeights=labelWeightsCorr, threshold=0.5)\n",
    "train_prec = WeightedPrecision(classWeights=labelWeightsCorr)\n",
    "train_rec = WeightedRecall(classWeights=labelWeightsCorr)\n",
    "\n",
    "val_acc_metric = WeightedAccuracy(classWeights=labelWeightsCorr)\n",
    "val_f1_metric = WeightedF1(classWeights=labelWeightsCorr, threshold=0.5)\n",
    "val_prec = WeightedPrecision(classWeights=labelWeightsCorr)\n",
    "val_rec = WeightedRecall(classWeights=labelWeightsCorr)\n",
    "\n",
    "batchedDataset = dataset.batch(BATCH_SIZE, drop_remainder=False)\n",
    "batchedDatasetVal = datasetVal.batch(BATCH_SIZE, drop_remainder=False)\n",
    "\n",
    "# batchedDataset = batchedDataset.cache(os.path.join(DATA_PATH, \"datasetCache\"+SO))\n",
    "# batchedDatasetVal = batchedDatasetVal.cache(os.path.join(DATA_PATH, \"datasetCacheVal\"+SO))\n",
    "\n",
    "@tf.function()\n",
    "def trainStep(x_batch_train, y_batch_train):\n",
    "    with tf.GradientTape() as tape:\n",
    "\n",
    "        probs = model(x_batch_train, training=True) \n",
    "        loss_value = loss_fn(y_batch_train, probs)\n",
    "\n",
    "    grads = tape.gradient(loss_value, model.trainable_weights)\n",
    "    #Gradient clipping\n",
    "    # grads = [tf.clip_by_norm(g, 2.0) for g in grads]\n",
    "\n",
    "    train_acc_metric.update_state(y_batch_train, probs)\n",
    "    train_f1_metric.update_state(y_batch_train, probs)\n",
    "    train_prec.update_state(y_batch_train, probs)\n",
    "    train_rec.update_state(y_batch_train, probs)\n",
    "\n",
    "    optimizer.apply_gradients(zip(grads, model.trainable_weights)) \n",
    "    return loss_value\n",
    "\n",
    "@tf.function()\n",
    "def valStep(x_batch_val, y_batch_val):\n",
    "    valProbs = model(x_batch_val, training=False)\n",
    "    # Update val metrics\n",
    "    val_acc_metric.update_state(y_batch_val, valProbs)\n",
    "    val_f1_metric.update_state(y_batch_val, valProbs)\n",
    "    val_prec.update_state(y_batch_val, valProbs)\n",
    "    val_rec.update_state(y_batch_val, valProbs)\n",
    "\n",
    "maxStep=0\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    print(\"\\nStart of epoch %d\" % (epoch+1,))\n",
    "\n",
    "    # Iterate over the batches of the dataset.\n",
    "    for step, (x_batch_train, y_batch_train) in enumerate(batchedDataset):\n",
    "\n",
    "        loss_value =trainStep(x_batch_train,y_batch_train)\n",
    "\n",
    "        # Log \n",
    "        if step % LOG_INTERVAL == 0:\n",
    "            template = 'Epoch {}/Step {}, Loss: {:.5f}, Accuracy: {:.5f}, F1: {:.4f}, Prec: {:.4f}, Rec: {:.4f} lr: {:.5f}'\n",
    "            print(template.format(epoch+1, step,loss_value.numpy(), \n",
    "                                    train_acc_metric.result(),train_f1_metric.result(),\n",
    "                                    train_prec.result(), train_rec.result(), optimizer.learning_rate.numpy()))\n",
    "            \n",
    "            with summary_writer.as_default():\n",
    "                tf.summary.scalar('loss', loss_value, step=maxStep*epoch+step)\n",
    "                tf.summary.scalar('accuracy', train_acc_metric.result(), step=maxStep*epoch+step)\n",
    "                tf.summary.scalar('f1', train_f1_metric.result(), step=maxStep*epoch+step)\n",
    "                tf.summary.scalar('prec', train_prec.result(), step=maxStep*epoch+step)\n",
    "                tf.summary.scalar('rec', train_rec.result(), step=maxStep*epoch+step)\n",
    "                tf.summary.scalar('learning rate', optimizer.learning_rate.numpy(), step=maxStep*epoch+step)\n",
    "                summary_writer.flush()\n",
    "\n",
    "    \n",
    "    train_acc_metric.reset_states()\n",
    "    train_f1_metric.reset_states()\n",
    "    train_prec.reset_states()\n",
    "    train_rec.reset_states()\n",
    "\n",
    "    maxStep=step\n",
    "\n",
    "    print(\"Epoch finished. Start validation\")\n",
    "    for x_batch_val, y_batch_val in batchedDatasetVal:\n",
    "        valStep(x_batch_val, y_batch_val)\n",
    "    val_acc = val_acc_metric.result()\n",
    "    val_acc_metric.reset_states()\n",
    "    val_f1 = val_f1_metric.result()\n",
    "    val_f1_metric.reset_states()\n",
    "    val_precision = val_prec.result()\n",
    "    val_prec.reset_states()\n",
    "    val_recall = val_rec.result()\n",
    "    val_rec.reset_states()\n",
    "    print(\"Validation acc: %.4f\" % (float(val_acc),))\n",
    "    print(\"Validation f1: %.4f\" % (float(val_f1),))\n",
    "    print(\"Validation precision: %.4f\" % (float(val_precision),))\n",
    "    print(\"Validation recall: %.4f\" % (float(val_recall),))\n",
    "    with summary_writer.as_default():\n",
    "        tf.summary.scalar('valAcc', float(val_acc), step=epoch)\n",
    "        tf.summary.scalar('valF1', float(val_f1), step=epoch)\n",
    "        tf.summary.scalar('valPrecision', float(val_precision), step=epoch)\n",
    "        tf.summary.scalar('valRecall', float(val_recall), step=epoch)\n",
    "        summary_writer.flush()\n",
    "    if saveModel:\n",
    "      model.save(os.path.join(DATA_PATH, \"model_\"+SO+\"_epoch_{}_valF1Score{:.3f}\".format(epoch, float(val_f1))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "374f7234",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# model.save(os.path.join(DATA_PATH, \"model_\"+SO+\"_epoch_{}_valf1Score{:.3f}\".format(epoch, float(val_f1))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cefc1f6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "probs= model.predict(tf.expand_dims(list(datasetVal.take(32))[10][0], 0))\n",
    "prediction= [1 if p > 0.5 else 0 for p in probs[0]]\n",
    "probabilities= probs[probs>0.5]\n",
    "# classes = np.argwhere(prediction)\n",
    "print(mlb.inverse_transform(np.array([prediction])))\n",
    "print(probabilities)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}