# -*- coding: utf-8 -*-
"""chapter5.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1FZ4JinLGstOWo72o-UtE1Q0A_V8k6unA
"""

# Chapter5 Biophysical ML
# we will explore in depth the problem of predicting how
# small drug-like molecules bind to a protein of interest in the human body.

# Our goal therefore is to design learning algorithms that can effectively predict when a
# given molecule is going to interact with a given protein. How can we do this?

#

# Commented out IPython magic to ensure Python compatibility.
##setup tensorflow v1
# %tensorflow_version 1.x

## this will install anaconda and deepchem, will add path, execution will take sometime

!wget -c https://repo.anaconda.com/archive/Anaconda3-2019.10-Linux-x86_64.sh
!chmod +x Anaconda3-2019.10-Linux-x86_64.sh
!bash ./Anaconda3-2019.10-Linux-x86_64.sh -b -f -p /usr/local
!conda install -y -c deepchem -c rdkit -c conda-forge -c omnia deepchem-gpu=2.3.0
import sys
sys.path.append('/usr/local/lib/python3.7/site-packages/')

############## check deepchem installation by 
import deepchem as dc
dc.__version__ #should match with the installed

#need to reinstall deepchem bcz of an openmm error

!curl -Lo deepchem_installer.py https://raw.githubusercontent.com/deepchem/deepchem/master/scripts/colab_install.py

import deepchem_installer

# Commented out IPython magic to ensure Python compatibility.
# %time deepchem_installer.install(additional_packages=['mdtraj'], version='2.3.0')

import deepchem as dc

grid_featurizer = dc.feat.RdkitGridFeaturizer(voxel_width=2.0,feature_types=['hbond', 'salt_bridge', 'pi_stack',
'cation_pi', 'ecfp', 'splif'], sanitize=True, flatten=True)

# do not really need to use this, feature parameter 
# in the next function will take care of it

#pdbbind dataset is large ~ 2GB, takes time

# bcz of an openmm error, I have to manually download
import deepchem as dc
from deepchem.utils import download_url

import os

data_dir = dc.utils.get_data_dir()
dataset_file = os.path.join(data_dir, "pdbbind_core_df.csv.gz")

if not os.path.exists(dataset_file):
    print('File does not exist. Downloading file...')
    download_url("https://s3-us-west-1.amazonaws.com/deepchem.io/datasets/pdbbind_core_df.csv.gz")
    print('File downloaded...')

raw_dataset = dc.utils.save.load_from_disk(dataset_file)

#raw_dataset

grid_featurizer = dc.feat.RdkitGridFeaturizer(
    voxel_width=16.0, feature_types=["ecfp", "splif", "hbond", "pi_stack", "cation_pi", "salt_bridge"], 
    ecfp_power=5, splif_power=5, parallel=True, flatten=True, sanitize=True)

compound_featurizer = dc.feat.CircularFingerprint(size=128)

#load data
pdbbind_tasks, (train_dataset, valid_dataset, test_dataset), transformers = dc.molnet.load_pdbbind_grid(
    featurizer="ECFP", subset="refined")

train_dataset.get_data_shape(),valid_dataset.get_data_shape(),test_dataset.get_data_shape()

# now comes model
from sklearn.ensemble import RandomForestRegressor

#RSF model
sklearn_model = RandomForestRegressor(n_estimators=100)
model = dc.models.SklearnModel(sklearn_model)
model.fit(train_dataset)

#MLP model
n_features = train_dataset.X.shape[1]
model_mlp = dc.models.MultitaskRegressor(n_tasks=len(pdbbind_tasks),n_features=n_features,
layer_sizes=[2000, 1000],dropouts=0.5,learning_rate=0.0003)
model_mlp.fit(train_dataset, nb_epoch=250)

metric = dc.metrics.Metric(dc.metrics.pearson_r2_score)

#checking the RF model
print("Evaluating model")
train_scores = model.evaluate(train_dataset, [metric], transformers)
test_scores = model.evaluate(test_dataset, [metric], transformers)
print("Train scores ", train_scores)
print("Test scores ", test_scores)

#RF train r2: 0.89423
#RF test r2: 0.5099
#overfitting evident

#checking the MLP model
print("Evaluating model")
train_scores = model_mlp.evaluate(train_dataset, [metric], transformers)
test_scores = model_mlp.evaluate(test_dataset, [metric], transformers)
print("Train scores ", train_scores)
print("Test scores ", test_scores)

#MLP train r2: 0.918
#MLP test r2: 0.369
#overfitting more than RF

#This is different than what the book says, in book RF is more ovefit

test_dataset.y.shape

rf_predicted_test = model.predict(test_dataset)

rf_predicted_test.shape

# Commented out IPython magic to ensure Python compatibility.
#plots
# %matplotlib inline

import matplotlib
import numpy as np
import matplotlib.pyplot as plt

rf_predicted_test = model.predict(test_dataset)
mlp_predicted_test = model_mlp.predict(test_dataset)
rf_true_test = test_dataset.y

plt.figure(figsize=(10,5))
plt.subplot(1,2,1)
plt.scatter(rf_predicted_test, rf_true_test)
plt.xlabel('Predicted pIC50s')
plt.ylabel('True IC50')
plt.title(r'RF predicted IC50 vs. True IC50')
plt.xlim([-2.5, 2.5])
plt.ylim([-2.5, 2.5])
plt.plot([-2.5, 2.5], [-2.5, 2.5], color='k')
plt.text(-2.5, 2, 'RF R-squared = %0.2f' % 0.5099)
plt.subplot(1,2,2)
plt.scatter(mlp_predicted_test, rf_true_test)
plt.xlabel('Predicted pIC50s')
plt.ylabel('True IC50')
plt.title(r'MLP predicted IC50 vs. True IC50')
plt.xlim([-2.5, 2.5])
plt.ylim([-2.5, 2.5])
plt.plot([-2.5, 2.5], [-2.5, 2.5], color='k')
plt.text(-2.5, 2, 'MLP R-squared = %0.2f' % 0.369)
plt.show()

#chapter 5 completed