# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + [markdown] id="view-in-github" colab_type="text"
# <a href="https://colab.research.google.com/github/csy99/dna-nn-theory/blob/master/kmer.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>

# + id="gad62cNa5Waf" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 125} outputId="4151c60e-dd7d-4857-e32f-1a1553ae7362"
# !curl https://cog-uk.s3.climb.ac.uk/2020-05-08/cog_2020-05-08_alignment.fasta -o cog_2020-05-08_alignment.fasta
# !curl https://cog-uk.s3.climb.ac.uk/2020-05-08/cog_2020-05-08_metadata.csv -o metadata.csv

# + id="Axh0xwMs5kSk" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 145} outputId="8b010a89-2759-4e0b-e254-30ac9c6365a8"
# !pip install biopython

# + id="1woDH1ec5j7g" colab_type="code" colab={}
import tensorflow as tf
from tensorflow import keras
import numpy as np
import pandas as pd
from Bio import SeqIO

# + id="alqpS2m-1sMq" colab_type="code" colab={}
from itertools import product
import re

# + id="wvygdm9I1iBE" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 1000} outputId="a1267d81-2241-4294-d86e-b6c5aa8f7818"
word_size = 3
words = [''.join(c) for c in product('ATCGN', repeat=word_size)]
word_to_idx = {word: i for i, word in enumerate(words)}
word_to_idx

# + id="Xht0R8Lk2O5Q" colab_type="code" colab={}
vocab_size = 5 ** word_size
word_to_idx_func = np.vectorize(lambda word: word_to_idx[word], otypes=[np.int8])

# + id="E0QlD0mJ4cQW" colab_type="code" colab={}


# + id="YPWVVClY5kiC" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 53} outputId="22ce0a5a-92f5-4ce2-82fa-5b60b68f5017"
with open('cog_2020-05-08_alignment.fasta', 'r') as f:
  # records = []
  # for seq in SeqIO.parse(f, 'fasta'):
    # records.append(" ".join([c for c in seq]))
  # records = np.array(records)
  records = np.array([record.seq._data.upper() for record in SeqIO.parse(f, 'fasta')])
seq_len = len(records[0])
print('Number of sequences :', len(records))
print('Length of sequences :', seq_len)

# + [markdown] id="5-29Vymb6880" colab_type="text"
# This process took 2 mins. 

# + id="4DdSLyd75kpY" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 143} outputId="14558f40-3bdb-4063-ecee-0cf2ca4f4ecd"
records[0]

# + id="OHEJfDfB5kwR" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 35} outputId="6e59e73d-6d53-42d1-8122-3848db32e6b1"
covid_df = pd.read_csv("metadata.csv")
label = covid_df["adm1"]
label_names = covid_df["adm1"].unique()
label_names

# + id="bJcxjzJfEPkz" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 269} outputId="a7b5e7e6-16ab-45f6-ad87-41466065556b"
print(covid_df["adm1"].value_counts())
print(covid_df[:2000]["adm1"].value_counts())
print(covid_df[-1000:]["adm1"].value_counts())

# + id="0ViB0yw_vKxt" colab_type="code" colab={}
# def getWeight(label_names):
#   return 1000.0/covid_df["adm1"].value_counts()[label_names]

# weights = []
# for name in label_names:
#   weight = getWeight(name)
#   weights.append(weight)

# weights

# + id="W-Jy2PPc5k3O" colab_type="code" colab={}
for i in range(len(label_names)):
  label = label.replace(label_names[i], i)

# + id="l0DJ51Yl5k-B" colab_type="code" colab={}
# should change the training/validation set size later
TRAIN_SIZE = 2000
Test_SIZE = 1000
xtrain = records[:TRAIN_SIZE]
ytrain = label[:TRAIN_SIZE]
xtest = records[-Test_SIZE:]
ytest = label[-Test_SIZE:]

# + id="O7meRbgwVqau" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 35} outputId="2dfde296-256e-4022-e013-e2809623934f"
xtrain = xtrain.reshape(xtrain.shape[0], 1)
xtrain.shape

# + id="2g3_WXDdLEnT" colab_type="code" colab={}
xtest = xtest.reshape(xtest.shape[0], 1)

# + id="yHbpa4SPVO3t" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 163} outputId="7144c399-12c0-4134-c55b-3a236aa83ac7"
xtrain[:6]

# + id="skCianA4HVfX" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 197} outputId="df5e3be4-ac00-4b4d-e2a9-013bbe481b75"
for x in xtrain[:10]:
  print(len(x[0]))


# + id="wgJX8TezZvKp" colab_type="code" colab={}
def x_to_encoding(x, word_size=3, region_size=2, expand=True, truncate=500):
  x = tf.strings.regex_replace(x, "-", "")
  x = tf.strings.regex_replace(x, 'K', 'N')
  x = tf.strings.regex_replace(x, 'Y', 'N')
  x = tf.strings.regex_replace(x, 'R', 'N')
  x = tf.strings.regex_replace(x, 'W', 'N')
  x = tf.strings.regex_replace(x, 'M', 'N')
  x = tf.strings.regex_replace(x, 'S', 'N')
  x = x.numpy()[0].decode("utf-8") 
  # truncate for 2 reasons: data cleaning and reduce parameters
  x = x[:truncate]
  # print("len of x after preprocess", len(x))
  arr_x = [c for c in x]
  words = tf.strings.ngrams(arr_x, ngram_width=word_size, separator='').numpy()
  idx = word_to_idx_func(words.astype('U'))
  encoding = np.zeros((len(idx), vocab_size))
  encoding[range(len(idx)), idx] = 1
  # print(encoding.shape)
  res = np.hstack([np.roll(encoding, -shift, axis=0) for shift in range(region_size+1)])
  # print(res.shape)
  if expand:
    return np.expand_dims(res, axis=-1) 
  return res


# + id="6OZ9y9JlTJoP" colab_type="code" colab={}
region_size = 0
truncated_len = 100
xtrain2 = np.apply_along_axis(x_to_encoding, axis=1, \
                arr=xtrain, region_size=region_size, truncate=truncated_len)

# + id="wGJij2C1Iq82" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 35} outputId="5557e696-c7bf-4336-f623-3bdd7183d6b2"
xtrain2.shape

# + id="yS-VJ57wKYrh" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 35} outputId="c3e1909f-92de-4478-c2de-32fe7efad4f5"
import gc
gc.collect()

# + id="PMWJDY3I7c7c" colab_type="code" colab={}
from tensorflow.keras import layers
shape = xtrain2.shape[1:]
model1 = keras.Sequential([
  layers.Conv2D(256, 3, activation="relu", input_shape=shape), 
  layers.MaxPooling2D(), 
  layers.Dropout(0.3), 
  layers.Flatten(), 
  layers.Dense(128, activation="relu"),
  layers.Dropout(0.5), 
  layers.Dense(4, activation="softmax"),
])

# + id="r_BFgFM_7dCG" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 557} outputId="218f9cd5-2cef-493b-fa47-77cb905b7dc2"
model1.compile(optimizer='adam', loss=tf.keras.losses.SparseCategoricalCrossentropy(), metrics=["accuracy"])
es_cb = tf.keras.callbacks.EarlyStopping(monitor='loss', patience=5, restore_best_weights=True)
hist1 = model1.fit(xtrain2, ytrain, epochs=15, callbacks=[es_cb])

# + id="QMKP62qtLTVv" colab_type="code" colab={}
xtest2 = np.apply_along_axis(x_to_encoding, axis=1, \
                arr=xtest, region_size=region_size, truncate=truncated_len)

# + id="BRUCVYVhE4Ao" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 53} outputId="38c3583e-6975-4426-b925-10d2622720be"
model1.evaluate(xtest2, ytest)

# + id="F4awV0eqLApH" colab_type="code" colab={}


# + id="o9GpcPxcc6tb" colab_type="code" colab={}
def evaluateModel(model, ds):
  ds_pred = np.argmax(model.predict(ds), axis=-1)
  print("prediction result shape", ds_pred.shape)
  ds_true = []
  label_ds = ds.map(lambda x, y: y)
  for item in label_ds:
    ds_true.append(item.numpy())
  ds_true = np.reshape(np.array(ds_true), (-1,)) 
  print("true label shape", ds_true.shape)
  return ds_pred, ds_true


# + id="X-YtFtFAfZgb" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 53} outputId="ddc31833-cf4f-4707-9283-a1c401beee02"
seq_ds_train_pred, seq_ds_train_true = evaluateModel(model1, seq_ds_train)

# + id="3jiuF3VNLPPh" colab_type="code" colab={}


# + id="ZOhAm9sMzD_L" colab_type="code" colab={}
import tensorflow_addons as tfa
confusion_mat = tfa.metrics.MultiLabelConfusionMatrix(num_classes=4)
confusion_mat.update_state(tf.one_hot(seq_ds_train_true, depth=4), tf.one_hot(seq_ds_train_pred, depth=4))

# + [markdown] id="-0aY4OuS0x-t" colab_type="text"
# true negatives for class i in M(0,0)  
# false positives for class i in M(0,1)  
# false negatives for class i in M(1,0)  
# true positives for class i in M(1,1)  

# + id="-qiPvm-P0I-n" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 233} outputId="a7a92b43-180d-4a8f-af0d-e11cedc06a40"
# pd.set_option('display.float_format',lambda x : '%.2f' % x)
np.set_printoptions(suppress=True)
print('Confusion matrix:\n', confusion_mat.result().numpy())

# + id="-KwdcB73CIJq" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 89} outputId="77052c3e-786c-49ab-c3ed-b51055ea33d2"
tmp = covid_df[:2000]["adm1"].value_counts()
# tmp["UK-NIR"]
for label in label_names:
  print(label, tmp[label])

# + id="7SCnhQBB7d74" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 269} outputId="7e546e50-3076-4c8f-c7d8-9568ea4f915c"
# seq_ds_test_pred, seq_ds_test_true = evaluateModel(model1, seq_ds_test)
# confusion_mat = tfa.metrics.MultiLabelConfusionMatrix(num_classes=4)
# confusion_mat.update_state(tf.one_hot(seq_ds_test_pred, depth=4), tf.one_hot(seq_ds_test_true, depth=4))
# print('Confusion matrix:\n', confusion_mat.result().numpy())

# + id="uXE9vJVsfzSQ" colab_type="code" colab={}


# + id="5hQPOs-zfzdM" colab_type="code" colab={}


# + id="LJkzN84-fz-k" colab_type="code" colab={}


# + id="JsoxAWtuf0JR" colab_type="code" colab={}


# + id="Nh5TR6AT9gJ7" colab_type="code" colab={}


# + id="CmfIKCZohFoK" colab_type="code" colab={}


# + id="zL6Jho8ihFwN" colab_type="code" colab={}


# + id="SZaFovzFhF39" colab_type="code" colab={}


# + id="ywiiUKoOhF_Y" colab_type="code" colab={}


StarcoderdataJupyterScriptsDedupFiltered