MMACE Paper: Random Forest for Blood-Brain Barrier

Hide code cell source
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import matplotlib as mpl
import rdkit, rdkit.Chem, rdkit.Chem.Draw
from rdkit.Chem.Draw import IPythonConsole
import numpy as np
import skunk
import mordred, mordred.descriptors
import exmol as exmol
from rdkit.Chem.Draw import rdDepictor
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score, plot_roc_curve

rdDepictor.SetPreferCoordGen(True)

IPythonConsole.ipython_useSVG = True
sns.set_context("notebook")
sns.set_style(
    "dark",
    {
        "xtick.bottom": True,
        "ytick.left": True,
        "xtick.color": "#666666",
        "ytick.color": "#666666",
        "axes.edgecolor": "#666666",
        "axes.linewidth": 0.8,
        "figure.dpi": 300,
    },
)
color_cycle = ["#1BBC9B", "#F06060", "#F3B562", "#6e5687", "#5C4B51"]
mpl.rcParams["axes.prop_cycle"] = mpl.cycler(color=color_cycle)
np.random.seed(0)
---------------------------------------------------------------------------
ImportError                               Traceback (most recent call last)
Cell In[1], line 14
     12 from sklearn.ensemble import RandomForestClassifier
     13 from sklearn.model_selection import train_test_split
---> 14 from sklearn.metrics import roc_auc_score, plot_roc_curve
     16 rdDepictor.SetPreferCoordGen(True)
     18 IPythonConsole.ipython_useSVG = True

ImportError: cannot import name 'plot_roc_curve' from 'sklearn.metrics' (/opt/hostedtoolcache/Python/3.11.10/x64/lib/python3.11/site-packages/sklearn/metrics/__init__.py)
data = pd.read_csv("BBBP.csv")
data.head()
# make object that can compute descriptors
calc = mordred.Calculator(mordred.descriptors, ignore_3D=True)
# make subsample from pandas df
molecules = [rdkit.Chem.MolFromSmiles(smi) for smi in data.smiles]

# the invalid molecules were None, so we'll just
# use the fact the None is False in Python
valid_mol_idx = [bool(m) for m in molecules]
valid_mols = [m for m in molecules if m]
try:
    raw_features = pd.read_pickle("raw_features.pb")
except FileNotFoundError as e:
    raw_features = calc.pandas(valid_mols, nproc=8, quiet=True)
    raw_features.to_pickle("raw_features.pb")
labels = data[valid_mol_idx].p_np
fm = raw_features.mean()
fs = raw_features.std()


def feature_convert(f):
    f -= fm
    f /= fs
    return f


features = feature_convert(raw_features)

# we have some nans in features, likely because std was 0
features = features.values.astype(float)
features_select = np.all(np.isfinite(features), axis=0)
features = features[:, features_select]
X_train, X_test, y_train, y_test = train_test_split(
    features, labels, test_size=0.2, shuffle=True
)

clf = RandomForestClassifier(max_depth=8, random_state=0)
clf.fit(X_train, y_train)
predicted = clf.predict(X_test)
print("AUC", roc_auc_score(y_test, clf.predict_proba(X_test)[:, 1]))
plt.figure(figsize=(4, 3), dpi=300)
plot_roc_curve(clf, X_test, y_test)
plt.plot([0, 1], [0, 1], linestyle="--")
plt.savefig("RF-ROC.png")
def model_eval(smiles, _=None):
    molecules = [rdkit.Chem.MolFromSmiles(smi) for smi in smiles]
    # input wrangling. Get some weird values from weird smiles
    raw_features = calc.pandas(molecules, nproc=8, quiet=True)
    features = feature_convert(raw_features)
    features = features.values.astype(float)
    features = features[:, features_select]
    labels = clf.predict(np.nan_to_num(features))
    return labels
    # return np.random.choice([True, False], size=labels.shape)


labels = model_eval(data.iloc[valid_mol_idx].smiles.values[:100])
example_neg = data.iloc[valid_mol_idx].smiles.values[np.argmin(labels)]
example_pos = data.iloc[valid_mol_idx].smiles.values[np.argmax(labels)]
example_neg_y, example_pos_y = model_eval([example_neg, example_pos])
print("neg:", example_neg, "\npos:", example_pos)
print(example_neg_y, example_pos_y)
space = exmol.sample_space(example_neg, model_eval, quiet=True)
exps = exmol.cf_explain(space)
print(exps)
fkw = {"figsize": (8, 6)}
mpl.rc("axes", titlesize=12)
exmol.plot_cf(exps, figure_kwargs=fkw, mol_size=(450, 400), nrows=1)

plt.savefig("rf-simple.png", dpi=180)
svg = exmol.insert_svg(exps, mol_fontsize=14)
with open("svg_figs/rf-simple.svg", "w") as f:
    f.write(svg)
font = {"family": "normal", "weight": "normal", "size": 22}
exmol.plot_space(
    space,
    exps,
    figure_kwargs=fkw,
    mol_size=(300, 200),
    offset=0,
    cartoon=True,
    rasterized=True,
)
plt.scatter([], [], label="Counterfactual", s=150, color=plt.get_cmap("viridis")(1.0))
plt.scatter([], [], label="Same Class", s=150, color=plt.get_cmap("viridis")(0.0))
plt.legend(fontsize=22)
plt.tight_layout()
plt.savefig("rf-space.png", dpi=180)
svg = exmol.insert_svg(exps, mol_fontsize=14)
with open("svg_figs/rf-space.svg", "w") as f:
    f.write(svg)
skunk.display(svg)

Schematic Plots

from rdkit.Chem import MolFromSmiles as smi2mol
from rdkit.Chem import MolToSmiles as mol2smi
from rdkit.Chem.Draw import MolToImage as mol2img

dos = rdkit.Chem.Draw.MolDrawOptions()
dos.useBWAtomPalette()
# dos.minFontSize = fontsize
img = mol2img(smi2mol(exps[0].smiles), options=dos)
# img.save("rf-schem-1.png")
fkw = {"figsize": (8, 4)}
font = {"family": "normal", "weight": "normal", "size": 22, "dpi": 300}
exmol.plot_space(
    space, exps[:2], figure_kwargs=fkw, mol_size=(300, 200), offset=0, cartoon=True
)
plt.scatter([], [], label="Counterfactual", s=150, color=plt.get_cmap("viridis")(1.0))
plt.scatter([], [], label="Same Class", s=150, color=plt.get_cmap("viridis")(0.0))
plt.legend(fontsize=22)
plt.tight_layout()
plt.savefig("rf-schem-3.png", bbox_inches="tight", dpi=180)
svg = exmol.insert_svg(exps[:2], mol_fontsize=10)
with open("rf-scheme.svg", "w") as f:
    f.write(svg)
skunk.display(svg)

Chemed

cspace = exmol.sample_space(
    "Cc1ccc(cc1Nc2nccc(n2)c3cccnc3)NC(=O)c4ccc(cc4)CN5CCN(CC5)C",
    model_eval,
    preset="medium",
    quiet=True,
)
kws = {"num_samples": 1500}
zspace = exmol.sample_space(
    "Cc1ccc(cc1Nc2nccc(n2)c3cccnc3)NC(=O)c4ccc(cc4)CN5CCN(CC5)C",
    model_eval,
    preset="chemed",
    method_kwargs=kws,
    quiet=True,
)
### Gleevec molecule
exps = exmol.cf_explain(zspace)
fkw = {"figsize": (8, 6)}
mpl.rc("axes", titlesize=12)
exmol.plot_cf(exps, figure_kwargs=fkw, mol_size=(450, 400), nrows=1)
fkw = {"figsize": (8, 6)}
mpl.rc("axes", titlesize=12)
cfs = exmol.cf_explain(cspace, nmols=4)
exmol.plot_cf(cfs, figure_kwargs=fkw, mol_fontsize=26, mol_size=(400, 400), nrows=1)
plt.savefig("gleevec-cs.png", bbox_inches="tight", dpi=180)
svg = exmol.insert_svg(cfs)
with open("svg_figs/gleevec-cs.svg", "w") as f:
    f.write(svg)
fkw = {"figsize": (8, 6)}
mpl.rc("axes", titlesize=12)
exmol.plot_cf(exps, figure_kwargs=fkw, mol_size=(450, 400), nrows=1)

plt.savefig("rf-simple.png", dpi=180)
svg = exmol.insert_svg(exps, mol_fontsize=14)
with open("svg_figs/gleevec-simple.svg", "w") as f:
    f.write(svg)
fkw = {"figsize": (10, 6)}
mpl.rc("axes", titlesize=12)
exmol.plot_cf(exps, figure_kwargs=fkw, mol_size=(450, 400), nrows=1)

zexps = exmol.cf_explain(zspace, nmols=5)