{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":67356,"databundleVersionId":8006601,"sourceType":"competition"}],"dockerImageVersionId":30684,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# BELKA: Molecule Representations for ML Tutorial\n\nAuthor: \n- Stephen Lee ([LinkedIn](https://www.linkedin.com/in/stephendongsoolee/))\n\n<a name=\"background\"></a>\n# 📖 Background\n---\n\n**SMILES (Simplified Molecular Input Line Entry System)** notation is a way of representing chemical structures as a string, allowing for easy input, storage, and manipulation of molecular data. The ability to convert these SMILES strings into meaningful numerical representations (embeddings) is crucial for applying machine learning techniques to chemical compounds.\n\nThis notebook aims to provide helpful, high-level material and demo code to generate various types of molecular embeddings that can be generated from SMILES and used for machine learning. Specifically, the following types of molecular embeddings are covered:\n- [Descriptors](#descriptors)\n- [Fingerprints](#fingerprints)\n- [Graphs](#graph)\n- [Mol2Vec Embeddings](#mol2vec)\n- [Chemical Language Model Embeddings](#clm)\n\nPlease leave a comment if you have questions or suggestions!","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19"}},{"cell_type":"markdown","source":"<a name=\"setup\"></a>\n# 🛠️ Environment Setup\n---","metadata":{}},{"cell_type":"code","source":"# for basic SMILES operations, descriptors and fingerprints\n!pip install -q rdkit-pypi\n\n# for graphs\n!pip install -q dgl\n\n# for mol2vec\n!pip install -q mol2vec","metadata":{"execution":{"iopub.status.busy":"2024-04-18T04:54:30.077608Z","iopub.execute_input":"2024-04-18T04:54:30.078189Z","iopub.status.idle":"2024-04-18T04:55:05.190410Z","shell.execute_reply.started":"2024-04-18T04:54:30.078120Z","shell.execute_reply":"2024-04-18T04:55:05.188802Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport seaborn as sns\n\n# for basic SMILES operations, descriptors and fingerprints\nfrom rdkit import Chem\nfrom rdkit.Chem import Descriptors, MACCSkeys, AllChem\nfrom rdkit.ML.Descriptors import MoleculeDescriptors\n\n# for graphs\nimport dgl\nimport torch\n\n# for mol2vec\nfrom gensim.models import word2vec\nfrom mol2vec.features import mol2alt_sentence, mol2sentence, MolSentence, DfVec, sentences2vec\n\n# for chemical language models\nfrom transformers import AutoModel, AutoTokenizer","metadata":{"execution":{"iopub.status.busy":"2024-04-18T04:55:16.958900Z","iopub.execute_input":"2024-04-18T04:55:16.961102Z","iopub.status.idle":"2024-04-18T04:55:17.057108Z","shell.execute_reply.started":"2024-04-18T04:55:16.961034Z","shell.execute_reply":"2024-04-18T04:55:17.055568Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<a name=\"data\"></a>\n# 🗃️ Load Data\n---","metadata":{}},{"cell_type":"code","source":"# load sample of training data for demo\ndf = pd.read_csv('/kaggle/input/leash-BELKA/train.csv', nrows=1000)\ndf.head()","metadata":{"execution":{"iopub.status.busy":"2024-04-18T04:55:30.303888Z","iopub.execute_input":"2024-04-18T04:55:30.304425Z","iopub.status.idle":"2024-04-18T04:55:30.349892Z","shell.execute_reply.started":"2024-04-18T04:55:30.304391Z","shell.execute_reply":"2024-04-18T04:55:30.348499Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# assign unique molecule_smiles values to variable. This is all we need for the purpose of this demo\nsmiles = df['molecule_smiles'].unique()\nprint(len(smiles))","metadata":{"execution":{"iopub.status.busy":"2024-04-18T04:55:31.145921Z","iopub.execute_input":"2024-04-18T04:55:31.146591Z","iopub.status.idle":"2024-04-18T04:55:31.159786Z","shell.execute_reply.started":"2024-04-18T04:55:31.146541Z","shell.execute_reply":"2024-04-18T04:55:31.158397Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"SMILES can be converted into standardized data structures called Mol objects which provide rich representations of molecular structures. Many of the embedding generation featured in this notebook relies on Mol objects as input.\n\nMol objects are generated using the RDKit package, a widely used cheminformatics library for molecular modeling, chemical analysis and computational chemistry.","metadata":{}},{"cell_type":"code","source":"# obtain RDKit mol objects from SMILES\nmols = pd.Series(smiles).apply(Chem.MolFromSmiles)","metadata":{"execution":{"iopub.status.busy":"2024-04-18T04:55:32.194458Z","iopub.execute_input":"2024-04-18T04:55:32.194992Z","iopub.status.idle":"2024-04-18T04:55:32.300370Z","shell.execute_reply.started":"2024-04-18T04:55:32.194952Z","shell.execute_reply":"2024-04-18T04:55:32.298881Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# mol objects can be readily viewed as 2d molecular drawings like so\nmols[0]","metadata":{"execution":{"iopub.status.busy":"2024-04-17T23:29:53.811463Z","iopub.execute_input":"2024-04-17T23:29:53.811828Z","iopub.status.idle":"2024-04-17T23:29:53.841418Z","shell.execute_reply.started":"2024-04-17T23:29:53.811796Z","shell.execute_reply":"2024-04-17T23:29:53.840072Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<a name=\"descriptors\"></a>\n# 🧪 Molecular Descriptors\n---\n\nMolecular descriptors are quantitative properties that capture information about the chemical structure that can be used effectively in cheminformatics, particularly for the modeling and analysis of chemical compounds. Descriptors can pertain to various aspects of a chemical structure, such as its constitution, topology, geometry, electron distribution, and hydrophobicity.\n\nThe following code demonstrates how to generate molecular descriptors from SMILES.\n\nReferences:\n- [RDKit Documentation](https://www.rdkit.org/docs/)\n- [RDKit GitHub](https://github.com/rdkit/rdkit)","metadata":{}},{"cell_type":"code","source":"# generate RDKit descriptors from Mol objects\ndef generate_descriptors(mols):\n    calc = MoleculeDescriptors.MolecularDescriptorCalculator([x[0] for x in Descriptors._descList])\n    \n    mol_descriptors = []\n    for mol in mols:\n        rdkit_descriptors = calc.CalcDescriptors(mol)\n        mol_descriptors.append(rdkit_descriptors)\n        \n    desc_names = calc.GetDescriptorNames()\n    return pd.DataFrame(mol_descriptors, columns=desc_names)\n\ndescriptor_df = generate_descriptors(mols)\ndescriptor_df.head()","metadata":{"execution":{"iopub.status.busy":"2024-04-17T23:29:53.843091Z","iopub.execute_input":"2024-04-17T23:29:53.843500Z","iopub.status.idle":"2024-04-17T23:30:05.028232Z","shell.execute_reply.started":"2024-04-17T23:29:53.843454Z","shell.execute_reply":"2024-04-17T23:30:05.026763Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Since molecular descriptors directly represent some characteristic of a chemical structure, they are readily usable for analysis. For example, you could examine the relationship between molecular weight and logP (measure of hydrophobicity):","metadata":{}},{"cell_type":"code","source":"sns.scatterplot(x='MolWt', y='MolLogP', data=descriptor_df)\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2024-04-17T23:30:05.029761Z","iopub.execute_input":"2024-04-17T23:30:05.030226Z","iopub.status.idle":"2024-04-17T23:30:05.349755Z","shell.execute_reply.started":"2024-04-17T23:30:05.030166Z","shell.execute_reply":"2024-04-17T23:30:05.348505Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<a name=\"fingerprints\"></a>\n# 🗝️ Molecular Fingerprints\n---\n**Molecular fingerprints** are fixed-length numerical representations that encode molecular structures, usually as a binary vector. Each element (called a bit) of a fingerprint vector represents the presence of specific atomic or structural features within the chemical compound.\n\nBelow is example code to generate two types of commonly used molecular fingerprints: \n- **MACCS Keys**: a 166-bit fingerprint based on a predefined list of molecular substructures or patterns\n- **Morgan Fingerprints**: a fixed-length bit fingerprint based on hashed topological features of atoms and their bond connectivities within a specified radius. Also referred to as Extended-Connectivity Fingerprints (ECFPs)\n\nReferences::\n- Durant, J. L., Leland, B. A., Henry, D. R., & Nourse, J. G. (2002). Reoptimization of MDL keys for use in drug discovery. Journal of Chemical Information and Computer Sciences, 42(6), 1273-1280. [Link to Article](https://pubs.acs.org/doi/10.1021/ci010132r)\n- Rogers, D., & Hahn, M. (2010). Extended-connectivity fingerprints. Journal of Chemical Information and Modeling, 50(5), 742-754. [Link to Article](https://pubs.acs.org/doi/10.1021/ci100050t)","metadata":{}},{"cell_type":"code","source":"# generate MACCS Keys from RDKit Mol objects as a dataframe\nmaccs_list = [MACCSkeys.GenMACCSKeys(mol) for mol in mols]\n\nmaccs_data = []\nfor maccs in maccs_list:\n    bit_array = list(map(int, maccs.ToBitString()))\n    maccs_data.append(bit_array[1:])  # skip 1st bit as it is unused in RDKit's MACCSkeys.GenMACCSKeys implementation\n\nmaccs_df = pd.DataFrame(maccs_data)\nmaccs_df.head()","metadata":{"execution":{"iopub.status.busy":"2024-04-17T23:47:59.368225Z","iopub.execute_input":"2024-04-17T23:47:59.368709Z","iopub.status.idle":"2024-04-17T23:47:59.891903Z","shell.execute_reply.started":"2024-04-17T23:47:59.368673Z","shell.execute_reply":"2024-04-17T23:47:59.890670Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# generate Morgan Fingerprints from RDKit Mol objects as a dataframe\nmfpt_list = [AllChem.GetMorganFingerprintAsBitVect(mol, radius=2, nBits=1024) for mol in mols]\n\nmfpt_df = pd.DataFrame([list(fpt) for fpt in mfpt_list])\nmfpt_df.head()","metadata":{"execution":{"iopub.status.busy":"2024-04-17T23:42:16.504499Z","iopub.execute_input":"2024-04-17T23:42:16.504927Z","iopub.status.idle":"2024-04-17T23:42:17.164201Z","shell.execute_reply.started":"2024-04-17T23:42:16.504895Z","shell.execute_reply":"2024-04-17T23:42:17.163070Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Note that different values can be specified for radius and nBits when generating Morgan Fingerprints.\n- **radius:** number of bond steps to include from each atom when calculating the fingerprint. A radius of 1 means that only the immediate neighbors are considered; a radius of 2 includes neighbors up to two bonds away, etc.\n- **nBits:** total number of bits in the fingerprint (i.e. resolution). Larger nBits values increases feature dimensionality and typically increases feature sparsity.\n\n**Bit Collision** refers to the scenario where different substructural features are mapped to the same bit position. This is undesirable as it reduces the fingerprint's ability to differentiate molecules. Increasing nBits increases the available feature space and thus lowers the chance of bit collisions.","metadata":{}},{"cell_type":"markdown","source":"<a name=\"graph\"></a>\n# ⚛️ Molecular Graph Embeddings\n---\nA graph is a type of data structure to represent some interconnected entitiy (e.g. social networks). Molecules can be represented as graphs with atoms as nodes and bonds as edges. \n\nBelow is example code to generate graphs from RDKit Mol objects using the Deep Graph Learning (DGL) package.\n\nReferences::\n- Wang, M., Zheng, D., Ye, Z., Gan, Q., Li, M., Song, X., Zhou, J., Ma, C., Yu, L., Gai, Y., Xiao, T., He, T., Karypis, G., Li, J., & Zhang, Z. (2019). Deep graph library: A graph-centric, highly-performant package for graph neural networks. [Link to Article](https://arxiv.org/abs/1909.01315)\n- [Deep Graph Libarary Docs](https://docs.dgl.ai/)","metadata":{}},{"cell_type":"code","source":"# generate molecular graphs from RDKit Mol objects\ndef mol_to_dgl_graph(mol):\n    graph = dgl.DGLGraph()\n\n    # add nodes\n    num_atoms = mol.GetNumAtoms()\n    graph.add_nodes(num_atoms)\n\n    # add edges\n    for bond in mol.GetBonds():\n        # DGL graphs are directed so both directions are added\n        graph.add_edges(bond.GetBeginAtomIdx(), bond.GetEndAtomIdx())\n        graph.add_edges(bond.GetEndAtomIdx(), bond.GetBeginAtomIdx())\n\n    # add node features - e.g. atomic number\n    node_feats = torch.tensor([[atom.GetAtomicNum()] for atom in mol.GetAtoms()], dtype=torch.float32)\n    graph.ndata['feat'] = node_feats\n\n    # add edge features (optional) - e.g. bond type: single/double/triple as 1, 2, 3\n    edge_feats = []\n    for bond in mol.GetBonds():\n        bond_type = bond.GetBondTypeAsDouble()\n        edge_feats.append([bond_type])\n        edge_feats.append([bond_type])  # add both directions\n    graph.edata['type'] = torch.tensor(edge_feats, dtype=torch.float32)\n\n    return graph\n\ngraphs = mols.apply(mol_to_dgl_graph)\nprint(graphs)","metadata":{"execution":{"iopub.status.busy":"2024-04-18T00:04:28.585901Z","iopub.execute_input":"2024-04-18T00:04:28.586381Z","iopub.status.idle":"2024-04-18T00:04:39.576007Z","shell.execute_reply.started":"2024-04-18T00:04:28.586344Z","shell.execute_reply":"2024-04-18T00:04:39.574819Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# graph neural networks typically process graphs in batches for efficiency\nbatched_graph = dgl.batch(graphs.tolist())\nbatched_graph","metadata":{"execution":{"iopub.status.busy":"2024-04-18T00:05:12.556340Z","iopub.execute_input":"2024-04-18T00:05:12.556775Z","iopub.status.idle":"2024-04-18T00:05:12.574069Z","shell.execute_reply.started":"2024-04-18T00:05:12.556739Z","shell.execute_reply":"2024-04-18T00:05:12.572417Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<a name=\"mol2vec\"></a>\n# 💬 Mol2Vec Embeddings\n---\nMol2Vec is a version of the widespread natural language processing (NLP) algorithm Word2vec, specifically adapted for molecular data. Mol2Vec represents molecules in a continuous vector space that captures chemical similarity in a mannner that is analagous to how Word2Vec captures semantic similarity between words. Mol2Vec treats chemical substructures (e.g. Morgan fingerprints with radius=1) as fundamental units of a 'molecular sentence', similar to how words are treated in NLP.\n\nA Mol2Vec model pre-trained on 20 million molecules is available to generate 300 dimensional embeddings based on learned molecular vector representations. Below is example code to generate Mol2Vec embeddings \n\nReferences:\n- Jaeger, S., Fulle, S., & Turk, S. (2018). Mol2vec: unsupervised machine learning approach with chemical intuition. Journal of Chemical Information and Modeling, 58(1), 27-35. [Link to Article](https://pubs.acs.org/doi/10.1021/acs.jcim.7b00616)\n- [Mol2Vec GitHub Repo](https://github.com/samoturk/mol2vec)","metadata":{}},{"cell_type":"code","source":"# load pre-trained mol2vec model\nmol2vec_url = 'https://github.com/samoturk/mol2vec/raw/master/examples/models/model_300dim.pkl'\nmol2vec_model = word2vec.Word2Vec.load(mol2vec_url)","metadata":{"execution":{"iopub.status.busy":"2024-04-18T04:48:10.761771Z","iopub.execute_input":"2024-04-18T04:48:10.762279Z","iopub.status.idle":"2024-04-18T04:48:35.845048Z","shell.execute_reply.started":"2024-04-18T04:48:10.762240Z","shell.execute_reply":"2024-04-18T04:48:35.843691Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# generate molecular sentences\nmol_sentences = mols.apply(lambda x: MolSentence(mol2alt_sentence(x, 1)))\n\n# version of sentences2vec() compatible with gensim v4.0 (source: https://github.com/samoturk/mol2vec/issues/14)\ndef sentences2vec(sentences, model, unseen=None):\n    \"\"\"Generate vectors for each sentence (list) in a list of sentences. Vector is simply a\n    sum of vectors for individual words.\n    \n    Parameters\n    ----------\n    sentences : list, array\n        List with sentences\n    model : word2vec.Word2Vec\n        Gensim word2vec model\n    unseen : None, str\n        Keyword for unseen words. If None, those words are skipped.\n        https://stats.stackexchange.com/questions/163005/how-to-set-the-dictionary-for-text-analysis-using-neural-networks/163032#163032\n\n    Returns\n    -------\n    np.array\n    \"\"\"\n    \n    keys = set(model.wv.key_to_index)\n    vec = []\n    \n    if unseen:\n        unseen_vec = model.wv.get_vector(unseen)\n\n    for sentence in sentences:\n        if unseen:\n            vec.append(sum([model.wv.get_vector(y) if y in set(sentence) & keys\n                       else unseen_vec for y in sentence]))\n        else:\n            vec.append(sum([model.wv.get_vector(y) for y in sentence \n                            if y in set(sentence) & keys]))\n    return np.array(vec)\n\n# generate vector embeddings from molecular sentences\nmol2vec_embeddings = np.array([DfVec(x) for x in sentences2vec(mol_sentences, mol2vec_model, unseen='UNK')])\nmol2vec_embeddings[:5]","metadata":{"execution":{"iopub.status.busy":"2024-04-18T05:00:27.628816Z","iopub.execute_input":"2024-04-18T05:00:27.629358Z","iopub.status.idle":"2024-04-18T05:00:28.592721Z","shell.execute_reply.started":"2024-04-18T05:00:27.629318Z","shell.execute_reply":"2024-04-18T05:00:28.591166Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<a name=\"clm\"></a>\n# 🤖 Chemical Language Model Embeddings\n---\n\nJust as traditional language models use self-attention mechanisms to compute the representation of each language element (e.g. word in a sentence) to every other element, chemical language models use the same principle in which elements are some chemical unit (e.g. atoms) instead of words.\n\nBelow is example code to obtain learned transformer-based embeddings from two chemical language models: \n- ChemBERTa: adapted for chemical SMILES from the RoBERTa architecture, trained on a dataset of 77 million molecules\n- MoLFormer: another transformer-based model adapted for SMILES but trained on a larger dataset (1.1 billion molecules!)\n\nReferences:\n- Chithrananda, S., Grand, G., & Ramsundar, B. (2020). ChemBERTa: Large-Scale Self-Supervised Pretraining for Molecular Property Prediction. [Link to Article](https://arxiv.org/abs/2010.09885)\n- Chithrananda, S., Grand, G., & Ramsundar, B. (2022). ChemBERTa-2: Towards Chemical Foundation Models. [Link to Article](https://arxiv.org/abs/2209.01712)\n- Ross, J., Belgodere, B., Chenthamarakshan, V., et al. (2022). Large-scale chemical language representations capture molecular structure and properties. Nature Machine Intelligence, 4, 1256-1264. [Link to Article](https://www.nature.com/articles/s42256-022-00580-7)\n- [ChemBERTa on HuggingFace Model Repo](https://huggingface.co/seyonec/ChemBERTa-zinc-base-v1)\n- [MolFormer on HuggingFace Model Repo](https://huggingface.co/ibm/MoLFormer-XL-both-10pct)\n- [MolFormer GitHub repo](https://github.com/IBM/molformer)","metadata":{}},{"cell_type":"code","source":"# load pre-trained ChemBERTa model checkpoint and tokenizer\ncb_tokenizer = AutoTokenizer.from_pretrained('DeepChem/ChemBERTa-10M-MLM')\ncb_model = AutoModel.from_pretrained('DeepChem/ChemBERTa-10M-MLM')\ncb_model.eval()\n\n# tokenize SMILES\ncb_encoded_inputs = cb_tokenizer(list(smiles), padding=True, truncation=True, return_tensors=\"pt\")\n\n# calculate embeddings\nwith torch.no_grad():\n    outputs = cb_model(**cb_encoded_inputs)\n\n# extract pooled output\ncb_embeddings = outputs.pooler_output\n\ncb_embeddings_df = pd.DataFrame(cb_embeddings.numpy())\ncb_embeddings_df.head()","metadata":{"execution":{"iopub.status.busy":"2024-04-18T05:24:20.352109Z","iopub.execute_input":"2024-04-18T05:24:20.352639Z","iopub.status.idle":"2024-04-18T05:24:43.766081Z","shell.execute_reply.started":"2024-04-18T05:24:20.352604Z","shell.execute_reply":"2024-04-18T05:24:43.764779Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# load pre-trained MolFormer model checkpoint and tokenizer\nmf_tokenizer = AutoTokenizer.from_pretrained('ibm/MoLFormer-XL-both-10pct', deterministic_eval=True, trust_remote_code=True)\nmf_model = AutoModel.from_pretrained('ibm/MoLFormer-XL-both-10pct', trust_remote_code=True)\n\n# tokenize SMILES\nmf_encoded_inputs = mf_tokenizer(list(smiles), padding=True, return_tensors=\"pt\")\n\n# calculate embeddings\nwith torch.no_grad():\n    outputs = mf_model(**mf_encoded_inputs)\n\n# extract embeddings\nmf_embeddings = outputs.pooler_output\n\nmf_embeddings_df = pd.DataFrame(mf_embeddings.numpy())\nmf_embeddings_df.head()","metadata":{"execution":{"iopub.status.busy":"2024-04-18T05:29:12.181270Z","iopub.execute_input":"2024-04-18T05:29:12.181765Z","iopub.status.idle":"2024-04-18T05:29:46.017239Z","shell.execute_reply.started":"2024-04-18T05:29:12.181724Z","shell.execute_reply":"2024-04-18T05:29:46.016235Z"},"trusted":true},"execution_count":null,"outputs":[]}]}