MMACE Paper: Random Forest for Blood-Brain Barrier

Show code cell source 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 10
      8 import skunk
      9 import mordred, mordred.descriptors
---> 10 import exmol as exmol
     11 from rdkit.Chem.Draw import rdDepictor
     12 from sklearn.ensemble import RandomForestClassifier

File /opt/hostedtoolcache/Python/3.8.18/x64/lib/python3.8/site-packages/exmol/__init__.py:3
      1 from .version import __version__
      2 from . import stoned
----> 3 from .exmol import *
      4 from .data import *
      5 from .stoned import sanitize_smiles

File /opt/hostedtoolcache/Python/3.8.18/x64/lib/python3.8/site-packages/exmol/exmol.py:30
     28 from rdkit.Chem import rdchem  # type: ignore
     29 from rdkit.DataStructs.cDataStructs import BulkTanimotoSimilarity, TanimotoSimilarity  # type: ignore
---> 30 import langchain.llms as llms
     31 import langchain.prompts as prompts
     33 from . import stoned

File /opt/hostedtoolcache/Python/3.8.18/x64/lib/python3.8/site-packages/langchain/llms/__init__.py:22
      1 """
      2 **LLM** classes provide
      3 access to the large language model (**LLM**) APIs and services.
   (...)
     18     AIMessage, BaseMessage
     19 """  # noqa: E501
     20 from typing import Any, Callable, Dict, Type
---> 22 from langchain.llms.base import BaseLLM
     25 def _import_ai21() -> Any:
     26     from langchain.llms.ai21 import AI21

File /opt/hostedtoolcache/Python/3.8.18/x64/lib/python3.8/site-packages/langchain/llms/base.py:2
      1 # Backwards compatibility.
----> 2 from langchain_core.language_models import BaseLanguageModel
      3 from langchain_core.language_models.llms import (
      4     LLM,
      5     BaseLLM,
   (...)
      9     update_cache,
     10 )
     12 __all__ = [
     13     "create_base_retry_decorator",
     14     "get_prompts",
   (...)
     19     "LLM",
     20 ]

File /opt/hostedtoolcache/Python/3.8.18/x64/lib/python3.8/site-packages/langchain_core/language_models/__init__.py:7
      1 from langchain_core.language_models.base import (
      2     BaseLanguageModel,
      3     LanguageModelInput,
      4     LanguageModelOutput,
      5     get_tokenizer,
      6 )
----> 7 from langchain_core.language_models.chat_models import BaseChatModel, SimpleChatModel
      8 from langchain_core.language_models.llms import LLM, BaseLLM
     10 __all__ = [
     11     "BaseLanguageModel",
     12     "BaseChatModel",
   (...)
     18     "LanguageModelOutput",
     19 ]

File /opt/hostedtoolcache/Python/3.8.18/x64/lib/python3.8/site-packages/langchain_core/language_models/chat_models.py:20
      7 from functools import partial
      8 from typing import (
      9     TYPE_CHECKING,
     10     Any,
   (...)
     17     cast,
     18 )
---> 20 from langchain_core.callbacks import (
     21     AsyncCallbackManager,
     22     AsyncCallbackManagerForLLMRun,
     23     BaseCallbackManager,
     24     CallbackManager,
     25     CallbackManagerForLLMRun,
     26     Callbacks,
     27 )
     28 from langchain_core.globals import get_llm_cache
     29 from langchain_core.language_models.base import BaseLanguageModel, LanguageModelInput

File /opt/hostedtoolcache/Python/3.8.18/x64/lib/python3.8/site-packages/langchain_core/callbacks/__init__.py:13
      1 from langchain_core.callbacks.base import (
      2     AsyncCallbackHandler,
      3     BaseCallbackHandler,
   (...)
     11     ToolManagerMixin,
     12 )
---> 13 from langchain_core.callbacks.manager import (
     14     AsyncCallbackManager,
     15     AsyncCallbackManagerForChainGroup,
     16     AsyncCallbackManagerForChainRun,
     17     AsyncCallbackManagerForLLMRun,
     18     AsyncCallbackManagerForRetrieverRun,
     19     AsyncCallbackManagerForToolRun,
     20     AsyncParentRunManager,
     21     AsyncRunManager,
     22     BaseRunManager,
     23     CallbackManager,
     24     CallbackManagerForChainGroup,
     25     CallbackManagerForChainRun,
     26     CallbackManagerForLLMRun,
     27     CallbackManagerForRetrieverRun,
     28     CallbackManagerForToolRun,
     29     ParentRunManager,
     30     RunManager,
     31 )
     32 from langchain_core.callbacks.stdout import StdOutCallbackHandler
     33 from langchain_core.callbacks.streaming_stdout import StreamingStdOutCallbackHandler

File /opt/hostedtoolcache/Python/3.8.18/x64/lib/python3.8/site-packages/langchain_core/callbacks/manager.py:26
      9 from typing import (
     10     TYPE_CHECKING,
     11     Any,
   (...)
     22     cast,
     23 )
     24 from uuid import UUID
---> 26 from langsmith.run_helpers import get_run_tree_context
     27 from tenacity import RetryCallState
     29 from langchain_core.callbacks.base import (
     30     BaseCallbackHandler,
     31     BaseCallbackManager,
   (...)
     37     ToolManagerMixin,
     38 )

File /opt/hostedtoolcache/Python/3.8.18/x64/lib/python3.8/site-packages/langsmith/__init__.py:10
      6 except metadata.PackageNotFoundError:
      7     # Case where package metadata is not available.
      8     __version__ = ""
---> 10 from langsmith.client import Client
     11 from langsmith.evaluation.evaluator import EvaluationResult, RunEvaluator
     12 from langsmith.run_helpers import trace, traceable

File /opt/hostedtoolcache/Python/3.8.18/x64/lib/python3.8/site-packages/langsmith/client.py:43
     41 from langsmith import schemas as ls_schemas
     42 from langsmith import utils as ls_utils
---> 43 from langsmith.evaluation import evaluator as ls_evaluator
     45 if TYPE_CHECKING:
     46     import pandas as pd

File /opt/hostedtoolcache/Python/3.8.18/x64/lib/python3.8/site-packages/langsmith/evaluation/__init__.py:4
      1 """Evaluation Helpers."""
      3 from langsmith.evaluation.evaluator import EvaluationResult, RunEvaluator
----> 4 from langsmith.evaluation.string_evaluator import StringEvaluator
      6 __all__ = ["EvaluationResult", "RunEvaluator", "StringEvaluator"]

File /opt/hostedtoolcache/Python/3.8.18/x64/lib/python3.8/site-packages/langsmith/evaluation/string_evaluator.py:3
      1 from typing import Callable, Dict, Optional
----> 3 from pydantic import BaseModel
      5 from langsmith.evaluation.evaluator import EvaluationResult, RunEvaluator
      6 from langsmith.schemas import Example, Run

File /opt/hostedtoolcache/Python/3.8.18/x64/lib/python3.8/site-packages/pydantic/__init__.py:372, in __getattr__(attr_name)
    370     return import_module(f'.{attr_name}', package=package)
    371 else:
--> 372     module = import_module(module_name, package=package)
    373     return getattr(module, attr_name)

File /opt/hostedtoolcache/Python/3.8.18/x64/lib/python3.8/importlib/__init__.py:127, in import_module(name, package)
    125             break
    126         level += 1
--> 127 return _bootstrap._gcd_import(name[level:], package, level)

File /opt/hostedtoolcache/Python/3.8.18/x64/lib/python3.8/site-packages/pydantic/main.py:15
     12 import typing_extensions
     13 from pydantic_core import PydanticUndefined
---> 15 from ._internal import (
     16     _config,
     17     _decorators,
     18     _fields,
     19     _forward_ref,
     20     _generics,
     21     _mock_val_ser,
     22     _model_construction,
     23     _repr,
     24     _typing_extra,
     25     _utils,
     26 )
     27 from ._migration import getattr_migration
     28 from .annotated_handlers import GetCoreSchemaHandler, GetJsonSchemaHandler

File /opt/hostedtoolcache/Python/3.8.18/x64/lib/python3.8/site-packages/pydantic/_internal/_decorators.py:15
     12 from typing_extensions import Literal, TypeAlias, is_typeddict
     14 from ..errors import PydanticUserError
---> 15 from ._core_utils import get_type_ref
     16 from ._internal_dataclass import slots_true
     17 from ._typing_extra import get_function_type_hints

File /opt/hostedtoolcache/Python/3.8.18/x64/lib/python3.8/site-packages/pydantic/_internal/_core_utils.py:16
     14 from pydantic_core import CoreSchema, core_schema
     15 from pydantic_core import validate_core_schema as _validate_core_schema
---> 16 from typing_extensions import TypeAliasType, TypeGuard, get_args, get_origin
     18 from . import _repr
     19 from ._typing_extra import is_generic_alias

ImportError: cannot import name 'TypeAliasType' from 'typing_extensions' (/opt/hostedtoolcache/Python/3.8.18/x64/lib/python3.8/site-packages/typing_extensions.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)