this is the source code of PointGAT
'''
Description: 
version: 
Email: rongzhangthu@yeah.net
Author: Rong Zhang
Date: 2022-10-19 08:31:57
LastEditors: Rong Zhang
LastEditTime: 2022-11-04 14:51:18
'''
from rdkit.Chem import MolFromSmiles
from rdkit.Chem import Draw
import numpy as np
import rdkit
from rdkit import Chem
from rdkit.Chem import AllChem


def one_of_k_encoding(x, allowable_set):
    if x not in allowable_set:
        raise Exception("input {0} not in allowable set{1}:".format(
            x, allowable_set))
    return [x == s for s in allowable_set]


def one_of_k_encoding_unk(x, allowable_set):
    """Maps inputs not in the allowable set to the last element."""
    if x not in allowable_set:
        x = allowable_set[-1]
    return [x == s for s in allowable_set]


def atom_features(atom,
                  bool_id_feat=False,
                  explicit_H=False,
                  use_chirality=True):
    if bool_id_feat:
        return np.array([atom_to_id(atom)])
    else:
        results = one_of_k_encoding_unk(
          atom.GetSymbol(),
          [
            'C'
          ]) + one_of_k_encoding(atom.GetDegree(),
                                 [0, 1, 2, 3,4]) + \
                  [atom.GetFormalCharge()] + [atom.IsInRing()] +\
                  one_of_k_encoding_unk(atom.GetHybridization(), [
                    Chem.rdchem.HybridizationType.SP2,
                    Chem.rdchem.HybridizationType.SP3
                  ])
        # In case of explicit hydrogen(QM8, QM9), avoid calling `GetTotalNumHs`
        if not explicit_H:
            results = results + one_of_k_encoding_unk(atom.GetTotalNumHs(),
                                                      [0, 1, 2, 3])
        if use_chirality:
            try:
                results = results + one_of_k_encoding_unk(
                    atom.GetProp('_CIPCode'),
                    ['R', 'S']) + [atom.HasProp('_ChiralityPossible')]
            except:
                results = results + [False, False
                                     ] + [atom.HasProp('_ChiralityPossible')]

        return np.array(results)


def bond_features(bond, use_chirality=True):
    bt = bond.GetBondType()
    bond_feats = [
        bt == Chem.rdchem.BondType.SINGLE, bt == Chem.rdchem.BondType.DOUBLE,
        bond.GetIsConjugated(),
        bond.IsInRing()
    ]
    if use_chirality:
        bond_feats = bond_feats + one_of_k_encoding_unk(
            str(bond.GetStereo()),
            ["STEREONONE", "STEREOANY", "STEREOZ", "STEREOE"])
    return np.array(bond_feats)


def num_atom_features():
    # Return length of feature vector using a very simple molecule.
    m = Chem.MolFromSmiles('CC')
    alist = m.GetAtoms()
    a = alist[0]
    return len(atom_features(a))


def num_bond_features():
    # Return length of feature vector using a very simple molecule.
    simple_mol = Chem.MolFromSmiles('CC')
    Chem.SanitizeMol(simple_mol)
    return len(bond_features(simple_mol.GetBonds()[0]))

src/PointGAT/Featurizer.py
sevencheung2021/PointGAT
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

class PointGAT_viz(nn.Module):

    def __init__(self, radius, T, input_feature_dim, input_bond_dim,\
            fingerprint_dim, output_units_num, p_dropout):
        super(PointGAT_viz, self).__init__()
        self.atom_fc = nn.Linear(input_feature_dim, fingerprint_dim)
        self.neighbor_fc = nn.Linear(input_feature_dim+input_bond_dim, fingerprint_dim)
        self.GRUCell = nn.ModuleList([nn.GRUCell(fingerprint_dim, fingerprint_dim) for r in range(radius)])
        self.align = nn.ModuleList([nn.Linear(2*fingerprint_dim,1) for r in range(radius)])
        self.attend = nn.ModuleList([nn.Linear(fingerprint_dim, fingerprint_dim) for r in range(radius)])
        
        self.mol_GRUCell = nn.GRUCell(fingerprint_dim, fingerprint_dim)
        self.mol_align = nn.Linear(2*fingerprint_dim,1)
        self.mol_attend = nn.Linear(fingerprint_dim, fingerprint_dim)
        
#         self.mol_GRUCell = nn.ModuleList([nn.GRUCell(fingerprint_dim, fingerprint_dim) for t in range(T)])
#         self.mol_align = nn.ModuleList([nn.Linear(2*fingerprint_dim,1) for t in range(T)])
#         self.mol_attend = nn.ModuleList([nn.Linear(fingerprint_dim, fingerprint_dim) for t in range(T)])
        
        self.dropout = nn.Dropout(p=p_dropout)
        self.output = nn.Linear(fingerprint_dim, output_units_num)

        self.radius = radius
        self.T = T


    def forward(self, atom_list, bond_list, atom_degree_list, bond_degree_list, atom_mask):
        atom_mask = atom_mask.unsqueeze(2)
        batch_size,mol_length,num_atom_feat = atom_list.size()
        atom_feature = F.leaky_relu(self.atom_fc(atom_list))

        atom_feature_viz = []
        atom_feature_viz.append(self.atom_fc(atom_list))

        bond_neighbor = [bond_list[i][bond_degree_list[i]] for i in range(batch_size)]
        bond_neighbor = torch.stack(bond_neighbor, dim=0)
        atom_neighbor = [atom_list[i][atom_degree_list[i]] for i in range(batch_size)]
        atom_neighbor = torch.stack(atom_neighbor, dim=0)
        #then catenate them
        neighbor_feature = torch.cat([atom_neighbor, bond_neighbor],dim=-1)
        neighbor_feature = F.leaky_relu(self.neighbor_fc(neighbor_feature))

        # generate mask to eliminate the influence of blank atoms
        attend_mask = atom_degree_list.clone()
        attend_mask[attend_mask != mol_length-1] = 1
        attend_mask[attend_mask == mol_length-1] = 0
        attend_mask = attend_mask.type(torch.cuda.FloatTensor).unsqueeze(-1)

        softmax_mask = atom_degree_list.clone()
        softmax_mask[softmax_mask != mol_length-1] = 0
        softmax_mask[softmax_mask == mol_length-1] = -9e8 # make the softmax value extremly small
        softmax_mask = softmax_mask.type(torch.cuda.FloatTensor).unsqueeze(-1)

        batch_size, mol_length, max_neighbor_num, fingerprint_dim = neighbor_feature.shape
        atom_feature_expand = atom_feature.unsqueeze(-2).expand(batch_size, mol_length, max_neighbor_num, fingerprint_dim)
        feature_attention = torch.cat([atom_feature_expand, neighbor_feature],dim=-1)
        
        align_score = self.dropout(F.leaky_relu(self.align[0](feature_attention)))
#             print(align_score)
        align_score = align_score + softmax_mask
        attention_weight = F.softmax(align_score,-2)
#             print(align_score)
        attention_weight = attention_weight * attend_mask
#         print(align_score)
        atom_attention_weight_viz = []
        atom_attention_weight_viz.append(attention_weight)

        neighbor_feature_transform = self.attend[0](self.dropout(neighbor_feature))
#             print(features_neighbor_transform.shape)
        context = torch.sum(torch.mul(attention_weight,neighbor_feature_transform),-2)
#             print(context.shape)
        context = F.elu(context)
        context_reshape = context.view(batch_size*mol_length, fingerprint_dim)
        atom_feature_reshape = atom_feature.view(batch_size*mol_length, fingerprint_dim)
        atom_feature_reshape = self.GRUCell[0](context_reshape, atom_feature_reshape)
        atom_feature = atom_feature_reshape.view(batch_size, mol_length, fingerprint_dim)
#         atom_feature_viz.append(atom_feature)

        #do nonlinearity
        activated_features = F.relu(atom_feature)
        atom_feature_viz.append(activated_features)
        
        for d in range(self.radius-1):
            # bonds_indexed = [bond_list[i][torch.cuda.LongTensor(bond_degree_list)[i]] for i in range(batch_size)]
            neighbor_feature = [activated_features[i][atom_degree_list[i]] for i in range(batch_size)]
            
            # neighbor_feature is a list of 3D tensor, so we need to stack them into a 4D tensor first
            neighbor_feature = torch.stack(neighbor_feature, dim=0)
            atom_feature_expand = activated_features.unsqueeze(-2).expand(batch_size, mol_length, max_neighbor_num, fingerprint_dim)

            feature_attention = torch.cat([atom_feature_expand, neighbor_feature],dim=-1)

            align_score = self.dropout(F.leaky_relu(self.align[d+1](feature_attention)))
    #             print(align_score)
            align_score = align_score + softmax_mask
            attention_weight = F.softmax(align_score,-2)
#             print(align_score)
            attention_weight = attention_weight * attend_mask
            atom_attention_weight_viz.append(attention_weight)
    
#             print(align_score)
            neighbor_feature_transform = self.attend[d+1](self.dropout(neighbor_feature))
    #             print(features_neighbor_transform.shape)
            context = torch.sum(torch.mul(attention_weight,neighbor_feature_transform),-2)
    #             print(context.shape)
            context = F.elu(context)
            context_reshape = context.view(batch_size*mol_length, fingerprint_dim)
            atom_feature_reshape = self.GRUCell[d+1](context_reshape, atom_feature_reshape)
            atom_feature = atom_feature_reshape.view(batch_size, mol_length, fingerprint_dim)
#             atom_feature_viz.append(atom_feature)
            
            #do nonlinearity
            activated_features = F.relu(atom_feature)
            atom_feature_viz.append(activated_features)

        
        # when the descriptor value are unbounded, like partial charge or LogP
        mol_feature_unbounded_viz = []
        mol_feature_unbounded_viz.append(torch.sum(atom_feature * atom_mask, dim=-2)) 
        
        mol_feature = torch.sum(activated_features * atom_mask, dim=-2)
        #do nonlinearity
        activated_features_mol = F.relu(mol_feature)
        
        # when the descriptor value has lower or upper bounds
        mol_feature_viz = []
        mol_feature_viz.append(mol_feature) 
        
        mol_attention_weight_viz = []
        mol_softmax_mask = atom_mask.clone()
        mol_softmax_mask[mol_softmax_mask == 0] = -9e8
        mol_softmax_mask[mol_softmax_mask == 1] = 0
        mol_softmax_mask = mol_softmax_mask.type(torch.cuda.FloatTensor)
        
        for t in range(self.T):
            
            mol_prediction_expand = activated_features_mol.unsqueeze(-2).expand(batch_size, mol_length, fingerprint_dim)
            mol_align = torch.cat([mol_prediction_expand, activated_features], dim=-1)
            mol_align_score = self.dropout(F.leaky_relu(self.mol_align(mol_align)))
            mol_align_score = mol_align_score + mol_softmax_mask
            mol_attention_weight = F.softmax(mol_align_score,-2)
            mol_attention_weight = mol_attention_weight * atom_mask
#             print(mol_attention_weight.shape,mol_attention_weight)
            mol_attention_weight_viz.append(mol_attention_weight)

            activated_features_transform = self.mol_attend(self.dropout(activated_features))
            mol_context = torch.sum(torch.mul(mol_attention_weight,activated_features_transform),-2)
#             print(mol_context.shape,mol_context)
            mol_context = F.elu(mol_context)
            mol_feature = self.mol_GRUCell(mol_context, mol_feature)
#             print(mol_feature.shape,mol_feature)

            mol_feature_unbounded_viz.append(mol_feature)
            #do nonlinearity
            activated_features_mol = F.relu(mol_feature)           
            mol_feature_viz.append(activated_features_mol)
            
        mol_prediction = self.output(self.dropout(mol_feature))
            
        return atom_feature_viz, atom_attention_weight_viz, mol_feature_viz, mol_feature_unbounded_viz, mol_attention_weight_viz, mol_prediction

src/PointGAT/Layers_viz.py
sevencheung2021/PointGAT
# from numpy import identity
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

import numpy as np
from torch.autograd import Variable
from torch_geometric.nn import GCNConv, GATConv, GINConv, global_add_pool
from torch_geometric.nn import global_mean_pool as gap, global_max_pool as gmp


class PointGAT(nn.Module):

    def __init__(self, radius, T, input_feature_dim, input_bond_dim,\
            fingerprint_dim, output_units_num, p_dropout,xyz_feature_dim):
        super(PointGAT, self).__init__()
        # graph attention for atom embedding
        self.atom_fc = nn.Linear(input_feature_dim, fingerprint_dim)
        self.neighbor_fc = nn.Linear(input_feature_dim+input_bond_dim, fingerprint_dim)
        self.GRUCell = nn.ModuleList([nn.GRUCell(fingerprint_dim, fingerprint_dim) for r in range(radius)])
        self.align = nn.ModuleList([nn.Linear(2*fingerprint_dim,1) for r in range(radius)])
        self.attend = nn.ModuleList([nn.Linear(fingerprint_dim, fingerprint_dim) for r in range(radius)])
        # graph attention for molecule embedding
        self.mol_GRUCell = nn.GRUCell(fingerprint_dim, fingerprint_dim)
        self.mol_align = nn.Linear(2*fingerprint_dim,1)
        self.mol_attend = nn.Linear(fingerprint_dim, fingerprint_dim)
        # you may alternatively assign a different set of parameter in each attentive layer for molecule embedding like in atom embedding process.
#         self.mol_GRUCell = nn.ModuleList([nn.GRUCell(fingerprint_dim, fingerprint_dim) for t in range(T)])
#         self.mol_align = nn.ModuleList([nn.Linear(2*fingerprint_dim,1) for t in range(T)])
#         self.mol_attend = nn.ModuleList([nn.Linear(fingerprint_dim, fingerprint_dim) for t in range(T)])
        
        self.dropout = nn.Dropout(p=p_dropout)
        self.output = nn.Linear(fingerprint_dim, output_units_num)


        self.radius = radius
        self.T = T

        self.relu = nn.ReLU()
        
        # self.transform = nn.Linear(xyz_feature_dim, 64)
        # self.conv1 = torch.nn.Conv1d(3, 64, 1)
        self.conv1 = torch.nn.Conv1d(6, 64, 1)   # feature_dim =5 [x,y,z]
        self.conv2 = torch.nn.Conv1d(64, 128, 1)
        self.conv3 = torch.nn.Conv1d(128, 1024, 1)
        self.fc1 = nn.Linear(1024, 512)
        self.fc2 = nn.Linear(512, fingerprint_dim)
        
        self.bn1 = nn.BatchNorm1d(64)
        self.bn2 = nn.BatchNorm1d(128)
        self.bn3 = nn.BatchNorm1d(1024)
        self.bn4 = nn.BatchNorm1d(512)
        self.bn5 = nn.BatchNorm1d(fingerprint_dim)

        self.norm = nn.LayerNorm(fingerprint_dim)

        self.fc3 = nn.Linear(fingerprint_dim*2, fingerprint_dim)
        self.fc4 = nn.Linear(fingerprint_dim,1)


    def forward(self, atom_list, bond_list, atom_degree_list, bond_degree_list, atom_mask,xyz_feature):
        atom_mask = atom_mask.unsqueeze(2)
        batch_size,mol_length,num_atom_feat = atom_list.size()
        atom_feature = F.leaky_relu(self.atom_fc(atom_list))
        
       
        identity = atom_feature
        # print(atom_feature.shape)

        bond_neighbor = [bond_list[i][bond_degree_list[i]] for i in range(batch_size)]
        bond_neighbor = torch.stack(bond_neighbor, dim=0)
        atom_neighbor = [atom_list[i][atom_degree_list[i]] for i in range(batch_size)]
        atom_neighbor = torch.stack(atom_neighbor, dim=0)
        # then concatenate them
        neighbor_feature = torch.cat([atom_neighbor, bond_neighbor],dim=-1)
        neighbor_feature = F.leaky_relu(self.neighbor_fc(neighbor_feature))

        # generate mask to eliminate the influence of blank atoms
        attend_mask = atom_degree_list.clone()
        attend_mask[attend_mask != mol_length-1] = 1
        attend_mask[attend_mask == mol_length-1] = 0
        attend_mask = attend_mask.type(torch.cuda.FloatTensor).unsqueeze(-1)

        softmax_mask = atom_degree_list.clone()
        softmax_mask[softmax_mask != mol_length-1] = 0
        softmax_mask[softmax_mask == mol_length-1] = -9e8 # make the softmax value extremly small
        softmax_mask = softmax_mask.type(torch.cuda.FloatTensor).unsqueeze(-1)

        batch_size, mol_length, max_neighbor_num, fingerprint_dim = neighbor_feature.shape
        atom_feature_expand = atom_feature.unsqueeze(-2).expand(batch_size, mol_length, max_neighbor_num, fingerprint_dim)
        feature_align = torch.cat([atom_feature_expand, neighbor_feature],dim=-1)
        
        align_score = F.leaky_relu(self.align[0](self.dropout(feature_align)))
#             print(attention_weight)
        align_score = align_score + softmax_mask
        attention_weight = F.softmax(align_score,-2)
#             print(attention_weight)
        attention_weight = attention_weight * attend_mask
#         print(attention_weight)
        neighbor_feature_transform = self.attend[0](self.dropout(neighbor_feature))
#             print(features_neighbor_transform.shape)
        context = torch.sum(torch.mul(attention_weight,neighbor_feature_transform),-2)
#             print(context.shape)
        context = F.elu(context)
        context_reshape = context.view(batch_size*mol_length, fingerprint_dim)
        atom_feature_reshape = atom_feature.view(batch_size*mol_length, fingerprint_dim)
        atom_feature_reshape = self.GRUCell[0](context_reshape, atom_feature_reshape)
        atom_feature = atom_feature_reshape.view(batch_size, mol_length, fingerprint_dim)

        #do nonlinearity
        activated_features = F.relu(atom_feature)

        for d in range(self.radius-1):
            # bonds_indexed = [bond_list[i][torch.cuda.LongTensor(bond_degree_list)[i]] for i in range(batch_size)]
            neighbor_feature = [activated_features[i][atom_degree_list[i]] for i in range(batch_size)]
            
            # neighbor_feature is a list of 3D tensor, so we need to stack them into a 4D tensor first
            neighbor_feature = torch.stack(neighbor_feature, dim=0)
            atom_feature_expand = activated_features.unsqueeze(-2).expand(batch_size, mol_length, max_neighbor_num, fingerprint_dim)

            feature_align = torch.cat([atom_feature_expand, neighbor_feature],dim=-1)

            align_score = F.leaky_relu(self.align[d+1](self.dropout(feature_align)))
    #             print(attention_weight)
            align_score = align_score + softmax_mask
            attention_weight = F.softmax(align_score,-2)
#             print(attention_weight)
            attention_weight = attention_weight * attend_mask
#             print(attention_weight)
            neighbor_feature_transform = self.attend[d+1](self.dropout(neighbor_feature))
    #             print(features_neighbor_transform.shape)
            context = torch.sum(torch.mul(attention_weight,neighbor_feature_transform),-2)
    #             print(context.shape)
            context = F.elu(context)
            context_reshape = context.view(batch_size*mol_length, fingerprint_dim)
#             atom_feature_reshape = atom_feature.view(batch_size*mol_length, fingerprint_dim)
            atom_feature_reshape = self.GRUCell[d+1](context_reshape, atom_feature_reshape)
            atom_feature = atom_feature_reshape.view(batch_size, mol_length, fingerprint_dim)
            
        #     print('atom_features shape')
        #     print(atom_feature.shape)

            # do nonlinearity
            activated_features = F.relu(atom_feature)
        # print('activated_features shape')
        # print(activated_features.shape)


        mol_feature = torch.sum(activated_features * atom_mask, dim=-2)
        # print('features_mol from sum atom feature shape')
        # print(mol_feature.shape)

        atom_level_mol_feature = mol_feature  # ZR add


        # do nonlinearity
        activated_features_mol = F.relu(mol_feature)           
        # print('activated_features_mol from sum atom feature shape')
        # print(activated_features_mol.shape)


        mol_softmax_mask = atom_mask.clone()
        mol_softmax_mask[mol_softmax_mask == 0] = -9e8
        mol_softmax_mask[mol_softmax_mask == 1] = 0
        mol_softmax_mask = mol_softmax_mask.type(torch.cuda.FloatTensor)
        
        for t in range(self.T):
            
            mol_prediction_expand = activated_features_mol.unsqueeze(-2).expand(batch_size, mol_length, fingerprint_dim)
            mol_align = torch.cat([mol_prediction_expand, activated_features], dim=-1)
            mol_align_score = F.leaky_relu(self.mol_align(mol_align))
            mol_align_score = mol_align_score + mol_softmax_mask
            mol_attention_weight = F.softmax(mol_align_score,-2)
            mol_attention_weight = mol_attention_weight * atom_mask
#             print(mol_attention_weight.shape,mol_attention_weight)
            activated_features_transform = self.mol_attend(self.dropout(activated_features))
#             aggregate embeddings of atoms in a molecule
            mol_context = torch.sum(torch.mul(mol_attention_weight,activated_features_transform),-2)
#             print(mol_context.shape,mol_context)
            mol_context = F.elu(mol_context)
            mol_feature = self.mol_GRUCell(mol_context, mol_feature)
            mol_feature = self.norm(mol_feature)+ atom_level_mol_feature

        #     print(mol_feature.shape,mol_feature)

            # do nonlinearity
            activated_features_mol = F.relu(mol_feature)           

        ##### xyz_feature #####
        batchsize = batch_size

        xyz_feature = xyz_feature.permute(0,2,1)
        
        xyz = F.leaky_relu(self.bn1(self.conv1(xyz_feature)))
        xyz = F.leaky_relu(self.bn2(self.conv2(xyz)))
        xyz = F.leaky_relu(self.bn3(self.conv3(xyz)))
        xyz = torch.max(xyz, 2, keepdim=True)[0]
        # xyz = xyz.view(-1, 1024)
        xyz = xyz.view(-1, 1024)
        

        xyz = F.relu(self.fc1(xyz))
        xyz = F.relu(self.fc2(xyz))
        

        cat = torch.cat([mol_feature,xyz],dim=1) #(N,256*2)
        
        final_mol = F.relu(self.fc3(cat))
        
        mol_prediction = self.fc4(final_mol)
        
        return mol_prediction
        
        


src/PointGAT/PointGAT_Layers.py
sevencheung2021/PointGAT
'''
Author: Seven Rong Cheung
Date: 2023-09-20 16:01:07
LastEditors: Seven Rong Cheung
LastEditTime: 2023-09-20 16:04:33
FilePath: /src/PointGAT/__init__.py
Description: 

Copyright (c) 2023 by {rongzhangthu@yeah.net}, All Rights Reserved. 
'''

from PointGAT.PointGAT_Layers import PointGAT
from PointGAT.Layers_viz import PointGAT_viz
from PointGAT.getFeatures import save_smiles_dicts, get_smiles_dicts, get_smiles_array, moltosvg_highlight
src/PointGAT/__init__.py
sevencheung2021/PointGAT

import matplotlib.pyplot as plt
plt.switch_backend('agg')

from rdkit.Chem import MolFromSmiles
from rdkit.Chem import Draw
import numpy as np
import rdkit
from rdkit import Chem
from rdkit.Chem import AllChem
import os
from PointGAT.Featurizer import *
import pickle
import time
from rdkit.Chem import rdDepictor
from rdkit.Chem.Draw import rdMolDraw2D
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import matplotlib
from rdkit.Chem.Draw import SimilarityMaps
from io import StringIO

smilesList = ['CC']
degrees = [0, 1, 2, 3, 4, 5]


class MolGraph(object):
    def __init__(self):
        self.nodes = {} # dict of lists of nodes, keyed by node type

    def new_node(self, ntype, features=None, rdkit_ix=None):
        new_node = Node(ntype, features, rdkit_ix)
        self.nodes.setdefault(ntype, []).append(new_node)
        return new_node

    def add_subgraph(self, subgraph):
        old_nodes = self.nodes
        new_nodes = subgraph.nodes
        for ntype in set(old_nodes.keys()) | set(new_nodes.keys()):
            old_nodes.setdefault(ntype, []).extend(new_nodes.get(ntype, []))

    def sort_nodes_by_degree(self, ntype):
        nodes_by_degree = {i : [] for i in degrees}
        for node in self.nodes[ntype]:
            nodes_by_degree[len(node.get_neighbors(ntype))].append(node)

        new_nodes = []
        for degree in degrees:
            cur_nodes = nodes_by_degree[degree]
            self.nodes[(ntype, degree)] = cur_nodes
            new_nodes.extend(cur_nodes)

        self.nodes[ntype] = new_nodes

    def feature_array(self, ntype):
        assert ntype in self.nodes
        return np.array([node.features for node in self.nodes[ntype]])

    def rdkit_ix_array(self):
        return np.array([node.rdkit_ix for node in self.nodes['atom']])

    def neighbor_list(self, self_ntype, neighbor_ntype):
        assert self_ntype in self.nodes and neighbor_ntype in self.nodes
        neighbor_idxs = {n : i for i, n in enumerate(self.nodes[neighbor_ntype])}
        return [[neighbor_idxs[neighbor]
                 for neighbor in self_node.get_neighbors(neighbor_ntype)]
                for self_node in self.nodes[self_ntype]]

class Node(object):
    __slots__ = ['ntype', 'features', '_neighbors', 'rdkit_ix']
    def __init__(self, ntype, features, rdkit_ix):
        self.ntype = ntype
        self.features = features
        self._neighbors = []
        self.rdkit_ix = rdkit_ix

    def add_neighbors(self, neighbor_list):
        for neighbor in neighbor_list:
            self._neighbors.append(neighbor)
            neighbor._neighbors.append(self)

    def get_neighbors(self, ntype):
        return [n for n in self._neighbors if n.ntype == ntype]

class memoize(object):
    def __init__(self, func):
        self.func = func
        self.cache = {}

    def __call__(self, *args):
        if args in self.cache:
            return self.cache[args]
        else:
            result = self.func(*args)
            self.cache[args] = result
            return result

    def __get__(self, obj, objtype):
        return partial(self.__call__, obj)

def graph_from_smiles_tuple(smiles_tuple):
    graph_list = [graph_from_smiles(s) for s in smiles_tuple]
    big_graph = MolGraph()
    for subgraph in graph_list:
        big_graph.add_subgraph(subgraph)

    # This sorting allows an efficient (but brittle!) indexing later on.
    big_graph.sort_nodes_by_degree('atom')
    return big_graph

def graph_from_smiles(smiles):
    graph = MolGraph()
    mol = MolFromSmiles(smiles)
    if not mol:
        raise ValueError("Could not parse SMILES string:", smiles)
    atoms_by_rd_idx = {}
    for atom in mol.GetAtoms():
        new_atom_node = graph.new_node('atom', features=atom_features(atom), rdkit_ix=atom.GetIdx())
        atoms_by_rd_idx[atom.GetIdx()] = new_atom_node

    for bond in mol.GetBonds():
        atom1_node = atoms_by_rd_idx[bond.GetBeginAtom().GetIdx()]
        atom2_node = atoms_by_rd_idx[bond.GetEndAtom().GetIdx()]
        new_bond_node = graph.new_node('bond', features=bond_features(bond))
        new_bond_node.add_neighbors((atom1_node, atom2_node))
        atom1_node.add_neighbors((atom2_node,))

    mol_node = graph.new_node('molecule')
    mol_node.add_neighbors(graph.nodes['atom'])
    return graph

def array_rep_from_smiles(molgraph):
    """Precompute everything we need from MolGraph so that we can free the memory asap."""
    #molgraph = graph_from_smiles_tuple(tuple(smiles))
    degrees = [0,1,2,3,4,5]
    arrayrep = {'atom_features' : molgraph.feature_array('atom'),
                'bond_features' : molgraph.feature_array('bond'),
                'atom_list'     : molgraph.neighbor_list('molecule', 'atom'),
                'rdkit_ix'      : molgraph.rdkit_ix_array()}

    for degree in degrees:
        arrayrep[('atom_neighbors', degree)] = \
            np.array(molgraph.neighbor_list(('atom', degree), 'atom'), dtype=int)
        arrayrep[('bond_neighbors', degree)] = \
            np.array(molgraph.neighbor_list(('atom', degree), 'bond'), dtype=int)
    return arrayrep


def num_atom_features():
    # Return length of feature vector using a very simple molecule.
    m = Chem.MolFromSmiles('CC')
    alist = m.GetAtoms()
    a = alist[0]
    return len(atom_features(a))

def num_bond_features():
    # Return length of feature vector using a very simple molecule.
    simple_mol = Chem.MolFromSmiles('CC')
    Chem.SanitizeMol(simple_mol)
    return len(bond_features(simple_mol.GetBonds()[0]))


def gen_descriptor_data(smilesList):

    smiles_to_fingerprint_array = {}

    for i,smiles in enumerate(smilesList):
#         if i > 5:
#             print("Due to the limited computational resource, submission with more than 5 molecules will not be processed")
#             break
        # smiles = Chem.MolToSmiles(Chem.MolFromSmiles(smiles), isomericSmiles=True)
        try:
            molgraph = graph_from_smiles(smiles)
            molgraph.sort_nodes_by_degree('atom')
            arrayrep = array_rep_from_smiles(molgraph)

            smiles_to_fingerprint_array[smiles] = arrayrep
            
        except:
            print('error preprocess{}'.format(smiles))
            time.sleep(3)
    return smiles_to_fingerprint_array



def get_smiles_dicts(smilesList):
    #first need to get the max atom length
    max_atom_len = 0
    max_bond_len = 0
    num_atom_features = 0
    num_bond_features = 0
    smiles_to_rdkit_list = {}

    smiles_to_fingerprint_features = gen_descriptor_data(smilesList)

    for smiles,arrayrep in smiles_to_fingerprint_features.items():
        atom_features = arrayrep['atom_features']
        bond_features = arrayrep['bond_features']

        rdkit_list = arrayrep['rdkit_ix']
        smiles_to_rdkit_list[smiles] = rdkit_list

        atom_len,num_atom_features = atom_features.shape
        bond_len,num_bond_features = bond_features.shape

        if atom_len > max_atom_len:
            max_atom_len = atom_len
        if bond_len > max_bond_len:
            max_bond_len = bond_len

    #then add 1 so I can zero pad everything
    max_atom_index_num = max_atom_len
    max_bond_index_num = max_bond_len

    max_atom_len += 1
    max_bond_len += 1

    smiles_to_atom_info = {}
    smiles_to_bond_info = {}

    smiles_to_atom_neighbors = {}
    smiles_to_bond_neighbors = {}

    smiles_to_atom_mask = {}

    degrees = [0,1,2,3,4,5]
    #then run through our numpy array again
    for smiles,arrayrep in smiles_to_fingerprint_features.items():
        mask = np.zeros((max_atom_len))

        #get the basic info of what
        #    my atoms and bonds are initialized
        atoms = np.zeros((max_atom_len,num_atom_features))
        bonds = np.zeros((max_bond_len,num_bond_features))

        #then get the arrays initlialized for the neighbors
        atom_neighbors = np.zeros((max_atom_len,len(degrees)))
        bond_neighbors = np.zeros((max_atom_len,len(degrees)))

        #now set these all to the last element of the list, which is zero padded
        atom_neighbors.fill(max_atom_index_num)
        bond_neighbors.fill(max_bond_index_num)

        atom_features = arrayrep['atom_features']
        bond_features = arrayrep['bond_features']

        for i,feature in enumerate(atom_features):
            mask[i] = 1.0
            atoms[i] = feature

        for j,feature in enumerate(bond_features):
            bonds[j] = feature

        atom_neighbor_count = 0
        bond_neighbor_count = 0
        working_atom_list = []
        working_bond_list = []
        for degree in degrees:
            atom_neighbors_list = arrayrep[('atom_neighbors', degree)]
            bond_neighbors_list = arrayrep[('bond_neighbors', degree)]

            if len(atom_neighbors_list) > 0:

                for i,degree_array in enumerate(atom_neighbors_list):
                    for j,value in enumerate(degree_array):
                        atom_neighbors[atom_neighbor_count,j] = value
                    atom_neighbor_count += 1

            if len(bond_neighbors_list) > 0:
                for i,degree_array in enumerate(bond_neighbors_list):
                    for j,value in enumerate(degree_array):
                        bond_neighbors[bond_neighbor_count,j] = value
                    bond_neighbor_count += 1

        #then add everything to my arrays
        smiles_to_atom_info[smiles] = atoms
        smiles_to_bond_info[smiles] = bonds

        smiles_to_atom_neighbors[smiles] = atom_neighbors
        smiles_to_bond_neighbors[smiles] = bond_neighbors
        
        smiles_to_atom_mask[smiles] = mask

    del smiles_to_fingerprint_features
    feature_dicts = {}
#     feature_dicts['smiles_to_atom_mask'] = smiles_to_atom_mask
#     feature_dicts['smiles_to_atom_info']= smiles_to_atom_info
    feature_dicts = {
        'smiles_to_atom_mask': smiles_to_atom_mask,
        'smiles_to_atom_info': smiles_to_atom_info,
        'smiles_to_bond_info': smiles_to_bond_info,
        'smiles_to_atom_neighbors': smiles_to_atom_neighbors,
        'smiles_to_bond_neighbors': smiles_to_bond_neighbors,
        'smiles_to_rdkit_list': smiles_to_rdkit_list
    }
    return feature_dicts
def save_smiles_dicts(smilesList,filename):
    #first need to get the max atom length
    max_atom_len = 0
    max_bond_len = 0
    num_atom_features = 0
    num_bond_features = 0
    smiles_to_rdkit_list = {}

    smiles_to_fingerprint_features = gen_descriptor_data(smilesList)

    for smiles,arrayrep in smiles_to_fingerprint_features.items():
        atom_features = arrayrep['atom_features']
        bond_features = arrayrep['bond_features']

        rdkit_list = arrayrep['rdkit_ix']
        smiles_to_rdkit_list[smiles] = rdkit_list

        atom_len,num_atom_features = atom_features.shape
        bond_len,num_bond_features = bond_features.shape

        if atom_len > max_atom_len:
            max_atom_len = atom_len
        if bond_len > max_bond_len:
            max_bond_len = bond_len

    #then add 1 so I can zero pad everything
    max_atom_index_num = max_atom_len
    max_bond_index_num = max_bond_len

    max_atom_len += 1
    max_bond_len += 1

    smiles_to_atom_info = {}
    smiles_to_bond_info = {}

    smiles_to_atom_neighbors = {}
    smiles_to_bond_neighbors = {}

    smiles_to_atom_mask = {}

    degrees = [0,1,2,3,4,5]
    #then run through our numpy array again
    for smiles,arrayrep in smiles_to_fingerprint_features.items():
        mask = np.zeros((max_atom_len))

        #get the basic info of what
        #    my atoms and bonds are initialized
        atoms = np.zeros((max_atom_len,num_atom_features))
        bonds = np.zeros((max_bond_len,num_bond_features))

        #then get the arrays initlialized for the neighbors
        atom_neighbors = np.zeros((max_atom_len,len(degrees)))
        bond_neighbors = np.zeros((max_atom_len,len(degrees)))

        #now set these all to the last element of the list, which is zero padded
        atom_neighbors.fill(max_atom_index_num)
        bond_neighbors.fill(max_bond_index_num)

        atom_features = arrayrep['atom_features']
        bond_features = arrayrep['bond_features']

        for i,feature in enumerate(atom_features):
            mask[i] = 1.0
            atoms[i] = feature

        for j,feature in enumerate(bond_features):
            bonds[j] = feature

        atom_neighbor_count = 0
        bond_neighbor_count = 0
        working_atom_list = []
        working_bond_list = []
        for degree in degrees:
            atom_neighbors_list = arrayrep[('atom_neighbors', degree)]
            bond_neighbors_list = arrayrep[('bond_neighbors', degree)]

            if len(atom_neighbors_list) > 0:

                for i,degree_array in enumerate(atom_neighbors_list):
                    for j,value in enumerate(degree_array):
                        atom_neighbors[atom_neighbor_count,j] = value
                    atom_neighbor_count += 1

            if len(bond_neighbors_list) > 0:
                for i,degree_array in enumerate(bond_neighbors_list):
                    for j,value in enumerate(degree_array):
                        bond_neighbors[bond_neighbor_count,j] = value
                    bond_neighbor_count += 1

        #then add everything to my arrays
        smiles_to_atom_info[smiles] = atoms
        smiles_to_bond_info[smiles] = bonds

        smiles_to_atom_neighbors[smiles] = atom_neighbors
        smiles_to_bond_neighbors[smiles] = bond_neighbors
        
        smiles_to_atom_mask[smiles] = mask

    del smiles_to_fingerprint_features
    feature_dicts = {}
#     feature_dicts['smiles_to_atom_mask'] = smiles_to_atom_mask
#     feature_dicts['smiles_to_atom_info']= smiles_to_atom_info
    feature_dicts = {
        'smiles_to_atom_mask': smiles_to_atom_mask,
        'smiles_to_atom_info': smiles_to_atom_info,
        'smiles_to_bond_info': smiles_to_bond_info,
        'smiles_to_atom_neighbors': smiles_to_atom_neighbors,
        'smiles_to_bond_neighbors': smiles_to_bond_neighbors,
        'smiles_to_rdkit_list': smiles_to_rdkit_list
    }
    pickle.dump(feature_dicts,open(filename+'.pickle',"wb"))
    print('feature dicts file saved as '+ filename+'.pickle')
    return feature_dicts

def get_smiles_array(smilesList, feature_dicts):
    x_mask = []
    x_atom = []
    x_bonds = []
    x_atom_index = []
    x_bond_index = []
    for smiles in smilesList:
        x_mask.append(feature_dicts['smiles_to_atom_mask'][smiles])
        x_atom.append(feature_dicts['smiles_to_atom_info'][smiles])
        x_bonds.append(feature_dicts['smiles_to_bond_info'][smiles])
        x_atom_index.append(feature_dicts['smiles_to_atom_neighbors'][smiles])
        x_bond_index.append(feature_dicts['smiles_to_bond_neighbors'][smiles])
    return np.asarray(x_atom),np.asarray(x_bonds),np.asarray(x_atom_index),\
        np.asarray(x_bond_index),np.asarray(x_mask),feature_dicts['smiles_to_rdkit_list']


def moltosvg(mol,molSize=(280,218),kekulize=False):
    mc = Chem.Mol(mol.ToBinary())
    if kekulize:
        try:
            Chem.Kekulize(mc)
        except:
            mc = Chem.Mol(mol.ToBinary())
    if not mc.GetNumConformers():
        rdDepictor.Compute2DCoords(mc)
    drawer = rdMolDraw2D.MolDraw2DSVG(molSize[0],molSize[1])
    drawer.DrawMolecule(mc)
    drawer.FinishDrawing()
    svg = drawer.GetDrawingText()
#    return svg
    return svg.replace('svg:','')

def rreplace(s, old, new, occurrence):
    li = s.rsplit(old, occurrence)
    return new.join(li)

def moltosvg_highlight(smiles, atom_list, atom_predictions, molecule_prediction,\
    molSize=(280,218),kekulize=False):
    
    mol = Chem.MolFromSmiles(smiles)
    min_pred = 0.05
    max_pred = 0.8
    # min_pred = np.amax([np.amin(all_atom_predictions),0])
    # min_pred = np.amin(all_atom_predictions)
    # min_pred = min(filter(lambda x:x>0,all_atom_predictions))
    # max_pred = np.amax(all_atom_predictions)
    note = 'y_pred: '+ str(molecule_prediction)

    norm = matplotlib.colors.Normalize(vmin=np.exp(0.068),vmax=np.exp(max_pred))
    cmap = cm.get_cmap('gray_r')


    plt_colors = cm.ScalarMappable(norm=norm,cmap=cmap)

    atom_colors = {}
    for i,atom in enumerate(atom_list):
        color_rgba = plt_colors.to_rgba(atom_predictions[i])
        atom_rgb = color_rgba #(color_rgba[0],color_rgba[1],color_rgba[2])
        atom_colors[atom] = atom_rgb

    rdDepictor.Compute2DCoords(mol)

    drawer = rdMolDraw2D.MolDraw2DSVG(molSize[0],molSize[1])
    #drawer.DrawMolecule(mc)
    drawer.DrawMolecule(mol,highlightAtoms=atom_list,highlightBonds=[],
        highlightAtomColors=atom_colors,legend=note)
    drawer.SetFontSize(68)
    drawer.FinishDrawing()
    svg = drawer.GetDrawingText()
    # svg = svg.replace('size:15px','size:18px')
    # svg = svg.replace('size:12px','size:28px')
    # svg = rreplace(svg,'size:11px','size:28px',1)
    # It seems that the svg renderer used doesn't quite hit the spec.
    # Here are some fixes to make it work in the notebook, although I think
    # the underlying issue needs to be resolved at the generation step
    return svg.replace('svg:','')

def moltosvg_highlight_known(smiles, atom_list, atom_predictions, molecule_prediction, molecule_experiment, Number,\
        molSize=(280,218),kekulize=False):
    
    mol = Chem.MolFromSmiles(smiles)
    min_pred = 0.05
    max_pred = 0.8
    # min_pred = np.amax([np.amin(all_atom_predictions),0])
    # min_pred = np.amin(all_atom_predictions)
    # min_pred = min(filter(lambda x:x>0,all_atom_predictions))
    # max_pred = np.amax(all_atom_predictions)
    note = '('+ str(Number) +') y-y\' : '+ str(round(molecule_experiment,2)) + '-' + str(round(molecule_prediction,2))

    norm = matplotlib.colors.Normalize(vmin=0.01*5,vmax=max_pred*6)
    cmap = cm.get_cmap('gray_r')


    plt_colors = cm.ScalarMappable(norm=norm,cmap=cmap)

    atom_colors = {}
    for i,atom in enumerate(atom_list):
        color_rgba = plt_colors.to_rgba(atom_predictions[i])
        atom_rgb = color_rgba #(color_rgba[0],color_rgba[1],color_rgba[2])
        atom_colors[atom] = atom_rgb

    rdDepictor.Compute2DCoords(mol)

    drawer = rdMolDraw2D.MolDraw2DSVG(molSize[0],molSize[1])
    mol = rdMolDraw2D.PrepareMolForDrawing(mol)
    #drawer.DrawMolecule(mc)
    drawer.DrawMolecule(mol,highlightAtoms=atom_list,highlightBonds=[],
        highlightAtomColors=atom_colors,legend=note)
    drawer.SetFontSize(68)
    drawer.FinishDrawing()
    svg = drawer.GetDrawingText()
    # svg = svg.replace('size:15px','size:18px')
    # svg = svg.replace('size:12px','size:28px')
    # svg = rreplace(svg,'size:11px','size:28px',1)
    # It seems that the svg renderer used doesn't quite hit the spec.
    # Here are some fixes to make it work in the notebook, although I think
    # the underlying issue needs to be resolved at the generation step
    return svg.replace('svg:','')

def weighted_highlight_known(smiles, atom_list, atom_predictions, molecule_prediction, molecule_experiment, Number,\
        molSize=(128,128)):
    
    mol = Chem.MolFromSmiles(smiles)
    note = '('+ str(Number) +') y-y\' : '+ str(round(molecule_experiment,2)) + '-' + str(round(molecule_prediction,2))

    contribs = [atom_predictions[m] for m in np.argsort(atom_list)]
    fig = SimilarityMaps.GetSimilarityMapFromWeights(mol, contribs, colorMap='bwr', contourLines=5, size=molSize)
    fig.axes[0].set_title(note)
    sio = StringIO()
    fig.savefig(sio, format="svg", bbox_inches='tight')
    svg = sio.getvalue()   
    return svg

def moltosvg_interaction_known(mol, atom_list, atom_predictions, molecule_prediction, molecule_experiment, max_atom_pred, min_atom_pred, Number):
    
    note = '('+ str(Number) +') y-y\' : '+ str(round(molecule_experiment,2)) + '-' + str(round(molecule_prediction,2))
    norm = matplotlib.colors.Normalize(vmin=min_atom_pred*0.9,vmax=max_atom_pred*1.1)
    cmap = cm.get_cmap('gray_r')

    plt_colors = cm.ScalarMappable(norm=norm,cmap=cmap)

    atom_colors = {}
    for i,atom in enumerate(atom_list):
        atom_colors[atom] = plt_colors.to_rgba(atom_predictions[i])          
    rdDepictor.Compute2DCoords(mol)

    drawer = rdMolDraw2D.MolDraw2DSVG(280,218)
    op = drawer.drawOptions()
    for i in range(mol.GetNumAtoms()):
        op.atomLabels[i]=mol.GetAtomWithIdx(i).GetSymbol() + str(i)

    mol = rdMolDraw2D.PrepareMolForDrawing(mol)
    drawer.DrawMolecule(mol,highlightAtoms=atom_list,highlightBonds=[],
        highlightAtomColors=atom_colors,legend=note)
    drawer.SetFontSize(68)
    drawer.FinishDrawing()
    svg = drawer.GetDrawingText()
    return svg.replace('svg:','')




src/PointGAT/getFeatures.py
sevencheung2021/PointGAT
import matplotlib.pyplot as plt
plt.switch_backend('agg')

from rdkit.Chem import MolFromSmiles
from rdkit.Chem import Draw
import numpy as np
import rdkit
from rdkit import Chem
from rdkit.Chem import AllChem
import os
from PointGAT.Featurizer_aromaticity_rm import *
import pickle
import time
from rdkit.Chem import rdDepictor
from rdkit.Chem.Draw import rdMolDraw2D
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import matplotlib
from rdkit.Chem.Draw import SimilarityMaps
from io import StringIO

smilesList = ['CC']
degrees = [0, 1, 2, 3, 4, 5]


class MolGraph(object):
    def __init__(self):
        self.nodes = {} # dict of lists of nodes, keyed by node type

    def new_node(self, ntype, features=None, rdkit_ix=None):
        new_node = Node(ntype, features, rdkit_ix)
        self.nodes.setdefault(ntype, []).append(new_node)
        return new_node

    def add_subgraph(self, subgraph):
        old_nodes = self.nodes
        new_nodes = subgraph.nodes
        for ntype in set(old_nodes.keys()) | set(new_nodes.keys()):
            old_nodes.setdefault(ntype, []).extend(new_nodes.get(ntype, []))

    def sort_nodes_by_degree(self, ntype):
        nodes_by_degree = {i : [] for i in degrees}
        for node in self.nodes[ntype]:
            nodes_by_degree[len(node.get_neighbors(ntype))].append(node)

        new_nodes = []
        for degree in degrees:
            cur_nodes = nodes_by_degree[degree]
            self.nodes[(ntype, degree)] = cur_nodes
            new_nodes.extend(cur_nodes)

        self.nodes[ntype] = new_nodes

    def feature_array(self, ntype):
        assert ntype in self.nodes
        return np.array([node.features for node in self.nodes[ntype]])

    def rdkit_ix_array(self):
        return np.array([node.rdkit_ix for node in self.nodes['atom']])

    def neighbor_list(self, self_ntype, neighbor_ntype):
        assert self_ntype in self.nodes and neighbor_ntype in self.nodes
        neighbor_idxs = {n : i for i, n in enumerate(self.nodes[neighbor_ntype])}
        return [[neighbor_idxs[neighbor]
                 for neighbor in self_node.get_neighbors(neighbor_ntype)]
                for self_node in self.nodes[self_ntype]]

class Node(object):
    __slots__ = ['ntype', 'features', '_neighbors', 'rdkit_ix']
    def __init__(self, ntype, features, rdkit_ix):
        self.ntype = ntype
        self.features = features
        self._neighbors = []
        self.rdkit_ix = rdkit_ix

    def add_neighbors(self, neighbor_list):
        for neighbor in neighbor_list:
            self._neighbors.append(neighbor)
            neighbor._neighbors.append(self)

    def get_neighbors(self, ntype):
        return [n for n in self._neighbors if n.ntype == ntype]

class memoize(object):
    def __init__(self, func):
        self.func = func
        self.cache = {}

    def __call__(self, *args):
        if args in self.cache:
            return self.cache[args]
        else:
            result = self.func(*args)
            self.cache[args] = result
            return result

    def __get__(self, obj, objtype):
        return partial(self.__call__, obj)

def graph_from_smiles_tuple(smiles_tuple):
    graph_list = [graph_from_smiles(s) for s in smiles_tuple]
    big_graph = MolGraph()
    for subgraph in graph_list:
        big_graph.add_subgraph(subgraph)

    # This sorting allows an efficient (but brittle!) indexing later on.
    big_graph.sort_nodes_by_degree('atom')
    return big_graph

def graph_from_smiles(smiles):
    graph = MolGraph()
    mol = MolFromSmiles(smiles)
    if not mol:
        raise ValueError("Could not parse SMILES string:", smiles)
    atoms_by_rd_idx = {}
    for atom in mol.GetAtoms():
        new_atom_node = graph.new_node('atom', features=atom_features(atom), rdkit_ix=atom.GetIdx())
        atoms_by_rd_idx[atom.GetIdx()] = new_atom_node

    for bond in mol.GetBonds():
        atom1_node = atoms_by_rd_idx[bond.GetBeginAtom().GetIdx()]
        atom2_node = atoms_by_rd_idx[bond.GetEndAtom().GetIdx()]
        new_bond_node = graph.new_node('bond', features=bond_features(bond))
        new_bond_node.add_neighbors((atom1_node, atom2_node))
        atom1_node.add_neighbors((atom2_node,))

    mol_node = graph.new_node('molecule')
    mol_node.add_neighbors(graph.nodes['atom'])
    return graph

def array_rep_from_smiles(molgraph):
    """Precompute everything we need from MolGraph so that we can free the memory asap."""
    #molgraph = graph_from_smiles_tuple(tuple(smiles))
    degrees = [0,1,2,3,4,5]
    arrayrep = {'atom_features' : molgraph.feature_array('atom'),
                'bond_features' : molgraph.feature_array('bond'),
                'atom_list'     : molgraph.neighbor_list('molecule', 'atom'),
                'rdkit_ix'      : molgraph.rdkit_ix_array()}

    for degree in degrees:
        arrayrep[('atom_neighbors', degree)] = \
            np.array(molgraph.neighbor_list(('atom', degree), 'atom'), dtype=int)
        arrayrep[('bond_neighbors', degree)] = \
            np.array(molgraph.neighbor_list(('atom', degree), 'bond'), dtype=int)
    return arrayrep


def num_atom_features():
    # Return length of feature vector using a very simple molecule.
    m = Chem.MolFromSmiles('CC')
    alist = m.GetAtoms()
    a = alist[0]
    return len(atom_features(a))

def num_bond_features():
    # Return length of feature vector using a very simple molecule.
    simple_mol = Chem.MolFromSmiles('CC')
    Chem.SanitizeMol(simple_mol)
    return len(bond_features(simple_mol.GetBonds()[0]))


def gen_descriptor_data(smilesList):

    smiles_to_fingerprint_array = {}

    for i,smiles in enumerate(smilesList):
#         if i > 5:
#             print("Due to the limited computational resource, submission with more than 5 molecules will not be processed")
#             break
        smiles = Chem.MolToSmiles(Chem.MolFromSmiles(smiles), isomericSmiles=True)
        try:
            molgraph = graph_from_smiles(smiles)
            molgraph.sort_nodes_by_degree('atom')
            arrayrep = array_rep_from_smiles(molgraph)

            smiles_to_fingerprint_array[smiles] = arrayrep
            
        except:
            print(smiles)
            time.sleep(3)
    return smiles_to_fingerprint_array



def get_smiles_dicts(smilesList):
    #first need to get the max atom length
    max_atom_len = 0
    max_bond_len = 0
    num_atom_features = 0
    num_bond_features = 0
    smiles_to_rdkit_list = {}

    smiles_to_fingerprint_features = gen_descriptor_data(smilesList)

    for smiles,arrayrep in smiles_to_fingerprint_features.items():
        atom_features = arrayrep['atom_features']
        bond_features = arrayrep['bond_features']

        rdkit_list = arrayrep['rdkit_ix']
        smiles_to_rdkit_list[smiles] = rdkit_list

        atom_len,num_atom_features = atom_features.shape
        bond_len,num_bond_features = bond_features.shape

        if atom_len > max_atom_len:
            max_atom_len = atom_len
        if bond_len > max_bond_len:
            max_bond_len = bond_len

    #then add 1 so I can zero pad everything
    max_atom_index_num = max_atom_len
    max_bond_index_num = max_bond_len

    max_atom_len += 1
    max_bond_len += 1

    smiles_to_atom_info = {}
    smiles_to_bond_info = {}

    smiles_to_atom_neighbors = {}
    smiles_to_bond_neighbors = {}

    smiles_to_atom_mask = {}

    degrees = [0,1,2,3,4,5]
    #then run through our numpy array again
    for smiles,arrayrep in smiles_to_fingerprint_features.items():
        mask = np.zeros((max_atom_len))

        #get the basic info of what
        #    my atoms and bonds are initialized
        atoms = np.zeros((max_atom_len,num_atom_features))
        bonds = np.zeros((max_bond_len,num_bond_features))

        #then get the arrays initlialized for the neighbors
        atom_neighbors = np.zeros((max_atom_len,len(degrees)))
        bond_neighbors = np.zeros((max_atom_len,len(degrees)))

        #now set these all to the last element of the list, which is zero padded
        atom_neighbors.fill(max_atom_index_num)
        bond_neighbors.fill(max_bond_index_num)

        atom_features = arrayrep['atom_features']
        bond_features = arrayrep['bond_features']

        for i,feature in enumerate(atom_features):
            mask[i] = 1.0
            atoms[i] = feature

        for j,feature in enumerate(bond_features):
            bonds[j] = feature

        atom_neighbor_count = 0
        bond_neighbor_count = 0
        working_atom_list = []
        working_bond_list = []
        for degree in degrees:
            atom_neighbors_list = arrayrep[('atom_neighbors', degree)]
            bond_neighbors_list = arrayrep[('bond_neighbors', degree)]

            if len(atom_neighbors_list) > 0:

                for i,degree_array in enumerate(atom_neighbors_list):
                    for j,value in enumerate(degree_array):
                        atom_neighbors[atom_neighbor_count,j] = value
                    atom_neighbor_count += 1

            if len(bond_neighbors_list) > 0:
                for i,degree_array in enumerate(bond_neighbors_list):
                    for j,value in enumerate(degree_array):
                        bond_neighbors[bond_neighbor_count,j] = value
                    bond_neighbor_count += 1

        #then add everything to my arrays
        smiles_to_atom_info[smiles] = atoms
        smiles_to_bond_info[smiles] = bonds

        smiles_to_atom_neighbors[smiles] = atom_neighbors
        smiles_to_bond_neighbors[smiles] = bond_neighbors
        
        smiles_to_atom_mask[smiles] = mask

    del smiles_to_fingerprint_features
    feature_dicts = {}
#     feature_dicts['smiles_to_atom_mask'] = smiles_to_atom_mask
#     feature_dicts['smiles_to_atom_info']= smiles_to_atom_info
    feature_dicts = {
        'smiles_to_atom_mask': smiles_to_atom_mask,
        'smiles_to_atom_info': smiles_to_atom_info,
        'smiles_to_bond_info': smiles_to_bond_info,
        'smiles_to_atom_neighbors': smiles_to_atom_neighbors,
        'smiles_to_bond_neighbors': smiles_to_bond_neighbors,
        'smiles_to_rdkit_list': smiles_to_rdkit_list
    }
    return feature_dicts
def save_smiles_dicts(smilesList,filename):
    #first need to get the max atom length
    max_atom_len = 0
    max_bond_len = 0
    num_atom_features = 0
    num_bond_features = 0
    smiles_to_rdkit_list = {}

    smiles_to_fingerprint_features = gen_descriptor_data(smilesList)

    for smiles,arrayrep in smiles_to_fingerprint_features.items():
        atom_features = arrayrep['atom_features']
        bond_features = arrayrep['bond_features']

        rdkit_list = arrayrep['rdkit_ix']
        smiles_to_rdkit_list[smiles] = rdkit_list

        atom_len,num_atom_features = atom_features.shape
        bond_len,num_bond_features = bond_features.shape

        if atom_len > max_atom_len:
            max_atom_len = atom_len
        if bond_len > max_bond_len:
            max_bond_len = bond_len

    #then add 1 so I can zero pad everything
    max_atom_index_num = max_atom_len
    max_bond_index_num = max_bond_len

    max_atom_len += 1
    max_bond_len += 1

    smiles_to_atom_info = {}
    smiles_to_bond_info = {}

    smiles_to_atom_neighbors = {}
    smiles_to_bond_neighbors = {}

    smiles_to_atom_mask = {}

    degrees = [0,1,2,3,4,5]
    #then run through our numpy array again
    for smiles,arrayrep in smiles_to_fingerprint_features.items():
        mask = np.zeros((max_atom_len))

        #get the basic info of what
        #    my atoms and bonds are initialized
        atoms = np.zeros((max_atom_len,num_atom_features))
        bonds = np.zeros((max_bond_len,num_bond_features))

        #then get the arrays initlialized for the neighbors
        atom_neighbors = np.zeros((max_atom_len,len(degrees)))
        bond_neighbors = np.zeros((max_atom_len,len(degrees)))

        #now set these all to the last element of the list, which is zero padded
        atom_neighbors.fill(max_atom_index_num)
        bond_neighbors.fill(max_bond_index_num)

        atom_features = arrayrep['atom_features']
        bond_features = arrayrep['bond_features']

        for i,feature in enumerate(atom_features):
            mask[i] = 1.0
            atoms[i] = feature

        for j,feature in enumerate(bond_features):
            bonds[j] = feature

        atom_neighbor_count = 0
        bond_neighbor_count = 0
        working_atom_list = []
        working_bond_list = []
        for degree in degrees:
            atom_neighbors_list = arrayrep[('atom_neighbors', degree)]
            bond_neighbors_list = arrayrep[('bond_neighbors', degree)]

            if len(atom_neighbors_list) > 0:

                for i,degree_array in enumerate(atom_neighbors_list):
                    for j,value in enumerate(degree_array):
                        atom_neighbors[atom_neighbor_count,j] = value
                    atom_neighbor_count += 1

            if len(bond_neighbors_list) > 0:
                for i,degree_array in enumerate(bond_neighbors_list):
                    for j,value in enumerate(degree_array):
                        bond_neighbors[bond_neighbor_count,j] = value
                    bond_neighbor_count += 1

        #then add everything to my arrays
        smiles_to_atom_info[smiles] = atoms
        smiles_to_bond_info[smiles] = bonds

        smiles_to_atom_neighbors[smiles] = atom_neighbors
        smiles_to_bond_neighbors[smiles] = bond_neighbors
        
        smiles_to_atom_mask[smiles] = mask

    del smiles_to_fingerprint_features
    feature_dicts = {}
#     feature_dicts['smiles_to_atom_mask'] = smiles_to_atom_mask
#     feature_dicts['smiles_to_atom_info']= smiles_to_atom_info
    feature_dicts = {
        'smiles_to_atom_mask': smiles_to_atom_mask,
        'smiles_to_atom_info': smiles_to_atom_info,
        'smiles_to_bond_info': smiles_to_bond_info,
        'smiles_to_atom_neighbors': smiles_to_atom_neighbors,
        'smiles_to_bond_neighbors': smiles_to_bond_neighbors,
        'smiles_to_rdkit_list': smiles_to_rdkit_list
    }
    pickle.dump(feature_dicts,open(filename+'.pickle',"wb"))
    print('feature dicts file saved as '+ filename+'.pickle')
    return feature_dicts

def get_smiles_array(smilesList, feature_dicts):
    x_mask = []
    x_atom = []
    x_bonds = []
    x_atom_index = []
    x_bond_index = []
    for smiles in smilesList:
        x_mask.append(feature_dicts['smiles_to_atom_mask'][smiles])
        x_atom.append(feature_dicts['smiles_to_atom_info'][smiles])
        x_bonds.append(feature_dicts['smiles_to_bond_info'][smiles])
        x_atom_index.append(feature_dicts['smiles_to_atom_neighbors'][smiles])
        x_bond_index.append(feature_dicts['smiles_to_bond_neighbors'][smiles])
    return np.asarray(x_atom),np.asarray(x_bonds),np.asarray(x_atom_index),\
        np.asarray(x_bond_index),np.asarray(x_mask),feature_dicts['smiles_to_rdkit_list']


def moltosvg(mol,molSize=(280,218),kekulize=False):
    mc = Chem.Mol(mol.ToBinary())
    if kekulize:
        try:
            Chem.Kekulize(mc)
        except:
            mc = Chem.Mol(mol.ToBinary())
    if not mc.GetNumConformers():
        rdDepictor.Compute2DCoords(mc)
    drawer = rdMolDraw2D.MolDraw2DSVG(molSize[0],molSize[1])
    drawer.DrawMolecule(mc)
    drawer.FinishDrawing()
    svg = drawer.GetDrawingText()
#    return svg
    return svg.replace('svg:','')

def rreplace(s, old, new, occurrence):
    li = s.rsplit(old, occurrence)
    return new.join(li)

def moltosvg_highlight(smiles, atom_list, atom_predictions, molecule_prediction,\
    molSize=(280,218),kekulize=False):
    
    mol = Chem.MolFromSmiles(smiles)
    min_pred = 0.05
    max_pred = 0.8
    # min_pred = np.amax([np.amin(all_atom_predictions),0])
    # min_pred = np.amin(all_atom_predictions)
    # min_pred = min(filter(lambda x:x>0,all_atom_predictions))
    # max_pred = np.amax(all_atom_predictions)
    note = 'y_pred: '+ str(molecule_prediction)

    norm = matplotlib.colors.Normalize(vmin=np.exp(0.068),vmax=np.exp(max_pred))
    cmap = cm.get_cmap('gray_r')


    plt_colors = cm.ScalarMappable(norm=norm,cmap=cmap)

    atom_colors = {}
    for i,atom in enumerate(atom_list):
        color_rgba = plt_colors.to_rgba(atom_predictions[i])
        atom_rgb = color_rgba #(color_rgba[0],color_rgba[1],color_rgba[2])
        atom_colors[atom] = atom_rgb

    rdDepictor.Compute2DCoords(mol)

    drawer = rdMolDraw2D.MolDraw2DSVG(molSize[0],molSize[1])
    #drawer.DrawMolecule(mc)
    drawer.DrawMolecule(mol,highlightAtoms=atom_list,highlightBonds=[],
        highlightAtomColors=atom_colors,legend=note)
    drawer.SetFontSize(68)
    drawer.FinishDrawing()
    svg = drawer.GetDrawingText()
    # svg = svg.replace('size:15px','size:18px')
    # svg = svg.replace('size:12px','size:28px')
    # svg = rreplace(svg,'size:11px','size:28px',1)
    # It seems that the svg renderer used doesn't quite hit the spec.
    # Here are some fixes to make it work in the notebook, although I think
    # the underlying issue needs to be resolved at the generation step
    return svg.replace('svg:','')

def moltosvg_highlight_known(smiles, atom_list, atom_predictions, molecule_prediction, molecule_experiment, Number,\
        molSize=(280,218),kekulize=False):
    
    mol = Chem.MolFromSmiles(smiles)
    min_pred = 0.05
    max_pred = 0.8
    # min_pred = np.amax([np.amin(all_atom_predictions),0])
    # min_pred = np.amin(all_atom_predictions)
    # min_pred = min(filter(lambda x:x>0,all_atom_predictions))
    # max_pred = np.amax(all_atom_predictions)
    note = '('+ str(Number) +') y-y\' : '+ str(round(molecule_experiment,2)) + '-' + str(round(molecule_prediction,2))

    norm = matplotlib.colors.Normalize(vmin=0.01*5,vmax=max_pred*6)
    cmap = cm.get_cmap('gray_r')


    plt_colors = cm.ScalarMappable(norm=norm,cmap=cmap)

    atom_colors = {}
    for i,atom in enumerate(atom_list):
        color_rgba = plt_colors.to_rgba(atom_predictions[i])
        atom_rgb = color_rgba #(color_rgba[0],color_rgba[1],color_rgba[2])
        atom_colors[atom] = atom_rgb

    rdDepictor.Compute2DCoords(mol)

    drawer = rdMolDraw2D.MolDraw2DSVG(molSize[0],molSize[1])
    mol = rdMolDraw2D.PrepareMolForDrawing(mol)
    #drawer.DrawMolecule(mc)
    drawer.DrawMolecule(mol,highlightAtoms=atom_list,highlightBonds=[],
        highlightAtomColors=atom_colors,legend=note)
    drawer.SetFontSize(68)
    drawer.FinishDrawing()
    svg = drawer.GetDrawingText()
    # svg = svg.replace('size:15px','size:18px')
    # svg = svg.replace('size:12px','size:28px')
    # svg = rreplace(svg,'size:11px','size:28px',1)
    # It seems that the svg renderer used doesn't quite hit the spec.
    # Here are some fixes to make it work in the notebook, although I think
    # the underlying issue needs to be resolved at the generation step
    return svg.replace('svg:','')

def weighted_highlight_known(smiles, atom_list, atom_predictions, molecule_prediction, molecule_experiment, Number,\
        molSize=(128,128)):
    
    mol = Chem.MolFromSmiles(smiles)
    note = '('+ str(Number) +') y-y\' : '+ str(round(molecule_experiment,2)) + '-' + str(round(molecule_prediction,2))

    contribs = [atom_predictions[m] for m in np.argsort(atom_list)]
    fig = SimilarityMaps.GetSimilarityMapFromWeights(mol, contribs, colorMap='bwr', contourLines=5, size=molSize)
    fig.axes[0].set_title(note)
    sio = StringIO()
    fig.savefig(sio, format="svg", bbox_inches='tight')
    svg = sio.getvalue()   
    return svg

def moltosvg_interaction_known(mol, atom_list, atom_predictions, molecule_prediction, molecule_experiment, max_atom_pred, min_atom_pred, Number):
    
    note = '('+ str(Number) +') y-y\' : '+ str(round(molecule_experiment,2)) + '-' + str(round(molecule_prediction,2))
    norm = matplotlib.colors.Normalize(vmin=min_atom_pred*0.9,vmax=max_atom_pred*1.1)
    cmap = cm.get_cmap('gray_r')

    plt_colors = cm.ScalarMappable(norm=norm,cmap=cmap)

    atom_colors = {}
    for i,atom in enumerate(atom_list):
        atom_colors[atom] = plt_colors.to_rgba(atom_predictions[i])          
    rdDepictor.Compute2DCoords(mol)

    drawer = rdMolDraw2D.MolDraw2DSVG(280,218)
    op = drawer.drawOptions()
    for i in range(mol.GetNumAtoms()):
        op.atomLabels[i]=mol.GetAtomWithIdx(i).GetSymbol() + str(i)

    mol = rdMolDraw2D.PrepareMolForDrawing(mol)
    drawer.DrawMolecule(mol,highlightAtoms=atom_list,highlightBonds=[],
        highlightAtomColors=atom_colors,legend=note)
    drawer.SetFontSize(68)
    drawer.FinishDrawing()
    svg = drawer.GetDrawingText()
    return svg.replace('svg:','')

src/PointGAT/getFeatures_aromaticity_rm.py
sevencheung2021/PointGAT
'''
Description: 
version: 
Email: rongzhangthu@yeah.net
Author: Rong Zhang
Date: 2022-04-13 16:51:50
LastEditors: Rong Zhang
LastEditTime: 2022-06-17 10:01:21
'''
import torch
import math
import torch.nn as nn
from rdkit import Chem
from rdkit import rdBase
rdBase.DisableLog('rdApp.*')

# Split SMILES into words
def split(sm):
    '''
    function: Split SMILES into words. 
    input: A SMILES
    output: A string with space between words
    '''
    arr = []
    i = 0
    while i < len(sm)-1:
        if not sm[i] in ['%', 'C']:
            arr.append(sm[i])
            i += 1
        elif sm[i]=='%':
            arr.append(sm[i:i+3])
            i += 3
        elif sm[i]=='C' and sm[i+1]=='+':
            arr.append(sm[i:i+2])
            i += 2
        elif sm[i]=='C' and sm[i+1]=='H' and sm[i+2]=='+':
            arr.append(sm[i:i+3])
            i += 3
        elif sm[i]=='C' and sm[i+1]=='H' and sm[i+2]=='2' and sm[i+3]=='+':
            arr.append(sm[i:i+4])
            i += 4
        else:
            arr.append(sm[i])
            i += 1
    if i == len(sm)-1:
        arr.append(sm[i])
    return ' '.join(arr) 

# activation
class GELU(nn.Module):
    def forward(self, x):
        return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))

# positional
class PositionwiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)
        self.activation = GELU()

    def forward(self, x):
        return self.w_2(self.dropout(self.activation(self.w_1(x))))
    
# regularization
class LayerNorm(nn.Module):
    def __init__(self, features, eps=1e-6):
        super(LayerNorm, self).__init__()
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps

    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2


class SublayerConnection(nn.Module):
    def __init__(self, size, dropout):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, sublayer):
        return x + self.dropout(sublayer(self.norm(x)))

# Sample SMILES from probablistic distribution
def sample(msms):
    ret = []
    for msm in msms:
        ret.append(torch.multinomial(msm.exp(), 1).squeeze())
    return torch.stack(ret)

def validity(smiles):
    loss = 0
    for sm in smiles:
        mol = Chem.MolFromSmiles(sm)
        if mol is None:
            loss += 1
    return 1-loss/len(smiles)


src/PointGAT/utils.py
sevencheung2021/PointGAT
'''
Description: 
version: 
Email: rongzhangthu@yeah.net
Author: Rong Zhang
Date: 2022-06-13 18:10:11
LastEditors: Seven Rong Cheung
LastEditTime: 2023-09-28 11:06:13
'''


import os
import copy
import json
import pickle
import time
import sys
import numpy as np
import pandas as pd


# torch
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim


sys.setrecursionlimit(50000)
torch.backends.cudnn.benchmark = True
torch.set_default_tensor_type('torch.cuda.FloatTensor')
torch.nn.Module.dump_patches = True


from PointGAT.PointGAT_Layers import PointGAT
from PointGAT import save_smiles_dicts, get_smiles_dicts, get_smiles_array, moltosvg_highlight

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


# log files to write
class Logger_print():
    def __init__(self,filename = 'default log',stream = sys.stdout):
        self.terminal = stream
        self.log = open(filename,'w')

    def write(self,message):
        self.terminal.write(message)
        self.log.write(message)

    def flush(self):
        pass

sys.stdout = Logger_print('C10_mmff_lowest_2023.log',sys.stdout)


##  generate feature rb files 
def generateGeoFeatureFile(file_path,file_name):
    raw_filename = os.path.join(file_path,file_name)
    feature_filename = raw_filename.replace('.csv','.pickle')
    filename = raw_filename.replace('.csv','')
    smiles_tasks_df = pd.read_csv(raw_filename,index_col=False)
    smiles_tasks_df.sample().reset_index(drop=True)
    smilesList = smiles_tasks_df.smiles.tolist()
    
    print("number of {} smiles {}: ".format(filename,len(smilesList)))
    
    if os.path.isfile(feature_filename):
        print('loading*****************************************')
        feature_dicts = pickle.load(open(feature_filename, "rb" ))
        # print(feature_dicts)
    else:
        print('generating*****************************************')
        feature_dicts = save_smiles_dicts(smilesList,filename)
    return smiles_tasks_df,feature_dicts


## train
def train(model, dataset, optimizer, loss_function,feature_dicts,tasks,batch_size,xyz_feature):
    model.train()
    np.random.seed(8)
    valList = np.arange(0,dataset.shape[0])
    #shuffle them
    np.random.shuffle(valList)
    batch_list = []
    for i in range(0, dataset.shape[0], batch_size):
        batch = valList[i:i+batch_size]
        batch_list.append(batch)   
    for counter, train_batch in enumerate(batch_list):
        batch_df = dataset.loc[train_batch,:]
        smiles_list = batch_df.smiles.values
        y_val = batch_df[tasks[0]].values
        
        
        # print(data)
        x_atom, x_bonds, x_atom_index, x_bond_index, x_mask, smiles_to_rdkit_lis = get_smiles_array(smiles_list,feature_dicts)
        
        x_xyz = batch_df[xyz_feature[0]].values
        x_xyz = [json.loads(val) for val in x_xyz]   
        x_xyz = np.asarray(x_xyz)
        
        # print(x_atom)
        mol_prediction = model(torch.Tensor(x_atom),torch.Tensor(x_bonds),torch.cuda.LongTensor(x_atom_index),torch.cuda.LongTensor(x_bond_index),torch.Tensor(x_mask),torch.Tensor(x_xyz))
        
        model.zero_grad()
        loss = loss_function(mol_prediction, torch.Tensor(y_val).view(-1,1))     
        loss.backward()
        optimizer.step()


## eval 
def eval(model, dataset,feature_dicts,tasks,batch_size,xyz_feature):
    model.eval()
    eval_MAE_list = []
    eval_MSE_list = []
    valList = np.arange(0,dataset.shape[0])
    batch_list = []
    for i in range(0, dataset.shape[0], batch_size):
        batch = valList[i:i+batch_size]
        batch_list.append(batch) 
    for counter, eval_batch in enumerate(batch_list):
        batch_df = dataset.loc[eval_batch,:]
        smiles_list = batch_df.smiles.values
#         print(batch_df)
        y_val = batch_df[tasks[0]].values

        x_atom, x_bonds, x_atom_index, x_bond_index, x_mask, smiles_to_rdkit_list= get_smiles_array(smiles_list,feature_dicts)

        x_xyz = batch_df[xyz_feature[0]].values
        x_xyz = [json.loads(val) for val in x_xyz]   
        x_xyz = np.asarray(x_xyz)

        mol_prediction = model(torch.Tensor(x_atom),torch.Tensor(x_bonds),torch.cuda.LongTensor(x_atom_index),torch.cuda.LongTensor(x_bond_index),torch.Tensor(x_mask),torch.Tensor(x_xyz))
        
        MAE = F.l1_loss(mol_prediction, torch.Tensor(y_val).view(-1,1), reduction='none')        
        MSE = F.mse_loss(mol_prediction, torch.Tensor(y_val).view(-1,1), reduction='none')
#      
        eval_MAE_list.extend(MAE.data.squeeze().cpu().numpy())
        # print(len(eval_MAE_list))
        eval_MSE_list.extend(MSE.data.squeeze().cpu().numpy())
    return np.array(eval_MAE_list).mean(), np.array(eval_MSE_list).mean()


def main():

    ## model parameters
    start_time = str(time.ctime()).replace(':','-').replace(' ','_')
    tasks = ['normalize_energy']
    xyz_feature = ['xyz_feature']
    batch_size = 32
    epochs = 800
    p_dropout= 0.1
    fingerprint_dim = 215
    weight_decay = 6.0 # also known as l2_regularization_lambda
    learning_rate = 3.5
    output_units_num = 1 # for regression model
    radius = 6
    T = 4
    xyz_feature_dim = 6

    ## data path
    data_file_path = '../Data/'
    train_file_name = 'c10_train.csv'
    valid_file_name = 'c10_valid.csv'
    test_file_name = 'c10_test.csv'
    
    ## load data
    train_df,train_feature_dicts = generateGeoFeatureFile(file_path=data_file_path,file_name=train_file_name)
    valid_df,valid_feature_dicts = generateGeoFeatureFile(file_path=data_file_path,file_name=valid_file_name)
    test_df,test_feature_dicts = generateGeoFeatureFile(file_path=data_file_path,file_name=test_file_name)
    
    x_atom, x_bonds, x_atom_index, x_bond_index, x_mask, smiles_to_rdkit_list = get_smiles_array([train_df.smiles.values[0]],train_feature_dicts)
    num_atom_features = x_atom.shape[-1]
    num_bond_features = x_bonds.shape[-1]
    
    
    # model setup
    loss_function = nn.MSELoss()
    model = PointGAT(radius, T, num_atom_features, num_bond_features,
            fingerprint_dim, output_units_num, p_dropout,xyz_feature_dim)
    model.cuda()
    optimizer = optim.Adam(model.parameters(), 10**-learning_rate, weight_decay=10**-weight_decay)
    model_parameters = filter(lambda p: p.requires_grad, model.parameters())
    params = sum([np.prod(p.size()) for p in model_parameters])
    # print('**** model params:{}'.format(params))

    ## perform training process
    best_param ={}
    best_param["train_epoch"] = 0
    best_param["valid_epoch"] = 0
    best_param["train_MAE"] = 9e8
    best_param["valid_MAE"] = 9e8

    collector= {k: [] for k in ['epoch',
                                    'train_loss',
                                    'valid_loss',
                                    'train_RMSE',
                                    'valid_RMSE',
                                    'train_MAE',
                                    'valid_MAE']}
    print('training*************************************************************')
    for epoch in range(epochs):
        
        train_MAE, train_MSE = eval(model, train_df,train_feature_dicts,tasks,batch_size,xyz_feature)
        valid_MAE, valid_MSE = eval(model, valid_df,valid_feature_dicts,tasks,batch_size,xyz_feature)

        if train_MAE < best_param["train_MAE"]:
            best_param["train_epoch"] = epoch
            best_param["train_MAE"] = train_MAE
        if valid_MAE < best_param["valid_MAE"]:
            best_param["valid_epoch"] = epoch
            best_param["valid_MAE"] = valid_MAE
            
            if valid_MAE < 0.025:

                torch.save(model, '../weights/model_'+'C10_'+start_time+'_'+str(best_param["valid_epoch"])+'.pt')

        if epoch - best_param["valid_epoch"] >100:        
            break

        print("Epoch",epoch)
        print('train_RMSE and valid_RMSE')
        print(epoch, np.sqrt(train_MSE), np.sqrt(valid_MSE))
        print('train_MAE and valid_MAE')
        print(epoch, train_MAE,valid_MAE)

        ## training 
        train(model, train_df, optimizer, loss_function,train_feature_dicts,tasks,batch_size,xyz_feature)

        collector['epoch'].append(epoch+1)
        collector['train_loss'].append(train_MSE)
        collector['train_RMSE'].append(np.sqrt(train_MSE))
        collector['train_MAE'].append(train_MAE)
        
        collector['valid_loss'].append(valid_MSE)
        collector['valid_RMSE'].append(np.sqrt(valid_MSE))
        collector['valid_MAE'].append(valid_MAE)

        pd.DataFrame(collector).to_csv('C10_pointgat_training_result.csv', index=False)
    

    # evaluate model
    print('evaluate*************************************************************')
    best_model = torch.load('../weights/model_'+'C10_'+start_time+'_'+str(best_param["valid_epoch"])+'.pt')

    best_model_dict = best_model.state_dict()
    best_model_wts = copy.deepcopy(best_model_dict)

    model.load_state_dict(best_model_wts)
    (best_model.align[0].weight == model.align[0].weight).all()
    test_MAE, test_MSE = eval(model, test_df,test_feature_dicts,tasks,batch_size,xyz_feature)

    print("best epoch:",best_param["valid_epoch"],"\n","test RMSE:",np.sqrt(test_MSE),"test MAE:",test_MAE)
    


if __name__ =='__main__':

    main()


src/pointgat.py
sevencheung2021/PointGAT