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

# + [markdown] id="xIl1ubvxyWdI" toc=true
# <h1>Table of Contents<span class="tocSkip"></span></h1>
# <div class="toc"><ul class="toc-item"></ul></div>

# + colab={"base_uri": "https://localhost:8080/", "height": 1000} id="M1tpfwM6_GeB" outputId="94edc658-2f7c-4052-e60b-6067a26d8b5e"
# !pip install lifelines
# !pip install scikit-plot

# + id="-S4tbupSz28Y"
import tensorflow as tf

from tensorflow import keras
from sklearn.model_selection import train_test_split
from keras import backend as K

from tqdm.keras import TqdmCallback
from tqdm.auto import tqdm
import random
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from keras.callbacks import CSVLogger

PREFIX = 'CNN-with-memory-'

csv_logger = CSVLogger(f'{PREFIX}.log', append=True, separator=';')

plt.style.use('default')
plt.style.use('seaborn-white')

models = tf.keras.models
layers = tf.keras.layers
preprocessing = tf.keras.preprocessing

tqdm.pandas()


# + id="kIaqre5vz3nq"
class LossWithMemoryCallback(tf.keras.callbacks.Callback):
    def __init__(self, kd_var, ki_var, ic_var, ec_var, ph_var, discount=0.1):
        self.kd_var = kd_var
        self.ki_var = ki_var
        self.ic_var = ic_var
        self.ec_var = ec_var
        self.ph_var = ph_var

        self.discount = discount

    def on_epoch_end(self, epoch, logs):
        K.set_value(self.kd_var, logs['Kd_loss'] * self.discount)
        K.set_value(self.ki_var, logs['Ki_loss'] * self.discount)
        K.set_value(self.ic_var, logs['IC50_loss'] * self.discount)
        K.set_value(self.ec_var, logs['EC50_loss'] * self.discount)
        K.set_value(self.ph_var, logs['pH_loss'] * self.discount)


def loss_wrapper(var):
    def mse_nan_with_memory(y_true, y_pred):
        masked_true = tf.where(tf.math.is_nan(y_true), tf.zeros_like(y_true), y_true)
        masked_pred = tf.where(tf.math.is_nan(y_true), tf.zeros_like(y_true), y_pred)
        is_empty = tf.equal(tf.size(masked_true), 0)

        if is_empty:
            error = var
            return error

        error = K.mean(K.square(masked_pred - masked_true), axis=-1)
        return error

    return mse_nan_with_memory


# + id="grzX2AuZyWdW"
def get_optimizer():
    optimizer = tf.keras.optimizers.Adadelta(learning_rate=1.0)
    return optimizer


def get_model(units=96, rate=0.3, loss_weights=None, kvars = {}):

    sm = layers.Input(shape=(100, ), name='D_Inp')
    aa = layers.Input(shape=(1000, ), name='T_Inp')

    emsm0 = layers.Embedding(53,
                             128,
                             trainable=True,
                             name='D_Emb',
                             mask_zero=True)(sm)
    emaa0 = layers.Embedding(22,
                             128,
                             trainable=True,
                             name='T_Emb',
                             mask_zero=True)(aa)

    cnvsm1 = layers.Conv1D(32, 3, name='D_L1')(emsm0)
    cnvaa1 = layers.Conv1D(32, 3, name='T_L1')(emaa0)

#     cnvsm2 = layers.Conv1D(64, 3, name='D_L2')(cnvsm1)
#     cnvaa2 = layers.Conv1D(64, 3, name='T_L2')(cnvaa1)

    cnvsm3 = layers.Conv1D(96, 3, name='D_L3')(cnvsm1)
    cnvaa3 = layers.Conv1D(96, 3, name='T_L3')(cnvaa1)

    gmpsm = layers.GlobalMaxPool1D(name='D_Gmp')(cnvsm3)
    gmpaa = layers.GlobalMaxPool1D(name='T_Gmp')(cnvaa3)

    C1 = layers.concatenate([gmpsm, gmpaa], axis=-1, name='C1')

    S1 = layers.Dense(512, activation='relu', name='S1')(C1)
    S1 = layers.Dropout(rate)(S1)

    S2 = layers.Dense(512, activation='relu', name='S2')(S1)
    S2 = layers.Dropout(rate)(S2)

    S3 = layers.Dense(512, activation='relu', name='S3')(S2)
    S3 = layers.Dropout(rate)(S3)

    Kd = layers.Dense(units, activation='relu', name='S1_Kd')(S3)
    Kd = layers.Dropout(rate)(Kd)

    Ki = layers.Dense(units, activation='relu', name='S1_Ki')(S3)
    Ki = layers.Dropout(rate)(Ki)

    IC50 = layers.Dense(units, activation='relu', name='S1_IC50')(S3)
    IC50 = layers.Dropout(rate)(IC50)

    EC50 = layers.Dense(units, activation='relu', name='S1_EC50')(S3)
    EC50 = layers.Dropout(rate)(EC50)

    IA = layers.Dense(units, activation='relu', name='S1_IA')(S3)
    IA = layers.Dropout(rate)(IA)

    pH = layers.Dense(units, activation='relu', name='S1_pH')(S3)
    pH = layers.Dropout(rate)(pH)

    out1 = layers.Dense(1, activation='linear', name='Kd')(Kd)
    out2 = layers.Dense(1, activation='linear', name='Ki')(Ki)
    out3 = layers.Dense(1, activation='linear', name='IC50')(IC50)
    out4 = layers.Dense(1, activation='linear', name='EC50')(EC50)
    out5 = layers.Dense(1, activation='sigmoid', name='IA')(IA)
    out6 = layers.Dense(1, activation='linear', name='pH')(pH)

    model = models.Model(inputs=[sm, aa],
                         outputs=[out1, out2, out3, out4, out5, out6])

    losses = {
        "Kd": loss_wrapper(kd_var),
        "Ki": loss_wrapper(ki_var ),
        "IC50": loss_wrapper(ic_var),
        "EC50": loss_wrapper(ec_var),
        "pH": loss_wrapper(ph_var),
        "IA": "binary_crossentropy",
    }


    metrics = {"IA": tf.keras.metrics.AUC()}

    model.compile(loss=losses, optimizer=get_optimizer(), metrics=metrics, loss_weights=loss_weights)
    model.summary()
    return model

# + colab={"base_uri": "https://localhost:8080/"} id="WBjkZsDq1TtY" outputId="7dc4a667-e40d-4d32-f846-dab1e3fbb824"
tf.keras.backend.clear_session()
np.random.seed(7)
tf.random.set_seed(7)

loss_weights = [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]

global kd_var, ki_var, ic_var, ec_var, ph_var
kd_var = K.variable(0.0)
ki_var = K.variable(0.0)
ic_var = K.variable(0.0)
ec_var = K.variable(0.0)
ph_var = K.variable(0.0)


lossCallback = LossWithMemoryCallback(kd_var, ki_var, ic_var, ec_var, ph_var)


model = get_model(rate=0.3, loss_weights=loss_weights)

# + colab={"base_uri": "https://localhost:8080/", "height": 367} id="QjJBESth1V-E" outputId="6752d73c-06b0-46f3-ff14-8f6694c72c25"
tf.keras.utils.plot_model(model, rankdir='LR',
                          show_shapes=True,
                          show_layer_activations=True)

# + id="_xK8JAXg7dS2"
CHARPROTSET = dict([('A', 1), ('G', 2), ('L', 3), ('M', 4), ('S', 5), ('T', 6),
                    ('E', 7), ('Q', 8), ('P', 9), ('F', 10), ('R', 11),
                    ('V', 12), ('D', 13), ('I', 14), ('N', 15), ('Y', 16),
                    ('H', 17), ('C', 18), ('K', 19), ('W', 20), ('X', 21)])

CHARCANSMISET = dict([(')', 1), ('(', 2), ('1', 3), ('C', 4), ('c', 5),
                      ('O', 6), ('2', 7), ('N', 8), ('=', 9), ('n', 10),
                      ('3', 11), ('-', 12), ('4', 13), ('F', 14), ('S', 15),
                      ('[', 16), (']', 17), ('l', 18), ('H', 19), ('s', 20),
                      ('#', 21), ('o', 22), ('5', 23), ('B', 24), ('r', 25),
                      ('+', 26), ('6', 27), ('P', 28), ('.', 29), ('I', 30),
                      ('7', 31), ('e', 32), ('i', 33), ('a', 34), ('8', 35),
                      ('K', 36), ('A', 37), ('9', 38), ('T', 39), ('g', 40),
                      ('R', 41), ('Z', 42), ('%', 43), ('0', 44), ('u', 45),
                      ('V', 46), ('b', 47), ('t', 48), ('L', 49), ('*', 50),
                      ('d', 51), ('W', 52)])


# + id="U3kEDO0i73Uj"
class Gen:
    def __init__(self,
                 data,
                 map_smiles,
                 map_aa,
                 shuffle=True,
                 test_only=False,
                 len_drug=100,
                 len_target=1000,
                 window=False):
        self.data = data
        self.map_smiles = map_smiles
        self.map_aa = map_aa
        self.shuffle = shuffle
        self.test_only = test_only
        self.len_drug = len_drug
        self.len_target = len_target
        self.size = self.data.shape[0]
        self.inds = list(range(self.size))
        if self.shuffle:
            random.shuffle(self.inds)

        self.window = window

        self.gen = self._get_inputs()

    def _get_inputs(self):
        seen = 0
        while seen < self.size:
            ind = self.inds[seen]
            sample = self.data.iloc[ind, :].values.tolist()
            sample[0] = self.map_smiles[sample[0]]
            sample[1] = self.map_aa[sample[1]]

            if self.window:
                ld = max(0, (len(sample[0]) - self.len_drug))
                lt = max(0, (len(sample[1]) - self.len_target))
                dstart = random.randint(0, ld)
                tstart = random.randint(0, lt)

                sample[0] = sample[0][dstart:dstart + self.len_drug]
                sample[1] = sample[1][tstart:dstart + self.len_target]

            yield sample
            seen += 1
            if seen == self.size:
                if self.shuffle:
                    random.shuffle(self.inds)
                seen = 0

    def get_batch(self, batch_size):
        while True:
            BATCH = []
            for _ in range(batch_size):
                sample = next(self.gen)
                for k, value in enumerate(sample):
                    if len(BATCH) < (k+1):
                        BATCH.append([])
                    BATCH[k].append(value)

            BATCH[0] = preprocessing.sequence.pad_sequences(BATCH[0], self.len_drug)
            BATCH[1] = preprocessing.sequence.pad_sequences(BATCH[1], self.len_target)
            
            for k in range(2, len(BATCH)):
                BATCH[k] = np.array(BATCH[k]).flatten()

            if not self.test_only:
                yield [BATCH[0],  BATCH[1]], [BATCH[k] for k in range(2, len(BATCH))]
            else:
                yield [BATCH[0],  BATCH[1]], [BATCH[k]*0 for k in range(2, len(BATCH))]


# + id="3tQtKlmk8TiV"
data = pd.read_csv("data_full_05_pH.zip", compression='zip')
order = [
    'smiles', 'target', 'p1Kd', 'p1Ki', 'p1IC50', 'p1EC50', 'is_active', 'pH'
]
data = data[order]

# + colab={"base_uri": "https://localhost:8080/", "height": 287} id="W38UR0acBU-i" outputId="c86cfca6-a7c8-4f5a-b4ba-bb99ab361dd3"
data = data.sample(frac=1, random_state = 7)
data.head()

# + colab={"base_uri": "https://localhost:8080/"} id="10ll6kWis2dT" outputId="0dd3d552-248a-4852-f676-25152848132c"
data.dropna().shape

# + colab={"base_uri": "https://localhost:8080/", "height": 81, "referenced_widgets": ["415da6e7e1454dcfaad83d5b5651a6b3", "1db48340efcb4a61a82672c3a542498a", "d98765fa55eb4ad9bb6ab73631bb9631", "ce935596b29b4abcad74cb7e09a33e00", "dea41680c43b4c92a022c01d76577e9d", "08abd662185143f89da9f9bf3fef2511", "2b2074bd630240cd9f2eda0ce38f147e", "d98d5e8fe59543f6b5f3823c78f1ea25", "58d38dd270ff42b597df46f665cf5feb", "2d26a4c295454845b04965e330ddcb49", "ca391d03802147c7a81533cfaf29a26e", "149745edeb31456abaef31673f5606c7", "0a1f1b1a2de94c89b8bec693a1389ba2", "3fdd379e80c5479dbede3d23367c3bae", "eba043f3e7b64a51a30c35d2dff4a310", "<KEY>", "c0c9fd436d0a4c4a98a42472134ccec1", "<KEY>", "35eee7e6d9604fee9de39784380d718e", "<KEY>", "e3c1bb16c0c44bd0a558ecce23aded23", "9a7652fd49a04b2fba24fd7899e6de7d"]} id="kZXswzVO8nmp" outputId="ac815ed4-53d3-45bb-8bc0-f5ca9307ef38"
SMILES = {}
for smiles in tqdm(data['smiles'].unique()):
    SMILES[smiles] = [CHARCANSMISET[s] for s in smiles]

AA = {}
for aa in tqdm(data['target'].unique()):
    AA[aa] = [CHARPROTSET[a.upper()] for a in aa]

# + colab={"base_uri": "https://localhost:8080/"} id="_S9hwivl8IiM" outputId="d550c40a-9bc6-443b-c5cd-ff535751eafd"
X_train, X_test = train_test_split(data,
                                   test_size=0.1,
                                   shuffle=True,
                                   random_state=7,
                                   stratify=data['is_active'])

X_train, X_valid = train_test_split(X_train,
                                    test_size=0.1,
                                    shuffle=True,
                                    random_state=7,
                                    stratify=X_train['is_active'])

X_train.shape[0], X_test.shape[0], X_valid.shape[0]

# + colab={"base_uri": "https://localhost:8080/", "height": 270} id="J8CbVClA8_Rf" outputId="d596c2c7-3ab6-4320-a1b4-88d260469ab0"
X_train.head()

# + id="5BhQxIns851d"
# It is stochastic on Windows ?
batch_size = 1024

trg = Gen(X_train, SMILES, AA)
trg = trg.get_batch(batch_size)

vag = Gen(X_valid, SMILES, AA)
vag = vag.get_batch(batch_size)

# + id="SnTU-fTirAbi"
# for batch in trg:
#     break

# batch

# + id="y6Jrd8xC9oBN"
steps_per_epoch = X_train.shape[0] // batch_size
valid_steps = X_valid.shape[0] // batch_size

# + id="o_lLDAXc9qls"
filepath = "CNN_basic_with_memory_{epoch:02d}-{val_loss:.2f}.h5"
checkpoint = tf.keras.callbacks.ModelCheckpoint(filepath,
                                                monitor='val_loss',
                                                verbose=1,
                                                save_best_only=False,
                                                mode='auto',
                                                save_weights_only=True)

# + colab={"base_uri": "https://localhost:8080/", "height": 1000, "referenced_widgets": ["043485b65027406eb818edc1f30453bc", "211f802a01c9497ca01b060f9ea2eaf1", "1732871292ae4f72a8ee02cf2507843e", "<KEY>", "7ed12a85ed674c47af1282a908a165ac", "<KEY>", "<KEY>", "<KEY>", "0a9d2f4a12b746f58c49e04b16d1ac52", "<KEY>", "2825f87870a94b1c8060945a006a4c78", "<KEY>", "<KEY>", "4c4adba52dff40f496ffd7468a9da7d2", "112433a0cc604458a32c00c541e3ac60", "<KEY>", "<KEY>", "<KEY>", "e1badb9cf2984e4db2db0de8d6cd3a49", "<KEY>", "156790ee47784da0b85557980217d960", "07e612eed3dd47bf9fbbe8053f1b7752"]} id="fT_ws69y9sTu" outputId="96787752-117e-4190-b5df-7af344f1f294"
history = model.fit(trg,
                    validation_data=vag,
                    steps_per_epoch=steps_per_epoch,
                    validation_steps=valid_steps,
                    verbose=0,
                    callbacks=[TqdmCallback(), checkpoint, csv_logger, lossCallback],
                    epochs=50)

# + id="vqtrfFk2Alzv"
model.load_weights('CNN_basic_with_memory_50-4.20.h5')
model.compile(optimizer=get_optimizer())
model.save("cnn_with_memory.hdf5")

# + id="y9NHd3-MItR1"
# # !rm *.h5 -r

# + id="wJXxv-nT-Uot"
# history.history

# + colab={"base_uri": "https://localhost:8080/", "height": 308} id="-BM529G697Fp" outputId="dc71b00a-2350-49dd-da76-d4c7080e7b0d"
plt.plot(history.history['loss'], label='train')
plt.plot(history.history['val_loss'], label='valid')
plt.xlabel('Epoch')
plt.title('Loss on train-valid subsets')
plt.legend()


# + id="dg3iIlQL97cA"
def get_batch_size(S):
    mbs = 1
    for i in range(1, min(64, S)):
        if S % i == 0:
            mbs = i
    assert S % mbs == 0

    return mbs


# + colab={"base_uri": "https://localhost:8080/"} id="dkGC2ZUK987w" outputId="6d1acb57-311a-440e-ac8b-78517b64fe4c"
mbs = get_batch_size(X_test.shape[0])
mbs

# + id="2rQ2J_Gl9-C4"
teg = Gen(X_test, SMILES, AA, shuffle=False, test_only=True)
teg = teg.get_batch(mbs)

# + colab={"base_uri": "https://localhost:8080/"} id="QzHF0cZk9_NG" outputId="c12e0f5b-4923-424c-f186-c9e07da811b3"
prediction = model.predict(teg, steps=X_test.shape[0] // mbs, verbose=1)

# + id="W5p6qCEb-Ac_"
# !from sklearn.metrics import mean_squared_error
from lifelines.utils import concordance_index
from scipy import stats


# + id="hndwzcfe-Bdm"
def get_scores(y_true, y_pred):
    mse = np.round(mean_squared_error(y_true, y_pred), 3)
    rmse = np.round(mse**0.5, 3)
    ci = np.round(concordance_index(y_true, y_pred), 3)
    pearson = np.round(stats.pearsonr(y_true, y_pred)[0], 3)
    spearman = np.round(stats.spearmanr(y_true, y_pred)[0], 3)

    res = f"rmse={rmse}, mse={mse},\npearson={pearson}, spearman={spearman},\nci={ci}"

    return res


# + colab={"base_uri": "https://localhost:8080/", "height": 1000} id="fIJo-WcN-Cl8" outputId="cf9b4079-17f1-4560-9c5f-2e248f8c103e"
for k, col in enumerate(
    ['p1Kd', 'p1Ki', 'p1IC50', 'p1EC50', 'is_active', 'pH']):
    plt.scatter(X_test[col], prediction[k], alpha=0.7, c='k')
    plt.xlabel('true')
    plt.ylabel('predicted')
    y_true = X_test[col][X_test[col].notna()]
    y_pred = prediction[k][X_test[col].notna()].ravel()
    plt.title(col + ":\n" + get_scores(y_true, y_pred))
    plt.show()  # 74.6

# + id="jAA6lozS-DsH"
import scikitplot as skplt

# + colab={"base_uri": "https://localhost:8080/", "height": 346} id="cropd2ga-EvM" outputId="60d8f76f-4be2-4f38-f73c-9d37fcd3efa9"
p = prediction[-2].ravel().tolist()
probas = np.zeros((len(p), 2))
probas[:, 1] = p
probas[:, 0] = 1
probas[:, 0] = probas[:, 0] - p
skplt.metrics.plot_roc_curve(X_test['is_active'].values.ravel().tolist(),
                             probas)
plt.show()

# + colab={"base_uri": "https://localhost:8080/", "height": 425} id="LnSdXhk0-FEk" outputId="b9357a31-3e3a-492f-ad3a-475057eacc63"
plt.hist(prediction[-2].ravel(), bins=32, edgecolor='w', color='k', alpha=0.7);

# + colab={"base_uri": "https://localhost:8080/"} id="DoEMgyBMrAbt" outputId="e82a4bf6-17b3-4c15-a942-26710da5e92a"
# !zip data.zip *.h5 -r

StarcoderdataJupyterScriptsDedupFiltered