Upload 11 files

app.py

@@ -0,0 +1,148 @@
+import streamlit as st
+import streamlit_ext as ste
+
+from inference import Inference
+import random
+from rdkit.Chem import Draw
+from rdkit import Chem
+from rdkit.Chem.Draw import IPythonConsole
+import io
+from PIL import Image
+
+class DrugGENConfig:
+    submodel='DrugGEN'
+    act='relu'
+    max_atom=45
+    dim=32
+    depth=1
+    heads=8
+    mlp_ratio=3
+    dropout=0.
+    features=False
+    inference_sample_num=1000
+    inf_batch_size=1
+    protein_data_dir='data/akt'
+    drug_index='data/drug_smiles.index'
+    drug_data_dir='data/akt'
+    mol_data_dir='data'
+    log_dir='experiments/logs'
+    model_save_dir='experiments/models'
+    sample_dir='experiments/samples'
+    result_dir="experiments/tboard_output"
+    inf_dataset_file="chembl45_test.pt"
+    inf_drug_dataset_file='akt_test.pt'
+    inf_raw_file='data/chembl_test.smi'
+    inf_drug_raw_file="data/akt_test.smi"
+    inference_model="experiments/models/DrugGEN"
+    log_sample_step=1000
+    set_seed=False
+    seed=1
+
+class NoTargetConfig(DrugGENConfig):
+    submodel="NoTarget"
+    dim=128
+    inference_model="experiments/models/NoTarget"
+
+
+model_configs = {
+    "DrugGEN": DrugGENConfig(),
+    "NoTarget": NoTargetConfig()
+}
+
+
+with st.sidebar:
+    st.title("DrugGEN: Target Centric De Novo Design of Drug Candidate Molecules with Graph Generative Deep Adversarial Networks")
+    st.write("[![arXiv](https://img.shields.io/badge/arXiv-2302.07868-b31b1b.svg)](https://arxiv.org/abs/2302.07868) [![github-repository](https://img.shields.io/badge/GitHub-black?logo=github)](https://github.com/HUBioDataLab/DrugGEN)")
+
+    with st.expander("Expand to display information about models"):
+        st.write("""
+### Model Variations
+- **DrugGEN-Prot**: composed of two GANs, incorporates protein features to the transformer decoder module of GAN2 (together with the de novo molecules generated by GAN1) to direct the target centric molecule design.
+- **DrugGEN-CrossLoss**: composed of one GAN, the input of the GAN1 generator is the real molecules dataset and the GAN1 discriminator compares the generated molecules with the real inhibitors of the given target.
+- **DrugGEN-NoTarget**: composed of one GAN, focuses on learning the chemical properties from the ChEMBL training dataset, no target-specific generation.
+
+        """)    
+
+    with st.form("model_selection_from"):
+        model_name = st.radio(
+            'Select a model to make inference (DrugGEN-Prot and DrugGEN-CrossLoss models design molecules to target the AKT1 protein)',
+            ('DrugGEN-Prot', 'DrugGEN-CrossLoss', 'DrugGEN-NoTarget')
+        )
+
+        model_name = model_name.replace("DrugGEN-", "")
+
+        molecule_num_input = st.number_input('Number of molecules to generate', min_value=1, max_value=100_000, value=1000, step=1)
+
+        seed_input = st.number_input("RNG seed value (can be used for reproducibility):", min_value=0, value=42, step=1)
+        
+        submitted = st.form_submit_button("Start Computing")
+
+
+
+if submitted:
+# if submitted or ("submitted" in st.session_state):
+    # st.session_state["submitted"] = True
+    config = model_configs[model_name]
+
+    config.inference_sample_num = molecule_num_input
+    config.seed = seed_input
+    
+    with st.spinner(f'Creating the trainer class instance for {model_name}...'):
+        trainer = Trainer(config)
+    with st.spinner(f'Running inference function of {model_name} (this may take a while) ...'):
+        results = trainer.inference()
+    st.success(f"Inference of {model_name} took {results['runtime']:.2f} seconds.")
+
+    with st.expander("Expand to see the generation performance scores"):
+        st.write("### Generation performance scores (novelty is calculated in comparison to the training dataset)")
+        st.success(f"Validity: {results['fraction_valid']}")
+        st.success(f"Uniqueness: {results['uniqueness']}")
+        st.success(f"Novelty: {results['novelty']}")
+
+    with open(f'experiments/inference/{model_name}/inference_drugs.txt') as f:
+        inference_drugs = f.read()
+    # st.download_button(label="Click to download generated molecules", data=inference_drugs, file_name=f'DrugGEN-{model_name}_denovo_mols.smi', mime="text/plain")
+    ste.download_button(label="Click to download generated molecules", data=inference_drugs, file_name=f'DrugGEN-{model_name}_denovo_mols.smi', mime="text/plain")
+
+    
+    st.write("Structures of randomly selected 12 de novo molecules from the inference set:")
+    # from rdkit.Chem import Draw
+# img = Draw.MolsToGridImage(mol_list, molsPerRow=5, subImgSize=(250, 250), maxMols=num_mols,
+                           # legends=None, useSVG=True)
+    generated_molecule_list = inference_drugs.split("\n")
+
+    selected_molecules = random.choices(generated_molecule_list,k=12)
+
+    selected_molecules = [Chem.MolFromSmiles(mol) for mol in selected_molecules]
+    # IPythonConsole.UninstallIPythonRenderer()
+    drawOptions = Draw.rdMolDraw2D.MolDrawOptions()
+    drawOptions.prepareMolsBeforeDrawing = False
+    drawOptions.bondLineWidth = 1.
+
+    molecule_image = Draw.MolsToGridImage(
+        selected_molecules,
+        molsPerRow=3,
+        subImgSize=(250, 250),
+        maxMols=len(selected_molecules),
+        # legends=None,
+        returnPNG=False,
+        # drawOptions=drawOptions,
+        highlightAtomLists=None,
+        highlightBondLists=None,
+        
+    )
+    print(type(molecule_image))
+    # print(type(molecule_image._data_and_metadata()))
+    molecule_image.save("result_grid.png") 
+    # png_data = io.BytesIO()
+    # molecule_image.save(png_data, format='PNG')
+    # png_data.seek(0)
+    
+    # Step 2: Read the PNG image data as a PIL image
+    # pil_image = Image.open(png_data)
+    # st.image(pil_image)
+    st.image(molecule_image)
+
+else:
+    st.warning("Please select a model to make inference")
+

gradio_app.py

@@ -0,0 +1,164 @@
+import gradio as gr
+from inference import Inference
+import PIL
+from PIL import Image
+import pandas as pd
+import random
+from rdkit import Chem
+from rdkit.Chem import Draw
+from rdkit.Chem.Draw import IPythonConsole
+import shutil
+
+class DrugGENConfig:
+    submodel='DrugGEN'
+    act='relu'
+    max_atom=45
+    dim=32
+    depth=1
+    heads=8
+    mlp_ratio=3
+    dropout=0.
+    features=False
+    inference_sample_num=1000
+    inf_batch_size=1
+    protein_data_dir='data/akt'
+    drug_index='data/drug_smiles.index'
+    drug_data_dir='data/akt'
+    mol_data_dir='data'
+    log_dir='experiments/logs'
+    model_save_dir='experiments/models'
+    inference_model="experiments/models/DrugGEN"
+    sample_dir='experiments/samples'
+    result_dir="experiments/tboard_output"
+    dataset_file="chembl45_train.pt"
+    drug_dataset_file="akt_train.pt"
+    raw_file='data/chembl_train.smi'
+    drug_raw_file="data/akt_train.smi"
+    inf_dataset_file="chembl45_test.pt"
+    inf_drug_dataset_file='akt_test.pt'
+    inf_raw_file='data/chembl_test.smi'
+    inf_drug_raw_file="data/akt_test.smi"
+    log_sample_step=1000
+    set_seed=False
+    seed=1
+
+
+class NoTargetConfig(DrugGENConfig):
+    submodel="NoTarget"
+    dim=128
+    inference_model="experiments/models/NoTarget"
+
+
+model_configs = {
+    "DrugGEN": DrugGENConfig(),
+    "NoTarget": NoTargetConfig(),
+}
+
+
+
+def function(model_name: str, mol_num: int, seed: int) -> tuple[PIL.Image, pd.DataFrame, str]:
+    '''
+    Returns:
+    image, score_df, file path
+    '''
+
+    config = model_configs[model_name]
+    config.inference_sample_num = mol_num
+    config.seed = seed
+
+    inferer = Inference(config)
+    scores = inferer.inference() # create scores_df out of this
+
+    score_df = pd.DataFrame(scores, index=[0])
+
+    output_file_path = f'experiments/inference/{model_name}/inference_drugs.txt'
+
+    import os
+    new_path = f'{model_name}_denovo_mols.smi'
+    os.rename(output_file_path, new_path)
+
+    with open(new_path) as f:
+        inference_drugs = f.read()
+
+    generated_molecule_list = inference_drugs.split("\n")
+
+    rng = random.Random(seed)
+    selected_molecules = rng.choices(generated_molecule_list,k=12)
+    selected_molecules = [Chem.MolFromSmiles(mol) for mol in selected_molecules]
+
+    drawOptions = Draw.rdMolDraw2D.MolDrawOptions()
+    drawOptions.prepareMolsBeforeDrawing = False
+    drawOptions.bondLineWidth = 0.5
+
+    molecule_image = Draw.MolsToGridImage(
+        selected_molecules,
+        molsPerRow=3,
+        subImgSize=(400, 400),
+        maxMols=len(selected_molecules),
+        # legends=None,
+        returnPNG=False,
+        drawOptions=drawOptions,
+        highlightAtomLists=None,
+        highlightBondLists=None,
+    )
+
+    return molecule_image, score_df, new_path
+
+
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.Markdown("# DrugGEN: Target Centric De Novo Design of Drug Candidate Molecules with Graph Generative Deep Adversarial Networks")
+            with gr.Row():
+                gr.Markdown("[![arXiv](https://img.shields.io/badge/arXiv-2302.07868-b31b1b.svg)](https://arxiv.org/abs/2302.07868)")
+                gr.Markdown("[![github-repository](https://img.shields.io/badge/GitHub-black?logo=github)](https://github.com/HUBioDataLab/DrugGEN)")
+            
+            with gr.Accordion("Expand to display information about models", open=False):
+                gr.Markdown("""
+### Model Variations
+- **DrugGEN** is the default model. The input of the generator is the real molecules (ChEMBL) dataset (to ease the learning process) and the discriminator compares the generated molecules with the real inhibitors of the given target protein.
+- **NoTarget** is the non-target-specific version of DrugGEN. This model only focuses on learning the chemical properties from the ChEMBL training dataset.
+        """)
+            model_name = gr.Radio(
+                choices=("DrugGEN", "NoTarget"),
+                value="DrugGEN",
+                label="Select a model to make inference",
+                info=" DrugGEN model design molecules to target the AKT1 protein"
+            )
+
+            num_molecules = gr.Number(
+                label="Number of molecules to generate",
+                precision=0, # integer input
+                minimum=1,
+                value=1000,
+                maximum=10_000,
+            )
+            seed_num = gr.Number(
+                label="RNG seed value (can be used for reproducibility):",
+                precision=0, # integer input
+                minimum=0,
+                value=42,
+            )
+
+            submit_button = gr.Button(
+                value="Start Generating"
+            )
+
+        with gr.Column(scale=2):
+            scores_df = gr.Dataframe(
+                label="Scores",
+                headers=["Runtime (seconds)", "Validity", "Uniqueness", "Novelty (Train)", "Novelty (Inference)", "Novelty (AKT)", "MaxLen", "MeanAtomType", "SNN (ChEMBL)", "SNN (AKT)"],
+            )
+            file_download = gr.File(
+                label="Click to download generated molecules",
+            )
+            image_output = gr.Image(
+                label="Structures of randomly selected 12 de novo molecules from the inference set:"
+            )
+
+
+    submit_button.click(function, inputs=[model_name, num_molecules, seed_num], outputs=[image_output, scores_df, file_download], api_name="inference")
+
+demo.queue(concurrency_count=1)
+demo.launch()

inference.py

@@ -0,0 +1,268 @@
+import os
+import time
+import pickle
+import random
+from tqdm import tqdm
+import argparse
+
+import torch
+from torch_geometric.loader import DataLoader
+import torch.utils.data
+from rdkit import RDLogger
+torch.set_num_threads(5)
+RDLogger.DisableLog('rdApp.*')
+
+from utils import *
+from models import Generator
+from new_dataloader import DruggenDataset
+from loss import generator_loss
+from training_data import load_molecules
+
+
+class Inference(object):
+    """Inference class for DrugGEN."""
+
+    def __init__(self, config):
+        if config.set_seed:
+            np.random.seed(config.seed)
+            random.seed(config.seed)
+            torch.manual_seed(config.seed)
+            torch.cuda.manual_seed_all(config.seed)
+
+            torch.backends.cudnn.deterministic = True
+            torch.backends.cudnn.benchmark = False
+
+            os.environ["PYTHONHASHSEED"] = str(config.seed)
+
+            print(f'Using seed {config.seed}')
+
+        self.device = torch.device("cuda" if torch.cuda.is_available() else 'cpu')
+
+        # Initialize configurations
+        self.submodel = config.submodel
+        
+        self.inference_model = config.inference_model
+        self.sample_num = config.sample_num
+
+        # Data loader.
+        self.inf_raw_file = config.inf_raw_file  # SMILES containing text file for first dataset. 
+                                         # Write the full path to file.
+        self.inf_dataset_file = config.inf_dataset_file    # Dataset file name for the first GAN. 
+                                                   # Contains large number of molecules.
+        self.inf_batch_size = config.inf_batch_size
+        self.mol_data_dir = config.mol_data_dir  # Directory where the dataset files are stored.
+        self.dataset_name = self.inf_dataset_file.split(".")[0]
+        self.max_atom = config.max_atom  # Model is based on one-shot generation. 
+                                         # Max atom number for molecules must be specified.
+        self.features = config.features  # Small model uses atom types as node features. (Boolean, False uses atom types only.)
+                                         # Additional node features can be added. Please check new_dataloarder.py Line 102.
+
+        self.inf_dataset = DruggenDataset(self.mol_data_dir,
+                                      self.inf_dataset_file,
+                                      self.inf_raw_file,
+                                      self.max_atom,
+                                      self.features) # Dataset for the first GAN. Custom dataset class from PyG parent class.
+                                                     # Can create any molecular graph dataset given smiles string.
+                                                     # Nonisomeric SMILES are suggested but not necessary.
+                                                     # Uses sparse matrix representation for graphs,
+                                                     # For computational and speed efficiency.
+        
+        self.inf_loader = DataLoader(self.inf_dataset,
+                                 shuffle=True,
+                                 batch_size=self.inf_batch_size,
+                                 drop_last=True)  # PyG dataloader for the first GAN.
+
+
+        # Atom and bond type dimensions for the construction of the model.
+        self.atom_decoders = self.decoder_load("atom")  # Atom type decoders for first GAN. 
+                                                        # eg. 0:0, 1:6 (C), 2:7 (N), 3:8 (O), 4:9 (F)
+        self.bond_decoders = self.decoder_load("bond")  # Bond type decoders for first GAN.
+                                                        # eg. 0: (no-bond), 1: (single), 2: (double), 3: (triple), 4: (aromatic)
+        self.m_dim = len(self.atom_decoders) if not self.features else int(self.inf_loader.dataset[0].x.shape[1]) # Atom type dimension.
+        self.b_dim = len(self.bond_decoders) # Bond type dimension.
+        self.vertexes = int(self.inf_loader.dataset[0].x.shape[0]) # Number of nodes in the graph.
+
+        # Transformer and Convolution configurations.
+        self.act = config.act
+        self.dim = config.dim
+        self.depth = config.depth
+        self.heads = config.heads
+        self.mlp_ratio = config.mlp_ratio
+        self.dropout = config.dropout
+
+        self.build_model()
+
+
+    def build_model(self):
+        """Create generators and discriminators."""
+        self.G = Generator(self.act,
+                           self.vertexes,
+                           self.b_dim,
+                           self.m_dim,
+                           self.dropout,
+                           dim=self.dim,
+                           depth=self.depth,
+                           heads=self.heads,
+                           mlp_ratio=self.mlp_ratio,
+                           submodel = self.submodel)
+
+        self.print_network(self.G, 'G')
+
+        self.G.to(self.device)
+
+
+    def decoder_load(self, dictionary_name):
+        ''' Loading the atom and bond decoders'''
+        with open("DrugGEN/data/decoders/" + dictionary_name + "_" + self.dataset_name + '.pkl', 'rb') as f:
+            return pickle.load(f)
+
+
+    def print_network(self, model, name):
+        """Print out the network information."""
+        num_params = 0
+        for p in model.parameters():
+            num_params += p.numel() 
+        print(model)
+        print(name)
+        print("The number of parameters: {}".format(num_params))
+
+
+    def restore_model(self, submodel, model_directory):
+        """Restore the trained generator and discriminator."""
+        print('Loading the model...')
+        G_path = os.path.join(model_directory, '{}-G.ckpt'.format(submodel))
+        self.G.load_state_dict(torch.load(G_path, map_location=lambda storage, loc: storage))
+
+
+    def inference(self):
+        # Load the trained generator.
+        self.restore_model(self.submodel, self.inference_model)
+
+        # smiles data for metrics calculation.
+        chembl_smiles = [line for line in open("DrugGEN/data/chembl_train.smi", 'r').read().splitlines()]
+        chembl_test = [line for line in open("DrugGEN/data/chembl_test.smi", 'r').read().splitlines()]
+        drug_smiles = [line for line in open("DrugGEN/data/akt_inhibitors.smi", 'r').read().splitlines()]
+        drug_mols = [Chem.MolFromSmiles(smi) for smi in drug_smiles]
+        drug_vecs = [AllChem.GetMorganFingerprintAsBitVect(x, 2, nBits=1024) for x in drug_mols if x is not None]
+
+
+        # Make directories if not exist.
+        if not os.path.exists("DrugGEN/experiments/inference/{}".format(self.submodel)):
+            os.makedirs("DrugGEN/experiments/inference/{}".format(self.submodel))
+
+
+        self.G.eval()
+
+        start_time = time.time()
+        metric_calc_dr = []
+        uniqueness_calc = []
+        real_smiles_snn = []
+        nodes_sample = torch.Tensor(size=[1,45,1]).to(self.device)
+
+        val_counter = 0
+        none_counter = 0
+        # Inference mode
+        with torch.inference_mode():
+            pbar = tqdm(range(self.sample_num))
+            pbar.set_description('Inference mode for {} model started'.format(self.submodel))
+            for i, data in enumerate(self.inf_loader):
+
+                val_counter += 1
+                # Preprocess dataset 
+                _, a_tensor, x_tensor = load_molecules(
+                    data=data, 
+                    batch_size=self.inf_batch_size,
+                    device=self.device,
+                    b_dim=self.b_dim,
+                    m_dim=self.m_dim,
+                )
+
+                _, _, node_sample, edge_sample = self.G(a_tensor, x_tensor)
+
+                g_edges_hat_sample = torch.max(edge_sample, -1)[1]
+                g_nodes_hat_sample = torch.max(node_sample, -1)[1]
+
+                fake_mol_g = [self.inf_dataset.matrices2mol_drugs(n_.data.cpu().numpy(), e_.data.cpu().numpy(), strict=True, file_name=self.dataset_name) 
+                        for e_, n_ in zip(g_edges_hat_sample, g_nodes_hat_sample)]
+
+                a_tensor_sample = torch.max(a_tensor, -1)[1]
+                x_tensor_sample = torch.max(x_tensor, -1)[1]
+                real_mols = [self.inf_dataset.matrices2mol_drugs(n_.data.cpu().numpy(), e_.data.cpu().numpy(), strict=True, file_name=self.dataset_name) 
+                        for e_, n_ in zip(a_tensor_sample, x_tensor_sample)]
+
+                inference_drugs = [None if line is None else Chem.MolToSmiles(line) for line in fake_mol_g]
+                inference_drugs = [None if x is None else max(x.split('.'), key=len) for x in inference_drugs]
+
+                for molecules in inference_drugs:
+                            if molecules is None:
+                                none_counter += 1
+
+                with open("DrugGEN/experiments/inference/{}/inference_drugs.txt".format(self.submodel), "a") as f:
+                    for molecules in inference_drugs:
+                        if molecules is not None:
+                            molecules = molecules.replace("*", "C") 
+                            f.write(molecules)
+                            f.write("\n")
+                            uniqueness_calc.append(molecules)
+                            nodes_sample = torch.cat((nodes_sample, g_nodes_hat_sample.view(1,45,1)), 0)
+                            pbar.update(1)
+                        metric_calc_dr.append(molecules)
+
+
+                generation_number = len([x for x in metric_calc_dr if x is not None])
+                if generation_number == self.sample_num or none_counter == self.sample_num:
+                    break
+                real_smiles_snn.append(real_mols[0])
+
+        et = time.time() - start_time
+        gen_vecs = [AllChem.GetMorganFingerprintAsBitVect(Chem.MolFromSmiles(x), 2, nBits=1024) for x in uniqueness_calc if Chem.MolFromSmiles(x) is not None]
+        real_vecs = [AllChem.GetMorganFingerprintAsBitVect(x, 2, nBits=1024) for x in real_smiles_snn if x is not None]
+        print("Inference mode is lasted for {:.2f} seconds".format(et))
+
+        print("Metrics calculation started using MOSES.")
+        # post-process * to Carbon atom in valid molecules
+
+        return{
+            "Validity": fraction_valid(metric_calc_dr),
+            "Uniqueness": fraction_unique(uniqueness_calc),
+            "Novelty (Train)": novelty(metric_calc_dr, chembl_smiles),
+            "Novelty (Inference)": novelty(metric_calc_dr, chembl_test),
+            "Novelty (AKT)": novelty(metric_calc_dr, drug_smiles),
+            "MaxLen": Metrics.max_component(uniqueness_calc, self.vertexes),
+            "MeanAtomType": Metrics.mean_atom_type(nodes_sample),
+            "SNN (ChEMBL)": average_agg_tanimoto(np.array(real_vecs), np.array(gen_vecs)),
+            "SNN (AKT)": average_agg_tanimoto(np.array(drug_vecs), np.array(gen_vecs)),
+        }
+
+
+if __name__=="__main__":
+    parser = argparse.ArgumentParser()
+
+    # Inference configuration.
+    parser.add_argument('--submodel', type=str, default="DrugGEN", help="Chose model subtype: DrugGEN, NoTarget", choices=['DrugGEN', 'NoTarget'])
+    parser.add_argument('--inference_model', type=str, help="Path to the model for inference")
+    parser.add_argument('--sample_num', type=int, default=10000, help='inference samples')
+
+    # Data configuration.
+    parser.add_argument('--inf_dataset_file', type=str, default='chembl45_test.pt')
+    parser.add_argument('--inf_raw_file', type=str, default='DrugGEN/data/chembl_test.smi')
+    parser.add_argument('--inf_batch_size', type=int, default=1, help='Batch size for inference')
+    parser.add_argument('--mol_data_dir', type=str, default='DrugGEN/data')
+    parser.add_argument('--features', type=str2bool, default=False, help='features dimension for nodes')
+
+    # Model configuration.
+    parser.add_argument('--act', type=str, default="relu", help="Activation function for the model.", choices=['relu', 'tanh', 'leaky', 'sigmoid'])
+    parser.add_argument('--max_atom', type=int, default=45, help='Max atom number for molecules must be specified.')
+    parser.add_argument('--dim', type=int, default=128, help='Dimension of the Transformer Encoder model for the GAN.')
+    parser.add_argument('--depth', type=int, default=1, help='Depth of the Transformer model from the GAN.')
+    parser.add_argument('--heads', type=int, default=8, help='Number of heads for the MultiHeadAttention module from the GAN.')
+    parser.add_argument('--mlp_ratio', type=int, default=3, help='MLP ratio for the Transformer.')
+    parser.add_argument('--dropout', type=float, default=0., help='dropout rate')
+
+    # Seed configuration.
+    parser.add_argument('--set_seed', type=bool, default=False, help='set seed for reproducibility')
+    parser.add_argument('--seed', type=int, default=1, help='seed for reproducibility')
+
+    config = parser.parse_args()
+    inference = Inference(config)
+    inference.inference()

layers.py

@@ -0,0 +1,106 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+import math
+
+class MLP(nn.Module):
+    def __init__(self, in_feat, hid_feat=None, out_feat=None, dropout=0.):
+        super().__init__()
+
+        if not hid_feat:
+            hid_feat = in_feat
+        if not out_feat:
+            out_feat = in_feat
+
+        self.fc1 = nn.Linear(in_feat, hid_feat)
+        self.act = torch.nn.ReLU()
+        self.fc2 = nn.Linear(hid_feat,out_feat)
+        self.droprateout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.fc2(x)
+        return self.droprateout(x)
+
+class Attention_new(nn.Module):
+    def __init__(self, dim, heads, attention_dropout=0.):
+        super().__init__()
+
+        assert dim % heads == 0
+
+        self.heads = heads
+        self.scale = 1./dim**0.5
+        self.q = nn.Linear(dim, dim)
+        self.k = nn.Linear(dim, dim)
+        self.v = nn.Linear(dim, dim)
+        self.e = nn.Linear(dim, dim)
+        self.d_k = dim // heads
+        self.heads = heads
+        self.out_e = nn.Linear(dim,dim)
+        self.out_n = nn.Linear(dim, dim)
+
+    def forward(self, node, edge):
+        b, n, c = node.shape
+
+        q_embed = self.q(node).view(-1, n, self.heads, c//self.heads)
+        k_embed = self.k(node).view(-1, n, self.heads, c//self.heads)
+        v_embed = self.v(node).view(-1, n, self.heads, c//self.heads)
+        e_embed = self.e(edge).view(-1, n, n, self.heads, c//self.heads)
+
+        q_embed = q_embed.unsqueeze(2)
+        k_embed = k_embed.unsqueeze(1)
+
+        attn = q_embed * k_embed
+        attn = attn/ math.sqrt(self.d_k)
+        attn = attn * (e_embed + 1) * e_embed
+
+        edge = self.out_e(attn.flatten(3))
+
+        attn = F.softmax(attn, dim=2)
+
+        v_embed = v_embed.unsqueeze(1)
+        v_embed = attn * v_embed
+        v_embed = v_embed.sum(dim=2).flatten(2)
+
+        node = self.out_n(v_embed)
+        return node, edge
+
+
+class Encoder_Block(nn.Module):
+    def __init__(self, dim, heads, act, mlp_ratio=4, drop_rate=0.):
+        super().__init__()
+
+        self.ln1 = nn.LayerNorm(dim)
+        self.attn = Attention_new(dim, heads, drop_rate)
+        self.ln3 = nn.LayerNorm(dim)
+        self.ln4 = nn.LayerNorm(dim)
+        self.mlp = MLP(dim, dim*mlp_ratio, dim, dropout=drop_rate)
+        self.mlp2 = MLP(dim, dim*mlp_ratio, dim, dropout=drop_rate)
+        self.ln5 = nn.LayerNorm(dim)
+        self.ln6 = nn.LayerNorm(dim)
+
+    def forward(self, x, y):
+        x1 = self.ln1(x)
+        x2,y1 = self.attn(x1,y)
+        x2 = x1 + x2
+        y2 = y1 + y
+        x2 = self.ln3(x2)
+        y2 = self.ln4(y2)
+        x = self.ln5(x2 + self.mlp(x2))
+        y = self.ln6(y2 + self.mlp2(y2))
+        return x, y
+
+
+class TransformerEncoder(nn.Module):
+    def __init__(self, dim, depth, heads, act, mlp_ratio=4, drop_rate=0.1):
+        super().__init__()
+
+        self.Encoder_Blocks = nn.ModuleList([
+            Encoder_Block(dim, heads, act, mlp_ratio, drop_rate)
+            for i in range(depth)])
+
+    def forward(self, x, y):
+        for Encoder_Block in self.Encoder_Blocks:
+            x, y = Encoder_Block(x,y)
+        return x, y

loss.py

@@ -0,0 +1,36 @@
+import torch
+
+
+def discriminator_loss(generator, discriminator, mol_graph, batch_size, device, grad_pen, lambda_gp, z_edge, z_node):
+    # Compute loss with real molecules.
+    logits_real_disc = discriminator(mol_graph)
+    prediction_real =  - torch.mean(logits_real_disc)
+
+    # Compute loss with fake molecules.
+    node, edge, node_sample, edge_sample = generator(z_edge, z_node)
+    graph = torch.cat((node_sample.view(batch_size, -1), edge_sample.view(batch_size, -1)), dim=-1)
+    logits_fake_disc = discriminator(graph.detach())
+    prediction_fake = torch.mean(logits_fake_disc)
+
+    # Compute gradient loss.
+    eps = torch.rand(mol_graph.size(0),1).to(device)
+    x_int0 = (eps * mol_graph + (1. - eps) * graph).requires_grad_(True)
+    grad0 = discriminator(x_int0)
+    d_loss_gp = grad_pen(grad0, x_int0) 
+
+    # Calculate total loss
+    d_loss = prediction_fake + prediction_real + d_loss_gp * lambda_gp
+    return node, edge, d_loss
+
+
+def generator_loss(generator, discriminator, adj, annot, batch_size):
+    # Compute loss with fake molecules.
+    node, edge, node_sample, edge_sample = generator(adj, annot)
+
+    graph = torch.cat((node_sample.view(batch_size, -1), edge_sample.view(batch_size, -1)), dim=-1)
+
+    logits_fake_disc = discriminator(graph)
+    prediction_fake = - torch.mean(logits_fake_disc)
+    g_loss = prediction_fake
+
+    return g_loss, node, edge, node_sample, edge_sample

models.py

@@ -0,0 +1,93 @@
+import torch
+import torch.nn as nn
+from layers import TransformerEncoder
+
+class Generator(nn.Module):
+    """Generator network."""
+
+    def __init__(self, act, vertexes, edges, nodes, dropout, dim, depth, heads, mlp_ratio, submodel):
+        super(Generator, self).__init__()
+        self.submodel = submodel
+        self.vertexes = vertexes
+        self.edges = edges
+        self.nodes = nodes
+        self.depth = depth
+        self.dim = dim
+        self.heads = heads
+        self.mlp_ratio = mlp_ratio
+        self.dropout = dropout
+
+        if act == "relu":
+            act = nn.ReLU()
+        elif act == "leaky":
+            act = nn.LeakyReLU()
+        elif act == "sigmoid":
+            act = nn.Sigmoid()
+        elif act == "tanh":
+            act = nn.Tanh()
+
+        self.features = vertexes * vertexes * edges + vertexes * nodes
+        self.transformer_dim = vertexes * vertexes * dim + vertexes * dim
+        self.pos_enc_dim = 5
+
+        self.node_layers = nn.Sequential(nn.Linear(nodes, 64), act, nn.Linear(64,dim), act, nn.Dropout(self.dropout))
+        self.edge_layers = nn.Sequential(nn.Linear(edges, 64), act, nn.Linear(64,dim), act, nn.Dropout(self.dropout))
+        self.TransformerEncoder = TransformerEncoder(dim=self.dim, depth=self.depth, heads=self.heads, act = act,
+                                                                    mlp_ratio=self.mlp_ratio, drop_rate=self.dropout)
+
+        self.readout_e = nn.Linear(self.dim, edges)
+        self.readout_n = nn.Linear(self.dim, nodes)
+        self.softmax = nn.Softmax(dim = -1)
+
+    def _generate_square_subsequent_mask(self, sz):
+        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
+        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
+        return mask
+
+    def laplacian_positional_enc(self, adj):
+        A = adj
+        D = torch.diag(torch.count_nonzero(A, dim=-1))
+        L = torch.eye(A.shape[0], device=A.device) - D * A * D
+
+        EigVal, EigVec = torch.linalg.eig(L)
+        idx = torch.argsort(torch.real(EigVal))
+        EigVal, EigVec = EigVal[idx], torch.real(EigVec[:,idx])
+        pos_enc = EigVec[:,1:self.pos_enc_dim + 1]
+        return pos_enc
+
+    def forward(self, z_e, z_n):
+        b, n, c = z_n.shape
+        _, _, _ , d = z_e.shape
+
+        node = self.node_layers(z_n)
+        edge = self.edge_layers(z_e)
+        edge = (edge + edge.permute(0, 2, 1, 3)) / 2
+
+        node, edge = self.TransformerEncoder(node,edge)
+
+        node_sample = self.readout_n(node)
+        edge_sample = self.readout_e(edge)
+        return node, edge, node_sample, edge_sample
+
+
+class simple_disc(nn.Module):
+    def __init__(self, act, m_dim, vertexes, b_dim):
+        super().__init__()
+
+        if act == "relu":
+            act = nn.ReLU()
+        elif act == "leaky":
+            act = nn.LeakyReLU()
+        elif act == "sigmoid":
+            act = nn.Sigmoid()
+        elif act == "tanh":
+            act = nn.Tanh()
+
+        features = vertexes * m_dim + vertexes * vertexes * b_dim
+        self.predictor = nn.Sequential(nn.Linear(features,256), act, nn.Linear(256,128), act, nn.Linear(128,64), act,
+                                       nn.Linear(64,32), act, nn.Linear(32,16), act,
+                                       nn.Linear(16,1))
+
+    def forward(self, x):
+        prediction = self.predictor(x)
+        return prediction

new_dataloader.py

@@ -0,0 +1,311 @@
+import pickle
+import numpy as np
+import torch
+from rdkit import Chem
+from torch_geometric.data import (Data, InMemoryDataset)
+import os.path as osp
+from tqdm import tqdm
+import re
+from rdkit import RDLogger
+import pandas as pd
+
+RDLogger.DisableLog('rdApp.*') 
+class DruggenDataset(InMemoryDataset):
+    
+    def __init__(self, root, dataset_file, raw_files, max_atom, features, transform=None, pre_transform=None, pre_filter=None):
+        self.dataset_name = dataset_file.split(".")[0]
+        self.dataset_file = dataset_file
+        self.raw_files = raw_files
+        self.max_atom = max_atom
+        self.features = features
+        super().__init__(root, transform, pre_transform, pre_filter)
+        path = osp.join(self.processed_dir, dataset_file)
+        self.data, self.slices = torch.load(path)
+        self.root = root
+
+        
+    @property
+    def processed_dir(self):
+        
+        return self.root
+    
+    @property
+    def raw_file_names(self):
+        return self.raw_files
+
+    @property
+    def processed_file_names(self):
+        return self.dataset_file
+
+    def _generate_encoders_decoders(self, data):
+        
+        self.data = data
+        print('Creating atoms and bonds  encoder and decoder..')
+        
+        atom_labels = set()
+        bond_labels = set()
+        max_length = 0
+        smiles_list = []
+        for smiles in tqdm(data):
+            mol = Chem.MolFromSmiles(smiles)
+            molecule_size = mol.GetNumAtoms()
+            if molecule_size > self.max_atom:
+                continue
+            smiles_list.append(smiles)
+            atom_labels.update([atom.GetAtomicNum() for atom in mol.GetAtoms()])
+            max_length = max(max_length, molecule_size)
+            bond_labels.update([bond.GetBondType() for bond in mol.GetBonds()])
+
+        atom_labels.update([0]) # add PAD symbol (for unknown atoms)
+        atom_labels = sorted(atom_labels) # turn set into list and sort it
+
+        bond_labels = sorted(bond_labels)
+        bond_labels = [Chem.rdchem.BondType.ZERO] + bond_labels
+
+        # atom_labels = sorted(set([atom.GetAtomicNum() for mol in self.data for atom in mol.GetAtoms()] + [0]))
+        self.atom_encoder_m = {l: i for i, l in enumerate(atom_labels)}
+        self.atom_decoder_m = {i: l for i, l in enumerate(atom_labels)}
+        self.atom_num_types = len(atom_labels)
+        print('Created atoms encoder and decoder with {} atom types and 1 PAD symbol!'.format(
+            self.atom_num_types - 1))
+        print("atom_labels", atom_labels)
+        # print('Creating bonds encoder and decoder..')
+        # bond_labels = [Chem.rdchem.BondType.ZERO] + list(sorted(set(bond.GetBondType()
+        #                                                             for mol in self.data
+        #                                                             for bond in mol.GetBonds())))
+        # bond_labels = [
+        #     Chem.rdchem.BondType.ZERO,
+        #     Chem.rdchem.BondType.SINGLE,
+        #     Chem.rdchem.BondType.DOUBLE,
+        #     Chem.rdchem.BondType.TRIPLE,
+        #     Chem.rdchem.BondType.AROMATIC,
+        # ]
+
+        print("bond labels", bond_labels)
+        self.bond_encoder_m = {l: i for i, l in enumerate(bond_labels)}
+        self.bond_decoder_m = {i: l for i, l in enumerate(bond_labels)}
+        self.bond_num_types = len(bond_labels)
+        print('Created bonds encoder and decoder with {} bond types and 1 PAD symbol!'.format(
+            self.bond_num_types - 1))        
+        #dataset_names = str(self.dataset_name)
+        with open("DrugGEN/data/encoders/" +"atom_" + self.dataset_name + ".pkl","wb") as atom_encoders:
+            pickle.dump(self.atom_encoder_m,atom_encoders)
+        
+        
+        with open("DrugGEN/data/decoders/" +"atom_" +  self.dataset_name + ".pkl","wb") as atom_decoders:
+            pickle.dump(self.atom_decoder_m,atom_decoders)
+        
+               
+        with open("DrugGEN/data/encoders/" +"bond_" +  self.dataset_name + ".pkl","wb") as bond_encoders:
+            pickle.dump(self.bond_encoder_m,bond_encoders)
+        
+              
+        with open("DrugGEN/data/decoders/" +"bond_" +  self.dataset_name + ".pkl","wb") as bond_decoders:
+            pickle.dump(self.bond_decoder_m,bond_decoders)
+         
+        return max_length, smiles_list # data is filtered now
+        
+    def _genA(self, mol, connected=True, max_length=None):
+
+        max_length = max_length if max_length is not None else mol.GetNumAtoms()
+
+        A = np.zeros(shape=(max_length, max_length))
+
+        begin, end = [b.GetBeginAtomIdx() for b in mol.GetBonds()], [b.GetEndAtomIdx() for b in mol.GetBonds()]
+        bond_type = [self.bond_encoder_m[b.GetBondType()] for b in mol.GetBonds()]
+
+        A[begin, end] = bond_type
+        A[end, begin] = bond_type
+
+        degree = np.sum(A[:mol.GetNumAtoms(), :mol.GetNumAtoms()], axis=-1)
+
+        return A if connected and (degree > 0).all() else None
+
+    def _genX(self, mol, max_length=None):
+
+        max_length = max_length if max_length is not None else mol.GetNumAtoms()
+
+        return np.array([self.atom_encoder_m[atom.GetAtomicNum()] for atom in mol.GetAtoms()] + [0] * (
+                    max_length - mol.GetNumAtoms()))
+
+    def _genF(self, mol, max_length=None):
+
+        max_length = max_length if max_length is not None else mol.GetNumAtoms()
+
+        features = np.array([[*[a.GetDegree() == i for i in range(5)],
+                              *[a.GetExplicitValence() == i for i in range(9)],
+                              *[int(a.GetHybridization()) == i for i in range(1, 7)],
+                              *[a.GetImplicitValence() == i for i in range(9)],
+                              a.GetIsAromatic(),
+                              a.GetNoImplicit(),
+                              *[a.GetNumExplicitHs() == i for i in range(5)],
+                              *[a.GetNumImplicitHs() == i for i in range(5)],
+                              *[a.GetNumRadicalElectrons() == i for i in range(5)],
+                              a.IsInRing(),
+                              *[a.IsInRingSize(i) for i in range(2, 9)]] for a in mol.GetAtoms()], dtype=np.int32)
+
+        return np.vstack((features, np.zeros((max_length - features.shape[0], features.shape[1]))))
+
+    def decoder_load(self, dictionary_name, file):
+        with open("DrugGEN/data/decoders/" + dictionary_name + "_" + file + '.pkl', 'rb') as f:
+            return pickle.load(f)    
+        
+    def drugs_decoder_load(self, dictionary_name):
+        with open("DrugGEN/data/decoders/" + dictionary_name +'.pkl', 'rb') as f:
+            return pickle.load(f)      
+        
+    def matrices2mol(self, node_labels, edge_labels, strict=True, file_name=None):
+        mol = Chem.RWMol()
+        RDLogger.DisableLog('rdApp.*') 
+        atom_decoders = self.decoder_load("atom", file_name)
+        bond_decoders = self.decoder_load("bond", file_name)
+        
+        for node_label in node_labels:
+            mol.AddAtom(Chem.Atom(atom_decoders[node_label]))
+
+        for start, end in zip(*np.nonzero(edge_labels)):
+            if start > end:
+                mol.AddBond(int(start), int(end), bond_decoders[edge_labels[start, end]])
+        #mol = self.correct_mol(mol)
+        if strict:
+            try:
+                
+                Chem.SanitizeMol(mol)
+            except:
+                mol = None
+
+        return mol
+    
+    def drug_decoder_load(self, dictionary_name, file):
+        
+        ''' Loading the atom and bond decoders '''
+        
+        with open("DrugGEN/data/decoders/" + dictionary_name +"_" + file +'.pkl', 'rb') as f:
+            
+            return pickle.load(f) 
+    def matrices2mol_drugs(self, node_labels, edge_labels, strict=True, file_name=None):
+        mol = Chem.RWMol()
+        RDLogger.DisableLog('rdApp.*') 
+        atom_decoders = self.drug_decoder_load("atom", file_name)
+        bond_decoders = self.drug_decoder_load("bond", file_name)
+        
+        for node_label in node_labels:
+            
+            mol.AddAtom(Chem.Atom(atom_decoders[node_label]))
+
+        for start, end in zip(*np.nonzero(edge_labels)):
+            if start > end:
+                mol.AddBond(int(start), int(end), bond_decoders[edge_labels[start, end]])
+        #mol = self.correct_mol(mol)
+        if strict:
+            try:
+                Chem.SanitizeMol(mol)
+            except:
+                mol = None
+
+        return mol
+    def check_valency(self,mol):
+        """
+        Checks that no atoms in the mol have exceeded their possible
+        valency
+        :return: True if no valency issues, False otherwise
+        """
+        try:
+            Chem.SanitizeMol(mol, sanitizeOps=Chem.SanitizeFlags.SANITIZE_PROPERTIES)
+            return True, None
+        except ValueError as e:
+            e = str(e)
+            p = e.find('#')
+            e_sub = e[p:]
+            atomid_valence = list(map(int, re.findall(r'\d+', e_sub)))
+            return False, atomid_valence
+
+
+    def correct_mol(self,x):
+        xsm = Chem.MolToSmiles(x, isomericSmiles=True)
+        mol = x
+        while True:
+            flag, atomid_valence = self.check_valency(mol)
+            if flag:
+                break
+            else:
+                assert len (atomid_valence) == 2
+                idx = atomid_valence[0]
+                v = atomid_valence[1]
+                queue = []
+                for b in mol.GetAtomWithIdx(idx).GetBonds():
+                    queue.append(
+                        (b.GetIdx(), int(b.GetBondType()), b.GetBeginAtomIdx(), b.GetEndAtomIdx())
+                    )
+                queue.sort(key=lambda tup: tup[1], reverse=True)
+                if len(queue) > 0:
+                    start = queue[0][2]
+                    end = queue[0][3]
+                    t = queue[0][1] - 1
+                    mol.RemoveBond(start, end)
+        
+                    #if t >= 1:
+                        
+                        #mol.AddBond(start, end, self.decoder_load('bond_decoders')[t])
+                    # if '.' in Chem.MolToSmiles(mol, isomericSmiles=True):
+                    #     mol.AddBond(start, end, self.decoder_load('bond_decoders')[t])
+                    #     print(tt)
+                    #     print(Chem.MolToSmiles(mol, isomericSmiles=True))
+
+        return mol
+    
+    
+        
+    def label2onehot(self, labels, dim):
+        
+        """Convert label indices to one-hot vectors."""
+        
+        out = torch.zeros(list(labels.size())+[dim])
+        out.scatter_(len(out.size())-1,labels.unsqueeze(-1),1.)
+        
+        return out.float() 
+       
+    def process(self, size= None):
+        smiles_list = pd.read_csv(self.raw_files, header=None)[0].tolist()
+        max_length, smiles_list = self._generate_encoders_decoders(smiles_list)
+
+        data_list = []
+      
+        self.m_dim = len(self.atom_decoder_m)
+        for smiles in tqdm(smiles_list, desc='Processing chembl dataset', total=len(smiles_list)):
+            mol = Chem.MolFromSmiles(smiles)
+            A = self._genA(mol, connected=True, max_length=max_length)
+            if A is not None:
+                
+
+                x = torch.from_numpy(self._genX(mol, max_length=max_length)).to(torch.long).view(1, -1)
+          
+                x = self.label2onehot(x,self.m_dim).squeeze()
+                if self.features: 
+                    f = torch.from_numpy(self._genF(mol, max_length=max_length)).to(torch.long).view(x.shape[0], -1)
+                    x = torch.concat((x,f), dim=-1)
+             
+                adjacency = torch.from_numpy(A)
+                
+                edge_index = adjacency.nonzero(as_tuple=False).t().contiguous()
+                edge_attr = adjacency[edge_index[0], edge_index[1]].to(torch.long)
+
+                data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr, smiles=smiles)
+
+                if self.pre_filter is not None and not self.pre_filter(data):
+                    continue
+
+                if self.pre_transform is not None:
+                    data = self.pre_transform(data)
+                    
+                data_list.append(data)
+
+
+        torch.save(self.collate(data_list), osp.join(self.processed_dir, self.dataset_file))
+
+
+
+   
+if __name__ == '__main__':
+    data = DruggenDataset("DrugGEN/data")
+    

packages.txt

@@ -0,0 +1 @@
+libcairo2-dev

requirements.txt

@@ -0,0 +1,12 @@
+torch
+rdkit-pypi
+tqdm
+numpy
+seaborn
+matplotlib
+pandas
+torch_geometric
+# demo related installs
+streamlit
+ipython
+streamlit-ext

training_data.py

@@ -0,0 +1,31 @@
+import torch
+import torch_geometric.utils as geoutils
+from utils import label2onehot
+
+def generate_z_values(batch_size=32, z_dim=32, vertexes=32, b_dim=32, m_dim=32, device=None):
+    z = torch.normal(mean=0, std=1, size=(batch_size, z_dim), device=device) # (batch,max_len)
+    z_edge = torch.normal(mean=0, std=1, size=(batch_size, vertexes, vertexes, b_dim), device=device) # (batch,max_len,max_len)
+    z_node = torch.normal(mean=0, std=1, size=(batch_size, vertexes, m_dim), device=device) # (batch,max_len)
+
+    z = z.float().requires_grad_(True)
+    z_edge = z_edge.float().requires_grad_(True)                                      # Edge noise.(batch,max_len,max_len)
+    z_node = z_node.float().requires_grad_(True)                                      # Node noise.(batch,max_len)
+    return z, z_edge, z_node
+
+
+def load_molecules(data=None, b_dim=32, m_dim=32, device=None, batch_size=32):
+    data = data.to(device)
+    a = geoutils.to_dense_adj(
+        edge_index = data.edge_index,
+        batch=data.batch,
+        edge_attr=data.edge_attr,
+        max_num_nodes=int(data.batch.shape[0]/batch_size)
+    )
+    x_tensor = data.x.view(batch_size,int(data.batch.shape[0]/batch_size),-1)
+    a_tensor = label2onehot(a, b_dim, device)
+
+    a_tensor_vec = a_tensor.reshape(batch_size,-1)
+    x_tensor_vec = x_tensor.reshape(batch_size,-1)
+    real_graphs = torch.concat((x_tensor_vec,a_tensor_vec),dim=-1)
+
+    return real_graphs, a_tensor, x_tensor

utils.py

@@ -0,0 +1,421 @@
+from statistics import mean
+import os
+import math
+import time
+import datetime
+from rdkit import DataStructs
+from rdkit import Chem
+from rdkit import RDLogger
+from rdkit.Chem import AllChem
+from rdkit.Chem import Draw
+from rdkit.Chem.Scaffolds import MurckoScaffold
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.lines import Line2D
+import torch 
+import wandb
+RDLogger.DisableLog('rdApp.*')
+import warnings
+from multiprocessing import Pool
+class Metrics(object):
+
+    @staticmethod
+    def valid(x):
+        return x is not None and Chem.MolToSmiles(x) != ''
+
+    @staticmethod
+    def tanimoto_sim_1v2(data1, data2):
+        min_len = data1.size if data1.size > data2.size else data2
+        sims = []
+        for i in range(min_len):
+            sim = DataStructs.FingerprintSimilarity(data1[i], data2[i])
+            sims.append(sim)
+        mean_sim = mean(sim)
+        return mean_sim
+
+    @staticmethod
+    def mol_length(x):
+        if x is not None:
+            return  len([char for char in max(x.split(sep =".")).upper() if char.isalpha()])
+        else:
+            return 0
+    
+    @staticmethod
+    def max_component(data, max_len):
+        
+        return ((np.array(list(map(Metrics.mol_length, data)), dtype=np.float32)/max_len).mean())
+
+    @staticmethod
+    def mean_atom_type(data):
+        atom_types_used = []
+        for i in data:
+
+            atom_types_used.append(len(i.unique().tolist()))
+        av_type = np.mean(atom_types_used) - 1
+        
+        return av_type
+
+
+def sim_reward(mol_gen, fps_r):
+    
+    gen_scaf = []
+    
+    for x in mol_gen: 
+        if x is not None:
+            try:
+                
+                gen_scaf.append(MurckoScaffold.GetScaffoldForMol(x))
+            except:
+                pass
+
+    if len(gen_scaf) == 0:
+        
+        rew = 1
+    else:
+        fps = [Chem.RDKFingerprint(x) for x in gen_scaf]
+            
+        
+        fps = np.array(fps)
+        fps_r = np.array(fps_r)
+    
+        rew =  average_agg_tanimoto(fps_r,fps)
+        if math.isnan(rew):
+            rew = 1
+    
+    return rew  ## change this to penalty
+
+##########################################
+##########################################
+##########################################
+
+def mols2grid_image(mols,path):
+    mols = [e if e is not None else Chem.RWMol() for e in mols]
+    
+    for i in range(len(mols)):
+        if Metrics.valid(mols[i]):
+            AllChem.Compute2DCoords(mols[i])
+            Draw.MolToFile(mols[i], os.path.join(path,"{}.png".format(i+1)), size=(1200,1200))
+            #wandb.save(os.path.join(path,"{}.png".format(i+1)))
+        else:
+            continue
+
+def save_smiles_matrices(mols,edges_hard, nodes_hard, path, data_source = None): 
+    mols = [e if e is not None else Chem.RWMol() for e in mols]
+    
+    for i in range(len(mols)):
+        if Metrics.valid(mols[i]):
+            save_path = os.path.join(path,"{}.txt".format(i+1))
+            with open(save_path, "a") as f:
+                np.savetxt(f, edges_hard[i].cpu().numpy(), header="edge matrix:\n",fmt='%1.2f')
+                f.write("\n")
+                np.savetxt(f, nodes_hard[i].cpu().numpy(), header="node matrix:\n", footer="\nsmiles:",fmt='%1.2f')
+                f.write("\n")
+                #f.write(m0)
+                f.write("\n")
+            print(Chem.MolToSmiles(mols[i]), file=open(save_path,"a"))
+            #wandb.save(save_path)
+        else:
+            continue
+
+
+##########################################
+##########################################
+##########################################
+
+
+def dense_to_sparse_with_attr(adj):
+    assert adj.dim() >= 2 and adj.dim() <= 3
+    assert adj.size(-1) == adj.size(-2)
+
+    index = adj.nonzero(as_tuple=True)
+    edge_attr = adj[index]
+
+    if len(index) == 3:
+        batch = index[0] * adj.size(-1)
+        index = (batch + index[1], batch + index[2])
+        #index = torch.stack(index, dim=0)
+    return index, edge_attr
+
+
+def label2onehot(labels, dim, device):
+    """Convert label indices to one-hot vectors."""
+    out = torch.zeros(list(labels.size())+[dim]).to(device)
+    out.scatter_(len(out.size())-1,labels.unsqueeze(-1),1.)
+
+    return out.float()
+
+
+def mol_sample(sample_directory, edges, nodes, idx, i,matrices2mol, dataset_name):
+    sample_path = os.path.join(sample_directory,"{}_{}-epoch_iteration".format(idx+1, i+1))
+    g_edges_hat_sample = torch.max(edges, -1)[1]
+    g_nodes_hat_sample = torch.max(nodes , -1)[1]
+    mol = [matrices2mol(n_.data.cpu().numpy(), e_.data.cpu().numpy(), strict=True, file_name=dataset_name)
+            for e_, n_ in zip(g_edges_hat_sample, g_nodes_hat_sample)]
+
+    if not os.path.exists(sample_path):
+        os.makedirs(sample_path)
+
+    mols2grid_image(mol,sample_path)
+    save_smiles_matrices(mol,g_edges_hat_sample.detach(), g_nodes_hat_sample.detach(), sample_path)
+
+    if len(os.listdir(sample_path)) == 0:
+        os.rmdir(sample_path)
+
+    print("Valid molecules are saved.")
+    print("Valid matrices and smiles are saved")
+
+
+def logging(log_path, start_time, i, idx, loss, save_path, drug_smiles, edge, node, 
+            matrices2mol, dataset_name, real_adj, real_annot, drug_vecs):
+
+    g_edges_hat_sample = torch.max(edge, -1)[1]
+    g_nodes_hat_sample = torch.max(node , -1)[1]
+
+    a_tensor_sample = torch.max(real_adj, -1)[1].float()
+    x_tensor_sample = torch.max(real_annot, -1)[1].float()
+
+    mols = [matrices2mol(n_.data.cpu().numpy(), e_.data.cpu().numpy(), strict=True, file_name=dataset_name)
+            for e_, n_ in zip(g_edges_hat_sample, g_nodes_hat_sample)]
+
+    real_mol = [matrices2mol(n_.data.cpu().numpy(), e_.data.cpu().numpy(), strict=True, file_name=dataset_name)
+            for e_, n_ in zip(a_tensor_sample, x_tensor_sample)]
+
+    atom_types_average = Metrics.mean_atom_type(g_nodes_hat_sample)
+    real_smiles = [Chem.MolToSmiles(x) for x in real_mol if x is not None]
+    gen_smiles = []
+    uniq_smiles = []
+    for line in mols:
+        if line is not None:
+            gen_smiles.append(Chem.MolToSmiles(line))
+            uniq_smiles.append(Chem.MolToSmiles(line))
+        elif line is None:
+            gen_smiles.append(None)
+
+    gen_smiles_saves = [None if x is None else max(x.split('.'), key=len) for x in gen_smiles]
+    uniq_smiles_saves = [None if x is None else max(x.split('.'), key=len) for x in uniq_smiles]
+
+    sample_save_dir = os.path.join(save_path, "samples.txt")
+    with open(sample_save_dir, "a") as f:
+        for idxs in range(len(gen_smiles_saves)):
+            if gen_smiles_saves[idxs] is not None:
+                f.write(gen_smiles_saves[idxs])
+                f.write("\n")
+
+    k = len(set(uniq_smiles_saves) - {None})
+    et = time.time() - start_time
+    et = str(datetime.timedelta(seconds=et))[:-7]
+    log = "Elapsed [{}], Epoch/Iteration [{}/{}]".format(et, idx, i+1)
+    gen_vecs = [AllChem.GetMorganFingerprintAsBitVect(x, 2, nBits=1024) for x in mols if x is not None]
+    chembl_vecs = [AllChem.GetMorganFingerprintAsBitVect(x, 2, nBits=1024) for x in real_mol if x is not None]
+
+    # Log update
+    #m0 = get_all_metrics(gen = gen_smiles, train = train_smiles, batch_size=batch_size, k = valid_mol_num, device=self.device)
+    valid = fraction_valid(gen_smiles_saves)
+    unique = fraction_unique(uniq_smiles_saves, k, check_validity=False)
+    novel_starting_mol = novelty(gen_smiles_saves, real_smiles)
+    novel_akt = novelty(gen_smiles_saves, drug_smiles)
+    if (len(uniq_smiles_saves) == 0):
+        snn_chembl = 0
+        snn_akt = 0
+        maxlen = 0
+    else:
+        snn_chembl = average_agg_tanimoto(np.array(chembl_vecs),np.array(gen_vecs))
+        snn_akt = average_agg_tanimoto(np.array(drug_vecs),np.array(gen_vecs))
+        maxlen = Metrics.max_component(uniq_smiles_saves, 45)
+
+    loss.update({'Validity': valid})
+    loss.update({'Uniqueness': unique})
+    loss.update({'Novelty': novel_starting_mol})
+    loss.update({'Novelty_akt': novel_akt})
+    loss.update({'SNN_chembl': snn_chembl})
+    loss.update({'SNN_akt': snn_akt})
+    loss.update({'MaxLen': maxlen})
+    loss.update({'Atom_types': atom_types_average})
+
+    wandb.log({"Validity": valid, "Uniqueness": unique, "Novelty": novel_starting_mol,
+                "Novelty_akt": novel_akt, "SNN_chembl": snn_chembl, "SNN_akt": snn_akt,
+                  "MaxLen": maxlen, "Atom_types": atom_types_average})
+
+    for tag, value in loss.items():
+        log += ", {}: {:.4f}".format(tag, value)
+    with open(log_path, "a") as f:
+        f.write(log)
+        f.write("\n")
+    print(log)
+    print("\n")
+
+
+def plot_grad_flow(named_parameters, model, itera, epoch,grad_flow_directory):
+    # Based on https://discuss.pytorch.org/t/check-gradient-flow-in-network/15063/10
+    '''Plots the gradients flowing through different layers in the net during training.
+    Can be used for checking for possible gradient vanishing / exploding problems.
+    
+    Usage: Plug this function in Trainer class after loss.backwards() as 
+    "plot_grad_flow(self.model.named_parameters())" to visualize the gradient flow'''
+    ave_grads = []
+    max_grads= []
+    layers = []
+    for n, p in named_parameters:
+        if(p.requires_grad) and ("bias" not in n):
+            #print(p.grad,n)
+            layers.append(n)
+            ave_grads.append(p.grad.abs().mean().cpu())
+            max_grads.append(p.grad.abs().max().cpu())
+    plt.bar(np.arange(len(max_grads)), max_grads, alpha=0.1, lw=1, color="c")
+    plt.bar(np.arange(len(max_grads)), ave_grads, alpha=0.1, lw=1, color="b")
+    plt.hlines(0, 0, len(ave_grads)+1, lw=2, color="k" )
+    plt.xticks(range(0,len(ave_grads), 1), layers, rotation="vertical")
+    plt.xlim(left=0, right=len(ave_grads))
+    plt.ylim(bottom = -0.001, top=1) # zoom in on the lower gradient regions
+    plt.xlabel("Layers")
+    plt.ylabel("average gradient")
+    plt.title("Gradient flow")
+    plt.grid(True)
+    plt.legend([Line2D([0], [0], color="c", lw=4),
+                Line2D([0], [0], color="b", lw=4),
+                Line2D([0], [0], color="k", lw=4)], ['max-gradient', 'mean-gradient', 'zero-gradient'])
+    pltsavedir = grad_flow_directory
+    plt.savefig(os.path.join(pltsavedir, "weights_" + model  + "_"  + str(itera) + "_" + str(epoch) +  ".png"), dpi= 500,bbox_inches='tight')
+
+
+def get_mol(smiles_or_mol):
+    '''
+    Loads SMILES/molecule into RDKit's object
+    '''
+    if isinstance(smiles_or_mol, str):
+        if len(smiles_or_mol) == 0:
+            return None
+        mol = Chem.MolFromSmiles(smiles_or_mol)
+        if mol is None:
+            return None
+        try:
+            Chem.SanitizeMol(mol)
+        except ValueError:
+            return None
+        return mol
+    return smiles_or_mol
+
+
+def mapper(n_jobs):
+    '''
+    Returns function for map call.
+    If n_jobs == 1, will use standard map
+    If n_jobs > 1, will use multiprocessing pool
+    If n_jobs is a pool object, will return its map function
+    '''
+    if n_jobs == 1:
+        def _mapper(*args, **kwargs):
+            return list(map(*args, **kwargs))
+
+        return _mapper
+    if isinstance(n_jobs, int):
+        pool = Pool(n_jobs)
+
+        def _mapper(*args, **kwargs):
+            try:
+                result = pool.map(*args, **kwargs)
+            finally:
+                pool.terminate()
+            return result
+
+        return _mapper
+    return n_jobs.map
+
+
+def remove_invalid(gen, canonize=True, n_jobs=1):
+    """
+    Removes invalid molecules from the dataset
+    """
+    if not canonize:
+        mols = mapper(n_jobs)(get_mol, gen)
+        return [gen_ for gen_, mol in zip(gen, mols) if mol is not None]
+    return [x for x in mapper(n_jobs)(canonic_smiles, gen) if
+            x is not None]
+
+
+def fraction_valid(gen, n_jobs=1):
+    """
+    Computes a number of valid molecules
+    Parameters:
+        gen: list of SMILES
+        n_jobs: number of threads for calculation
+    """
+    gen = mapper(n_jobs)(get_mol, gen)
+    return 1 - gen.count(None) / len(gen)
+def canonic_smiles(smiles_or_mol):
+    mol = get_mol(smiles_or_mol)
+    if mol is None:
+        return None
+    return Chem.MolToSmiles(mol)
+def fraction_unique(gen, k=None, n_jobs=1, check_validity=False):
+    """
+    Computes a number of unique molecules
+    Parameters:
+        gen: list of SMILES
+        k: compute unique@k
+        n_jobs: number of threads for calculation
+        check_validity: raises ValueError if invalid molecules are present
+    """
+    if k is not None:
+        if len(gen) < k:
+            warnings.warn(
+                "Can't compute unique@{}.".format(k) +
+                "gen contains only {} molecules".format(len(gen))
+            )
+        gen = gen[:k]
+    canonic = set(mapper(n_jobs)(canonic_smiles, gen))
+    if None in canonic and check_validity:
+        #canonic = [i for i in canonic if i is not None]
+        raise ValueError("Invalid molecule passed to unique@k")
+    return 0 if len(gen) == 0 else len(canonic) / len(gen)
+
+def novelty(gen, train, n_jobs=1):
+    gen_smiles = mapper(n_jobs)(canonic_smiles, gen)
+    gen_smiles_set = set(gen_smiles) - {None}
+    train_set = set(train)
+    return 0 if len(gen_smiles_set) == 0 else len(gen_smiles_set - train_set) / len(gen_smiles_set)
+
+
+
+def average_agg_tanimoto(stock_vecs, gen_vecs,
+                         batch_size=5000, agg='max',
+                         device='cpu', p=1):
+    """
+    For each molecule in gen_vecs finds closest molecule in stock_vecs.
+    Returns average tanimoto score for between these molecules
+
+    Parameters:
+        stock_vecs: numpy array <n_vectors x dim>
+        gen_vecs: numpy array <n_vectors' x dim>
+        agg: max or mean
+        p: power for averaging: (mean x^p)^(1/p)
+    """
+    assert agg in ['max', 'mean'], "Can aggregate only max or mean"
+    agg_tanimoto = np.zeros(len(gen_vecs))
+    total = np.zeros(len(gen_vecs))
+    for j in range(0, stock_vecs.shape[0], batch_size):
+        x_stock = torch.tensor(stock_vecs[j:j + batch_size]).to(device).float()
+        for i in range(0, gen_vecs.shape[0], batch_size):
+            
+            y_gen = torch.tensor(gen_vecs[i:i + batch_size]).to(device).float()
+            y_gen = y_gen.transpose(0, 1)
+            tp = torch.mm(x_stock, y_gen)
+            jac = (tp / (x_stock.sum(1, keepdim=True) +
+                         y_gen.sum(0, keepdim=True) - tp)).cpu().numpy()
+            jac[np.isnan(jac)] = 1
+            if p != 1:
+                jac = jac**p
+            if agg == 'max':
+                agg_tanimoto[i:i + y_gen.shape[1]] = np.maximum(
+                    agg_tanimoto[i:i + y_gen.shape[1]], jac.max(0))
+            elif agg == 'mean':
+                agg_tanimoto[i:i + y_gen.shape[1]] += jac.sum(0)
+                total[i:i + y_gen.shape[1]] += jac.shape[0]
+    if agg == 'mean':
+        agg_tanimoto /= total
+    if p != 1:
+        agg_tanimoto = (agg_tanimoto)**(1/p)
+    return np.mean(agg_tanimoto)
+
+def str2bool(v):
+    return v.lower() in ('true')