dmlc--dgl
Mufei Li
828a5e5bc6
* First commit * Update * Update splitters * Update * Update * Update * Update * Update * Update * Migrate ACNN * Fix * Fix * Update * Update * Update * Update * Update * Update * Finish classification * Update * Fix * Update * Update * Update * Fix * Fix * Fix * Update * Update * Update * trigger CI * Fix CI * Update * Update * Update * Add default values * Rename * Update deprecation message
1041 行
32 KiB
Python
1041 行
32 KiB
Python
import datetime
|
|
import math
|
|
import numpy as np
|
|
import os
|
|
import pickle
|
|
import random
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.nn as nn
|
|
from collections import defaultdict
|
|
from datetime import timedelta
|
|
from dgl.data.utils import download, _get_dgl_url
|
|
from dgllife.model.model_zoo.dgmg import MoleculeEnv
|
|
from multiprocessing import Pool
|
|
from pprint import pprint
|
|
from rdkit import Chem
|
|
from rdkit.Chem import AllChem
|
|
from rdkit.Chem.Crippen import MolLogP
|
|
from rdkit.Chem.QED import qed
|
|
from torch.utils.data import Dataset
|
|
|
|
from sascorer import calculateScore
|
|
|
|
########################################################################################################################
|
|
# configuration #
|
|
########################################################################################################################
|
|
|
|
def mkdir_p(path, log=True):
|
|
"""Create a directory for the specified path.
|
|
|
|
Parameters
|
|
----------
|
|
path : str
|
|
Path name
|
|
log : bool
|
|
Whether to print result for directory creation
|
|
"""
|
|
import errno
|
|
try:
|
|
os.makedirs(path)
|
|
if log:
|
|
print('Created directory {}'.format(path))
|
|
except OSError as exc:
|
|
if exc.errno == errno.EEXIST and os.path.isdir(path) and log:
|
|
print('Directory {} already exists.'.format(path))
|
|
else:
|
|
raise
|
|
|
|
def get_date_postfix():
|
|
"""Get a date based postfix for directory name.
|
|
|
|
Returns
|
|
-------
|
|
post_fix : str
|
|
"""
|
|
dt = datetime.datetime.now()
|
|
post_fix = '{}_{:02d}-{:02d}-{:02d}'.format(
|
|
dt.date(), dt.hour, dt.minute, dt.second)
|
|
|
|
return post_fix
|
|
|
|
def setup_log_dir(args):
|
|
"""Name and create directory for logging.
|
|
|
|
Parameters
|
|
----------
|
|
args : dict
|
|
Configuration
|
|
|
|
Returns
|
|
-------
|
|
log_dir : str
|
|
Path for logging directory
|
|
"""
|
|
date_postfix = get_date_postfix()
|
|
log_dir = os.path.join(
|
|
args['log_dir'],
|
|
'{}_{}_{}'.format(args['dataset'], args['order'], date_postfix))
|
|
mkdir_p(log_dir)
|
|
return log_dir
|
|
|
|
def save_arg_dict(args, filename='settings.txt'):
|
|
"""Save all experiment settings in a file.
|
|
|
|
Parameters
|
|
----------
|
|
args : dict
|
|
Configuration
|
|
filename : str
|
|
Name for the file to save settings
|
|
"""
|
|
def _format_value(v):
|
|
if isinstance(v, float):
|
|
return '{:.4f}'.format(v)
|
|
elif isinstance(v, int):
|
|
return '{:d}'.format(v)
|
|
else:
|
|
return '{}'.format(v)
|
|
|
|
save_path = os.path.join(args['log_dir'], filename)
|
|
with open(save_path, 'w') as f:
|
|
for key, value in args.items():
|
|
f.write('{}\t{}\n'.format(key, _format_value(value)))
|
|
print('Saved settings to {}'.format(save_path))
|
|
|
|
def configure(args):
|
|
"""Use default hyperparameters.
|
|
|
|
Parameters
|
|
----------
|
|
args : dict
|
|
Old configuration
|
|
|
|
Returns
|
|
-------
|
|
args : dict
|
|
Updated configuration
|
|
"""
|
|
configure = {
|
|
'node_hidden_size': 128,
|
|
'num_propagation_rounds': 2,
|
|
'lr': 1e-4,
|
|
'dropout': 0.2,
|
|
'nepochs': 400,
|
|
'batch_size': 1,
|
|
}
|
|
args.update(configure)
|
|
return args
|
|
|
|
def set_random_seed(seed):
|
|
"""Fix random seed for reproducible results.
|
|
|
|
Parameters
|
|
----------
|
|
seed : int
|
|
Random seed to use.
|
|
"""
|
|
random.seed(seed)
|
|
np.random.seed(seed)
|
|
torch.manual_seed(seed)
|
|
|
|
def setup_dataset(args):
|
|
"""Dataset setup
|
|
|
|
For unsupported dataset, we need to perform data preprocessing.
|
|
|
|
Parameters
|
|
----------
|
|
args : dict
|
|
Configuration
|
|
"""
|
|
if args['dataset'] in ['ChEMBL', 'ZINC']:
|
|
print('Built-in support for dataset {} exists.'.format(args['dataset']))
|
|
else:
|
|
print('Configure for new dataset {}...'.format(args['dataset']))
|
|
configure_new_dataset(args['dataset'], args['train_file'], args['val_file'])
|
|
|
|
def setup(args, train=True):
|
|
"""Setup
|
|
|
|
Parameters
|
|
----------
|
|
args : argparse.Namespace
|
|
Configuration
|
|
train : bool
|
|
Whether the setup is for training or evaluation
|
|
"""
|
|
# Convert argparse.Namespace into a dict
|
|
args = args.__dict__.copy()
|
|
# Dataset
|
|
args = configure(args)
|
|
|
|
# Log
|
|
print('Prepare logging directory...')
|
|
log_dir = setup_log_dir(args)
|
|
args['log_dir'] = log_dir
|
|
save_arg_dict(args)
|
|
|
|
if train:
|
|
setup_dataset(args)
|
|
args['checkpoint_dir'] = os.path.join(log_dir, 'checkpoint.pth')
|
|
pprint(args)
|
|
|
|
return args
|
|
|
|
########################################################################################################################
|
|
# multi-process #
|
|
########################################################################################################################
|
|
|
|
def synchronize(num_processes):
|
|
"""Synchronize all processes.
|
|
|
|
Parameters
|
|
----------
|
|
num_processes : int
|
|
Number of subprocesses used
|
|
"""
|
|
if num_processes > 1:
|
|
dist.barrier()
|
|
|
|
def launch_a_process(rank, args, target, minutes=720):
|
|
"""Launch a subprocess for training.
|
|
|
|
Parameters
|
|
----------
|
|
rank : int
|
|
Subprocess id
|
|
args : dict
|
|
Configuration
|
|
target : callable
|
|
Target function for the subprocess
|
|
minutes : int
|
|
Timeout minutes for operations executed against the process group
|
|
"""
|
|
dist_init_method = 'tcp://{master_ip}:{master_port}'.format(
|
|
master_ip=args['master_ip'], master_port=args['master_port'])
|
|
dist.init_process_group(backend='gloo',
|
|
init_method=dist_init_method,
|
|
# If you have a larger dataset, you will need to increase it.
|
|
timeout=timedelta(minutes=minutes),
|
|
world_size=args['num_processes'],
|
|
rank=rank)
|
|
assert torch.distributed.get_rank() == rank
|
|
target(rank, args)
|
|
|
|
########################################################################################################################
|
|
# optimization #
|
|
########################################################################################################################
|
|
|
|
class Optimizer(nn.Module):
|
|
"""Wrapper for optimization
|
|
|
|
Parameters
|
|
----------
|
|
lr : float
|
|
Initial learning rate
|
|
optimizer
|
|
model optimizer
|
|
"""
|
|
def __init__(self, lr, optimizer):
|
|
super(Optimizer, self).__init__()
|
|
self.lr = lr
|
|
self.optimizer = optimizer
|
|
self._reset()
|
|
|
|
def _reset(self):
|
|
self.optimizer.zero_grad()
|
|
|
|
def backward_and_step(self, loss):
|
|
"""Backward and update model.
|
|
|
|
Parameters
|
|
----------
|
|
loss : torch.tensor consisting of a float only
|
|
"""
|
|
loss.backward()
|
|
self.optimizer.step()
|
|
self._reset()
|
|
|
|
def decay_lr(self, decay_rate=0.99):
|
|
"""Decay learning rate.
|
|
|
|
Parameters
|
|
----------
|
|
decay_rate : float
|
|
Multiply the current learning rate by the decay_rate
|
|
"""
|
|
self.lr *= decay_rate
|
|
for param_group in self.optimizer.param_groups:
|
|
param_group['lr'] = self.lr
|
|
|
|
class MultiProcessOptimizer(Optimizer):
|
|
"""Wrapper for optimization with multiprocess
|
|
|
|
Parameters
|
|
----------
|
|
n_processes : int
|
|
Number of processes used
|
|
lr : float
|
|
Initial learning rate
|
|
optimizer
|
|
model optimizer
|
|
"""
|
|
def __init__(self, n_processes, lr, optimizer):
|
|
super(MultiProcessOptimizer, self).__init__(lr=lr, optimizer=optimizer)
|
|
self.n_processes = n_processes
|
|
|
|
def _sync_gradient(self):
|
|
"""Average gradients across all subprocesses."""
|
|
for param_group in self.optimizer.param_groups:
|
|
for p in param_group['params']:
|
|
if p.requires_grad and p.grad is not None:
|
|
dist.all_reduce(p.grad.data, op=dist.ReduceOp.SUM)
|
|
p.grad.data /= self.n_processes
|
|
|
|
def backward_and_step(self, loss):
|
|
"""Backward and update model.
|
|
|
|
Parameters
|
|
----------
|
|
loss : torch.tensor consisting of a float only
|
|
"""
|
|
loss.backward()
|
|
self._sync_gradient()
|
|
self.optimizer.step()
|
|
self._reset()
|
|
|
|
########################################################################################################################
|
|
# data #
|
|
########################################################################################################################
|
|
|
|
def initialize_neuralization_reactions():
|
|
"""Reference neuralization reactions
|
|
|
|
Code adapted from RDKit Cookbook, by Hans de Winter.
|
|
"""
|
|
patts = (
|
|
# Imidazoles
|
|
('[n+;H]', 'n'),
|
|
# Amines
|
|
('[N+;!H0]', 'N'),
|
|
# Carboxylic acids and alcohols
|
|
('[$([O-]);!$([O-][#7])]', 'O'),
|
|
# Thiols
|
|
('[S-;X1]', 'S'),
|
|
# Sulfonamides
|
|
('[$([N-;X2]S(=O)=O)]', 'N'),
|
|
# Enamines
|
|
('[$([N-;X2][C,N]=C)]', 'N'),
|
|
# Tetrazoles
|
|
('[n-]', '[n]'),
|
|
# Sulfoxides
|
|
('[$([S-]=O)]', 'S'),
|
|
# Amides
|
|
('[$([N-]C=O)]', 'N'),
|
|
)
|
|
return [(Chem.MolFromSmarts(x), Chem.MolFromSmiles(y, False)) for x, y in patts]
|
|
|
|
def neutralize_charges(mol, reactions=None):
|
|
"""Deprotonation for molecules.
|
|
|
|
Code adapted from RDKit Cookbook, by Hans de Winter.
|
|
|
|
DGMG currently cannot generate protonated molecules.
|
|
For example, it can only generate
|
|
CC(C)(C)CC1CCC[NH+]1Cc1nnc(-c2ccccc2F)o1
|
|
from
|
|
CC(C)(C)CC1CCCN1Cc1nnc(-c2ccccc2F)o1
|
|
even with correct decisions.
|
|
|
|
Deprotonation is therefore an important step to avoid
|
|
false novel molecules.
|
|
|
|
Parameters
|
|
----------
|
|
mol : Chem.rdchem.Mol
|
|
reactions : list of 2-tuples
|
|
Rules for deprotonation
|
|
|
|
Returns
|
|
-------
|
|
mol : Chem.rdchem.Mol
|
|
Deprotonated molecule
|
|
"""
|
|
if reactions is None:
|
|
reactions = initialize_neuralization_reactions()
|
|
for i, (reactant, product) in enumerate(reactions):
|
|
while mol.HasSubstructMatch(reactant):
|
|
rms = AllChem.ReplaceSubstructs(mol, reactant, product)
|
|
mol = rms[0]
|
|
return mol
|
|
|
|
def standardize_mol(mol):
|
|
"""Standardize molecule to avoid false novel molecule.
|
|
|
|
Kekulize and deprotonate molecules to avoid false novel molecules.
|
|
|
|
In addition to deprotonation, we also kekulize molecules to avoid
|
|
explicit Hs in the SMILES. Otherwise we will get false novel molecules
|
|
as well. For example, DGMG can only generate
|
|
O=S(=O)(NC1=CC=CC(C(F)(F)F)=C1)C1=CNC=N1
|
|
from
|
|
O=S(=O)(Nc1cccc(C(F)(F)F)c1)c1c[nH]cn1.
|
|
|
|
One downside is that we remove all explicit aromatic rings and to
|
|
explicitly predict aromatic bond might make the learning easier for
|
|
the model.
|
|
"""
|
|
reactions = initialize_neuralization_reactions()
|
|
Chem.Kekulize(mol, clearAromaticFlags=True)
|
|
mol = neutralize_charges(mol, reactions)
|
|
return mol
|
|
|
|
def smiles_to_standard_mol(s):
|
|
"""Convert SMILES to a standard molecule.
|
|
|
|
Parameters
|
|
----------
|
|
s : str
|
|
SMILES
|
|
|
|
Returns
|
|
-------
|
|
Chem.rdchem.Mol
|
|
Standardized molecule
|
|
"""
|
|
mol = Chem.MolFromSmiles(s)
|
|
return standardize_mol(mol)
|
|
|
|
def mol_to_standard_smile(mol):
|
|
"""Standardize a molecule and convert it to a SMILES.
|
|
|
|
Parameters
|
|
----------
|
|
mol : Chem.rdchem.Mol
|
|
|
|
Returns
|
|
-------
|
|
str
|
|
SMILES
|
|
"""
|
|
return Chem.MolToSmiles(standardize_mol(mol))
|
|
|
|
def get_atom_and_bond_types(smiles, log=True):
|
|
"""Identify the atom types and bond types
|
|
appearing in this dataset.
|
|
|
|
Parameters
|
|
----------
|
|
smiles : list
|
|
List of smiles
|
|
log : bool
|
|
Whether to print the process of pre-processing.
|
|
|
|
Returns
|
|
-------
|
|
atom_types : list
|
|
E.g. ['C', 'N']
|
|
bond_types : list
|
|
E.g. [Chem.rdchem.BondType.SINGLE, Chem.rdchem.BondType.DOUBLE,
|
|
Chem.rdchem.BondType.TRIPLE, Chem.rdchem.BondType.AROMATIC]
|
|
"""
|
|
atom_types = set()
|
|
bond_types = set()
|
|
n_smiles = len(smiles)
|
|
|
|
for i, s in enumerate(smiles):
|
|
if log:
|
|
print('Processing smiles {:d}/{:d}'.format(i + 1, n_smiles))
|
|
|
|
mol = smiles_to_standard_mol(s)
|
|
if mol is None:
|
|
continue
|
|
|
|
for atom in mol.GetAtoms():
|
|
a_symbol = atom.GetSymbol()
|
|
if a_symbol not in atom_types:
|
|
atom_types.add(a_symbol)
|
|
|
|
for bond in mol.GetBonds():
|
|
b_type = bond.GetBondType()
|
|
if b_type not in bond_types:
|
|
bond_types.add(b_type)
|
|
|
|
return list(atom_types), list(bond_types)
|
|
|
|
def eval_decisions(env, decisions):
|
|
"""This function mimics the way DGMG generates a molecule and is
|
|
helpful for debugging and verification in data preprocessing.
|
|
|
|
Parameters
|
|
----------
|
|
env : MoleculeEnv
|
|
MDP environment for generating molecules
|
|
decisions : list of 2-tuples of int
|
|
A decision sequence for generating a molecule
|
|
|
|
Returns
|
|
-------
|
|
str
|
|
SMILES for the molecule generated with decisions
|
|
"""
|
|
env.reset(rdkit_mol=True)
|
|
t = 0
|
|
|
|
def whether_to_add_atom(t):
|
|
assert decisions[t][0] == 0
|
|
atom_type = decisions[t][1]
|
|
t += 1
|
|
return t, atom_type
|
|
|
|
def whether_to_add_bond(t):
|
|
assert decisions[t][0] == 1
|
|
bond_type = decisions[t][1]
|
|
t += 1
|
|
return t, bond_type
|
|
|
|
def decide_atom2(t):
|
|
assert decisions[t][0] == 2
|
|
dst = decisions[t][1]
|
|
t += 1
|
|
return t, dst
|
|
|
|
t, atom_type = whether_to_add_atom(t)
|
|
while atom_type != len(env.atom_types):
|
|
env.add_atom(atom_type)
|
|
t, bond_type = whether_to_add_bond(t)
|
|
while bond_type != len(env.bond_types):
|
|
t, dst = decide_atom2(t)
|
|
env.add_bond((env.num_atoms() - 1), dst, bond_type)
|
|
t, bond_type = whether_to_add_bond(t)
|
|
t, atom_type = whether_to_add_atom(t)
|
|
assert t == len(decisions)
|
|
|
|
return env.get_current_smiles()
|
|
|
|
def get_DGMG_smile(env, mol):
|
|
"""Mimics the reproduced SMILES with DGMG for a molecule.
|
|
|
|
Given a molecule, we are interested in what SMILES we will
|
|
get if we want to generate it with DGMG. This is an important
|
|
step to check false novel molecules.
|
|
|
|
Parameters
|
|
----------
|
|
env : MoleculeEnv
|
|
MDP environment for generating molecules
|
|
mol : Chem.rdchem.Mol
|
|
A molecule
|
|
|
|
Returns
|
|
-------
|
|
canonical_smile : str
|
|
SMILES of the generated molecule with a canonical decision sequence
|
|
random_smile : str
|
|
SMILES of the generated molecule with a random decision sequence
|
|
"""
|
|
canonical_decisions = env.get_decision_sequence(mol, list(range(mol.GetNumAtoms())))
|
|
canonical_smile = eval_decisions(env, canonical_decisions)
|
|
|
|
order = list(range(mol.GetNumAtoms()))
|
|
random.shuffle(order)
|
|
random_decisions = env.get_decision_sequence(mol, order)
|
|
random_smile = eval_decisions(env, random_decisions)
|
|
|
|
return canonical_smile, random_smile
|
|
|
|
def preprocess_dataset(atom_types, bond_types, smiles, max_num_atoms=23):
|
|
"""Preprocess the dataset
|
|
|
|
1. Standardize the SMILES of the dataset
|
|
2. Only keep the SMILES that DGMG can reproduce
|
|
3. Drop repeated SMILES
|
|
|
|
Parameters
|
|
----------
|
|
atom_types : list
|
|
The types of atoms appearing in a dataset. E.g. ['C', 'N']
|
|
bond_types : list
|
|
The types of bonds appearing in a dataset.
|
|
E.g. [Chem.rdchem.BondType.SINGLE, Chem.rdchem.BondType.DOUBLE,
|
|
Chem.rdchem.BondType.TRIPLE, Chem.rdchem.BondType.AROMATIC]
|
|
|
|
Returns
|
|
-------
|
|
valid_smiles : list of str
|
|
SMILES left after preprocessing
|
|
"""
|
|
valid_smiles = []
|
|
env = MoleculeEnv(atom_types, bond_types)
|
|
|
|
for id, s in enumerate(smiles):
|
|
print('Processing {:d}/{:d}'.format(id + 1, len(smiles)))
|
|
|
|
raw_s = s.strip()
|
|
mol = smiles_to_standard_mol(raw_s)
|
|
if mol is None:
|
|
continue
|
|
standard_s = Chem.MolToSmiles(mol)
|
|
|
|
if (max_num_atoms is not None) and (mol.GetNumAtoms() > max_num_atoms):
|
|
continue
|
|
|
|
canonical_s, random_s = get_DGMG_smile(env, mol)
|
|
canonical_mol = Chem.MolFromSmiles(canonical_s)
|
|
random_mol = Chem.MolFromSmiles(random_s)
|
|
|
|
if (standard_s != canonical_s) or (canonical_s != random_s) or (canonical_mol is None) or (random_mol is None):
|
|
continue
|
|
|
|
valid_smiles.append(standard_s)
|
|
|
|
valid_smiles = list(set(valid_smiles))
|
|
return valid_smiles
|
|
|
|
def download_data(dataset, fname):
|
|
"""Download dataset if built-in support exists
|
|
|
|
Parameters
|
|
----------
|
|
dataset : str
|
|
Dataset name
|
|
fname : str
|
|
Name of dataset file
|
|
"""
|
|
if dataset not in ['ChEMBL', 'ZINC']:
|
|
# For dataset without built-in support, they should be locally processed.
|
|
return
|
|
|
|
data_path = fname
|
|
download(_get_dgl_url(os.path.join('dataset', fname)), path=data_path)
|
|
|
|
def load_smiles_from_file(f_name):
|
|
"""Load dataset into a list of SMILES
|
|
|
|
Parameters
|
|
----------
|
|
f_name : str
|
|
Path to a file of molecules, where each line of the file
|
|
is a molecule in SMILES format.
|
|
|
|
Returns
|
|
-------
|
|
smiles : list of str
|
|
List of molecules as SMILES
|
|
"""
|
|
with open(f_name, 'r') as f:
|
|
smiles = f.read().splitlines()
|
|
return smiles
|
|
|
|
def write_smiles_to_file(f_name, smiles):
|
|
"""Write dataset to a file.
|
|
|
|
Parameters
|
|
----------
|
|
f_name : str
|
|
Path to create a file of molecules, where each line of the file
|
|
is a molecule in SMILES format.
|
|
smiles : list of str
|
|
List of SMILES
|
|
"""
|
|
with open(f_name, 'w') as f:
|
|
for s in smiles:
|
|
f.write(s + '\n')
|
|
|
|
def configure_new_dataset(dataset, train_file, val_file):
|
|
"""Configure for a new dataset.
|
|
|
|
Parameters
|
|
----------
|
|
dataset : str
|
|
Dataset name
|
|
train_file : str
|
|
Path to a file with one SMILES a line for training data
|
|
val_file : str
|
|
Path to a file with one SMILES a line for validation data
|
|
"""
|
|
assert train_file is not None, 'Expect a file of SMILES for training, got None.'
|
|
assert val_file is not None, 'Expect a file of SMILES for validation, got None.'
|
|
train_smiles = load_smiles_from_file(train_file)
|
|
val_smiles = load_smiles_from_file(val_file)
|
|
all_smiles = train_smiles + val_smiles
|
|
|
|
# Get all atom and bond types in the dataset
|
|
path_to_atom_and_bond_types = '_'.join([dataset, 'atom_and_bond_types.pkl'])
|
|
if not os.path.exists(path_to_atom_and_bond_types):
|
|
atom_types, bond_types = get_atom_and_bond_types(all_smiles)
|
|
with open(path_to_atom_and_bond_types, 'wb') as f:
|
|
pickle.dump({'atom_types': atom_types, 'bond_types': bond_types}, f)
|
|
else:
|
|
with open(path_to_atom_and_bond_types, 'rb') as f:
|
|
type_info = pickle.load(f)
|
|
atom_types = type_info['atom_types']
|
|
bond_types = type_info['bond_types']
|
|
|
|
# Standardize training data
|
|
path_to_processed_train_data = '_'.join([dataset, 'DGMG', 'train.txt'])
|
|
if not os.path.exists(path_to_processed_train_data):
|
|
processed_train_smiles = preprocess_dataset(atom_types, bond_types, train_smiles, None)
|
|
write_smiles_to_file(path_to_processed_train_data, processed_train_smiles)
|
|
|
|
path_to_processed_val_data = '_'.join([dataset, 'DGMG', 'val.txt'])
|
|
if not os.path.exists(path_to_processed_val_data):
|
|
processed_val_smiles = preprocess_dataset(atom_types, bond_types, val_smiles, None)
|
|
write_smiles_to_file(path_to_processed_val_data, processed_val_smiles)
|
|
|
|
class MoleculeDataset(object):
|
|
"""Initialize and split the dataset.
|
|
|
|
Parameters
|
|
----------
|
|
dataset : str
|
|
Dataset name
|
|
order : None or str
|
|
Order to extract a decision sequence for generating a molecule. Default to be None.
|
|
modes : None or list
|
|
List of subsets to use, which can contain 'train', 'val', corresponding to
|
|
training and validation. Default to be None.
|
|
subset_id : int
|
|
With multiprocess training, we partition the training set into multiple subsets and
|
|
each process will use one subset only. This subset_id corresponds to subprocess id.
|
|
n_subsets : int
|
|
With multiprocess training, this corresponds to the number of total subprocesses.
|
|
"""
|
|
def __init__(self, dataset, order=None, modes=None, subset_id=0, n_subsets=1):
|
|
super(MoleculeDataset, self).__init__()
|
|
|
|
if modes is None:
|
|
modes = []
|
|
else:
|
|
assert order is not None, 'An order should be specified for extracting ' \
|
|
'decision sequences.'
|
|
|
|
assert order in ['random', 'canonical', None], \
|
|
"Unexpected order option to get sequences of graph generation decisions"
|
|
assert len(set(modes) - {'train', 'val'}) == 0, \
|
|
"modes should be a list, representing a subset of ['train', 'val']"
|
|
|
|
self.dataset = dataset
|
|
self.order = order
|
|
self.modes = modes
|
|
self.subset_id = subset_id
|
|
self.n_subsets = n_subsets
|
|
self._setup()
|
|
|
|
def collate(self, samples):
|
|
"""PyTorch's approach to batch multiple samples.
|
|
|
|
For auto-regressive generative models, we process one sample at a time.
|
|
|
|
Parameters
|
|
----------
|
|
samples : list
|
|
A list of length 1 that consists of decision sequence to generate a molecule.
|
|
|
|
Returns
|
|
-------
|
|
list
|
|
List of 2-tuples, a decision sequence to generate a molecule
|
|
"""
|
|
assert len(samples) == 1
|
|
return samples[0]
|
|
|
|
def _create_a_subset(self, smiles):
|
|
"""Create a dataset from a subset of smiles.
|
|
|
|
Parameters
|
|
----------
|
|
smiles : list of str
|
|
List of molecules in SMILES format
|
|
"""
|
|
# We evenly divide the smiles into multiple susbets with multiprocess
|
|
subset_size = len(smiles) // self.n_subsets
|
|
return Subset(smiles[self.subset_id * subset_size: (self.subset_id + 1) * subset_size],
|
|
self.order, self.env)
|
|
|
|
def _setup(self):
|
|
"""
|
|
1. Instantiate an MDP environment for molecule generation
|
|
2. Download the dataset, which is a file of SMILES
|
|
3. Create subsets for training and validation
|
|
"""
|
|
if self.dataset == 'ChEMBL':
|
|
# For new datasets, get_atom_and_bond_types can be used to
|
|
# identify the atom and bond types in them.
|
|
self.atom_types = ['O', 'Cl', 'C', 'S', 'F', 'Br', 'N']
|
|
self.bond_types = [Chem.rdchem.BondType.SINGLE,
|
|
Chem.rdchem.BondType.DOUBLE,
|
|
Chem.rdchem.BondType.TRIPLE]
|
|
|
|
elif self.dataset == 'ZINC':
|
|
self.atom_types = ['Br', 'S', 'C', 'P', 'N', 'O', 'F', 'Cl', 'I']
|
|
self.bond_types = [Chem.rdchem.BondType.SINGLE,
|
|
Chem.rdchem.BondType.DOUBLE,
|
|
Chem.rdchem.BondType.TRIPLE]
|
|
|
|
else:
|
|
path_to_atom_and_bond_types = '_'.join([self.dataset, 'atom_and_bond_types.pkl'])
|
|
with open(path_to_atom_and_bond_types, 'rb') as f:
|
|
type_info = pickle.load(f)
|
|
self.atom_types = type_info['atom_types']
|
|
self.bond_types = type_info['bond_types']
|
|
self.env = MoleculeEnv(self.atom_types, self.bond_types)
|
|
|
|
dataset_prefix = self._dataset_prefix()
|
|
|
|
if 'train' in self.modes:
|
|
fname = '_'.join([dataset_prefix, 'train.txt'])
|
|
download_data(self.dataset, fname)
|
|
smiles = load_smiles_from_file(fname)
|
|
self.train_set = self._create_a_subset(smiles)
|
|
|
|
if 'val' in self.modes:
|
|
fname = '_'.join([dataset_prefix, 'val.txt'])
|
|
download_data(self.dataset, fname)
|
|
smiles = load_smiles_from_file(fname)
|
|
# We evenly divide the smiles into multiple susbets with multiprocess
|
|
self.val_set = self._create_a_subset(smiles)
|
|
|
|
def _dataset_prefix(self):
|
|
"""Get the prefix for the data files of supported datasets.
|
|
|
|
Returns
|
|
-------
|
|
str
|
|
Prefix for dataset file name
|
|
"""
|
|
return '_'.join([self.dataset, 'DGMG'])
|
|
|
|
class Subset(Dataset):
|
|
"""A set of molecules which can be used for training, validation, test.
|
|
|
|
Parameters
|
|
----------
|
|
smiles : list
|
|
List of SMILES for the dataset
|
|
order : str
|
|
Specifies how decision sequences for molecule generation
|
|
are obtained, can be either "random" or "canonical"
|
|
env : MoleculeEnv object
|
|
MDP environment for generating molecules
|
|
"""
|
|
def __init__(self, smiles, order, env):
|
|
super(Subset, self).__init__()
|
|
self.smiles = smiles
|
|
self.order = order
|
|
self.env = env
|
|
self._setup()
|
|
|
|
def _setup(self):
|
|
"""Convert SMILES into rdkit molecule objects.
|
|
|
|
Decision sequences are extracted if we use a fixed order.
|
|
"""
|
|
smiles_ = []
|
|
mols = []
|
|
for s in self.smiles:
|
|
m = smiles_to_standard_mol(s)
|
|
if m is None:
|
|
continue
|
|
smiles_.append(s)
|
|
mols.append(m)
|
|
self.smiles = smiles_
|
|
self.mols = mols
|
|
|
|
if self.order is 'random':
|
|
return
|
|
|
|
self.decisions = []
|
|
for m in self.mols:
|
|
self.decisions.append(
|
|
self.env.get_decision_sequence(m, list(range(m.GetNumAtoms())))
|
|
)
|
|
|
|
def __len__(self):
|
|
"""Get number of molecules in the dataset."""
|
|
return len(self.mols)
|
|
|
|
def __getitem__(self, item):
|
|
"""Get the decision sequence for generating the molecule indexed by item."""
|
|
if self.order == 'canonical':
|
|
return self.decisions[item]
|
|
else:
|
|
m = self.mols[item]
|
|
nodes = list(range(m.GetNumAtoms()))
|
|
random.shuffle(nodes)
|
|
return self.env.get_decision_sequence(m, nodes)
|
|
|
|
########################################################################################################################
|
|
# progress tracking #
|
|
########################################################################################################################
|
|
|
|
class Printer(object):
|
|
def __init__(self, num_epochs, dataset_size, batch_size, writer=None):
|
|
"""Wrapper to track the learning progress.
|
|
|
|
Parameters
|
|
----------
|
|
num_epochs : int
|
|
Number of epochs for training
|
|
dataset_size : int
|
|
batch_size : int
|
|
writer : None or SummaryWriter
|
|
If not None, tensorboard will be used to visualize learning curves.
|
|
"""
|
|
super(Printer, self).__init__()
|
|
|
|
self.num_epochs = num_epochs
|
|
self.batch_size = batch_size
|
|
self.num_batches = math.ceil(dataset_size / batch_size)
|
|
self.count = 0
|
|
self.batch_count = 0
|
|
self.writer = writer
|
|
self._reset()
|
|
|
|
def _reset(self):
|
|
"""Reset when an epoch is completed."""
|
|
self.batch_loss = 0
|
|
self.batch_prob = 0
|
|
|
|
def _get_current_batch(self):
|
|
"""Get current batch index."""
|
|
remainer = self.batch_count % self.num_batches
|
|
if (remainer == 0):
|
|
return self.num_batches
|
|
else:
|
|
return remainer
|
|
|
|
def update(self, epoch, loss, prob):
|
|
"""Update learning progress.
|
|
|
|
Parameters
|
|
----------
|
|
epoch : int
|
|
loss : float
|
|
prob : float
|
|
"""
|
|
self.count += 1
|
|
self.batch_loss += loss
|
|
self.batch_prob += prob
|
|
|
|
if self.count % self.batch_size == 0:
|
|
self.batch_count += 1
|
|
if self.writer is not None:
|
|
self.writer.add_scalar('train_log_prob', self.batch_loss, self.batch_count)
|
|
self.writer.add_scalar('train_prob', self.batch_prob, self.batch_count)
|
|
|
|
print('epoch {:d}/{:d}, batch {:d}/{:d}, loss {:.4f}, prob {:.4f}'.format(
|
|
epoch, self.num_epochs, self._get_current_batch(),
|
|
self.num_batches, self.batch_loss, self.batch_prob))
|
|
self._reset()
|
|
|
|
########################################################################################################################
|
|
# eval #
|
|
########################################################################################################################
|
|
|
|
def summarize_a_molecule(smile, checklist=None):
|
|
"""Get information about a molecule.
|
|
|
|
Parameters
|
|
----------
|
|
smile : str
|
|
Molecule in SMILES format
|
|
checklist : dict
|
|
Things to learn about the molecule
|
|
"""
|
|
if checklist is None:
|
|
checklist = {
|
|
'HBA': Chem.rdMolDescriptors.CalcNumHBA,
|
|
'HBD': Chem.rdMolDescriptors.CalcNumHBD,
|
|
'logP': MolLogP,
|
|
'SA': calculateScore,
|
|
'TPSA': Chem.rdMolDescriptors.CalcTPSA,
|
|
'QED': qed,
|
|
'NumAtoms': lambda mol: mol.GetNumAtoms(),
|
|
'NumBonds': lambda mol: mol.GetNumBonds()
|
|
}
|
|
|
|
summary = dict()
|
|
mol = Chem.MolFromSmiles(smile)
|
|
if mol is None:
|
|
summary.update({
|
|
'smile': smile,
|
|
'valid': False
|
|
})
|
|
for k in checklist.keys():
|
|
summary[k] = None
|
|
else:
|
|
mol = standardize_mol(mol)
|
|
summary.update({
|
|
'smile': Chem.MolToSmiles(mol),
|
|
'valid': True
|
|
})
|
|
Chem.SanitizeMol(mol)
|
|
for k, f in checklist.items():
|
|
summary[k] = f(mol)
|
|
|
|
return summary
|
|
|
|
def summarize_molecules(smiles, num_processes):
|
|
"""Summarize molecules with multiprocess.
|
|
|
|
Parameters
|
|
----------
|
|
smiles : list of str
|
|
List of molecules in SMILES for summarization
|
|
num_processes : int
|
|
Number of processes to use for summarization
|
|
|
|
Returns
|
|
-------
|
|
summary_for_valid : dict
|
|
Summary of all valid molecules, where
|
|
summary_for_valid[k] gives the values of all
|
|
valid molecules on item k.
|
|
"""
|
|
with Pool(processes=num_processes) as pool:
|
|
result = pool.map(summarize_a_molecule, smiles)
|
|
|
|
items = list(result[0].keys())
|
|
items.remove('valid')
|
|
|
|
summary_for_valid = defaultdict(list)
|
|
for summary in result:
|
|
if summary['valid']:
|
|
for k in items:
|
|
summary_for_valid[k].append(summary[k])
|
|
return summary_for_valid
|
|
|
|
def get_unique_smiles(smiles):
|
|
"""Given a list of smiles, return a list consisting of unique elements in it.
|
|
|
|
Parameters
|
|
----------
|
|
smiles : list of str
|
|
Molecules in SMILES
|
|
|
|
Returns
|
|
-------
|
|
list of str
|
|
Sublist where each SMIES occurs exactly once
|
|
"""
|
|
unique_set = set()
|
|
for mol_s in smiles:
|
|
if mol_s not in unique_set:
|
|
unique_set.add(mol_s)
|
|
|
|
return list(unique_set)
|
|
|
|
def get_novel_smiles(new_unique_smiles, reference_unique_smiles):
|
|
"""Get novel smiles which do not appear in the reference set.
|
|
|
|
Parameters
|
|
----------
|
|
new_unique_smiles : list of str
|
|
List of SMILES from which we want to identify novel ones
|
|
reference_unique_smiles : list of str
|
|
List of reference SMILES that we already have
|
|
"""
|
|
return set(new_unique_smiles).difference(set(reference_unique_smiles))
|