policies.py

import os
import random
import numpy as np
import tensorflow as tf
from rdkit import Chem, DataStructs
from rdkit.Chem import AllChem
from tqdm import tqdm, trange
from collections import defaultdict


def fps_to_arr(fps):
    """Faster conversion to ndarray"""
    arrs = []
    for fp in fps:
        onbits = list(fp.GetOnBits())
        arr = np.zeros(fp.GetNumBits())
        arr[onbits] = 1
        arrs.append(arr)
    arrs = np.array(arrs)
    return arrs


def highway_layer(x, activation, carry_bias=-1.0):
    size = x.shape[-1].value
    W_T = tf.Variable(tf.truncated_normal((size, size), stddev=0.1), name='weight_transform')
    b_T = tf.Variable(tf.constant(carry_bias, shape=(size,)), name='bias_transform')
    W = tf.Variable(tf.truncated_normal((size, size), stddev=0.1), name='weight')
    b = tf.Variable(tf.constant(0.1, shape=(size,)), name='bias')
    T = tf.sigmoid(tf.matmul(x, W_T) + b_T, name='transform_gate')
    H = activation(tf.matmul(x, W) + b, name='activation')
    C = tf.subtract(1.0, T, name='carry_gate')
    return tf.add(tf.multiply(H, T), tf.multiply(x, C))


def fingerprint_mols(mols, fp_dim):
    fps = []
    for mol in mols:
        mol = Chem.MolFromSmiles(mol)

        # Necessary for fingerprinting
        # Chem.GetSymmSSSR(mol)

        # "When comparing the ECFP/FCFP fingerprints and
        # the Morgan fingerprints generated by the RDKit,
        # remember that the 4 in ECFP4 corresponds to the
        # diameter of the atom environments considered,
        # while the Morgan fingerprints take a radius parameter.
        # So the examples above, with radius=2, are roughly
        # equivalent to ECFP4 and FCFP4."
        # <http://www.rdkit.org/docs/GettingStartedInPython.html>
        fp = AllChem.GetMorganFingerprintAsBitVect(mol, 2, nBits=int(fp_dim))
        # fold_factor = fp.GetNumBits()//fp_dim
        # fp = DataStructs.FoldFingerprint(fp, fold_factor)
        fps.append(fp)
    return fps


def fingerprint_reactions(reactions, fp_dim):
    fps = []
    for r in reactions:
        rxn = AllChem.ReactionFromSmarts(r)
        # fp = AllChem.CreateDifferenceFingerprintForReaction(rxn)
        fp = AllChem.CreateStructuralFingerprintForReaction(rxn)
        fold_factor = fp.GetNumBits()//fp_dim
        fp = DataStructs.FoldFingerprint(fp, fold_factor)
        fps.append(fp)
    return fps


def train(sess, net, X, y, batch_size=16, epochs=10):
    losses = []
    accuracy = []
    it = trange(epochs)
    n_steps = int(np.ceil(len(X)/batch_size))
    for e in it:
        # Shuffle
        # p = np.random.permutation(len(X))
        # X, y = X[p], y[p]
        xy = list(zip(X, y))
        random.shuffle(xy)
        X, y = zip(*xy)

        # Iterate batches
        for i in tqdm(range(n_steps)):
            l = i*batch_size
            u = l + batch_size
            X_batch, y_batch = X[l:u], y[l:u]
            X_batch = net.preprocess(X_batch)
            _, err, acc = sess.run(
                [net.train_op, net.loss_op, net.acc_op],
                feed_dict={
                    net.keep_prob: 0.4,
                    net.X: X_batch,
                    net.y: y_batch
                }
            )
            losses.append(err)
            accuracy.append(acc)
            it.set_postfix(
                loss=np.mean(losses[-10:]) if losses else None,
                acc=np.mean(accuracy[-10:]) if accuracy else None)
    return losses


class RolloutPolicyNet:
    def __init__(self, n_rules, fp_dim=8912, k=10):
        self.fp_dim = fp_dim
        self.n_rules = n_rules
        self.X = tf.placeholder(tf.float32, shape=(None, fp_dim))
        self.y = tf.placeholder(tf.int64, shape=(None,))
        self.keep_prob = tf.placeholder(tf.float32, shape=())

        inp = tf.math.log(self.X+1)
        net = tf.layers.dense(inp, 512, activation=tf.nn.elu)
        net = tf.nn.dropout(net, keep_prob=self.keep_prob)
        net = tf.layers.dense(net, n_rules, activation=None)
        self.pred_op = tf.nn.softmax(net)
        # self.pred = tf.nn.top_k(pred, k=k)
        self.loss_op = tf.losses.sparse_softmax_cross_entropy(self.y, net)
        self.train_op = tf.train.AdamOptimizer(learning_rate=1e-4).minimize(self.loss_op)

        correct_pred = tf.equal(tf.argmax(self.pred_op, 1), self.y)
        self.acc_op = tf.reduce_mean(tf.cast(correct_pred, tf.float32))

    def preprocess(self, X):
        # Compute fingerprints
        return fps_to_arr(fingerprint_mols(X, self.fp_dim))


class ExpansionPolicyNet:
    def __init__(self, idx, n_rules, fp_dim=1e6):
        self.idx = idx
        self.fp_dim = fp_dim
        self.n_rules = n_rules

        # self.X = tf.placeholder(tf.float32, shape=(None, self.idx.shape[-1]))
        self.X = tf.placeholder(tf.float32, shape=(None, fp_dim))
        self.y = tf.placeholder(tf.int64, shape=(None,))
        self.keep_prob = tf.placeholder(tf.float32, shape=())

        # inp = self.X
        inp = tf.math.log(self.X+1)
        net = tf.layers.dense(inp, 512, activation=tf.nn.elu)
        net = tf.nn.dropout(net, keep_prob=self.keep_prob)
        for _ in range(5):
            net = highway_layer(net, activation=tf.nn.elu)
            net = tf.nn.dropout(net, keep_prob=self.keep_prob)

        net = tf.layers.dense(net, n_rules, activation=None)
        self.pred_op = tf.nn.softmax(net)
        self.loss_op = tf.losses.sparse_softmax_cross_entropy(self.y, net)
        self.train_op = tf.train.AdamOptimizer(learning_rate=1e-4).minimize(self.loss_op)

        correct_pred = tf.equal(tf.argmax(self.pred_op, 1), self.y)
        self.acc_op = tf.reduce_mean(tf.cast(correct_pred, tf.float32))

    def preprocess(self, X):
        # Compute fingerprints
        X = fingerprint_mols(X, self.fp_dim)
        X = fps_to_arr(X)

        # Apply variance threshold
        # return np.log(X[:,self.idx] + 1)
        return X


class InScopeFilterNet:
    def __init__(self, product_fp_dim=16384, reaction_fp_dim=2048):
        self.prod_fp_dim = product_fp_dim
        self.react_fp_dim = reaction_fp_dim
        self.X = tf.placeholder(tf.float32, shape=(None, product_fp_dim+reaction_fp_dim))
        self.X_prod = self.X[:,:product_fp_dim]
        self.X_react = self.X[:,product_fp_dim:]
        self.y = tf.placeholder(tf.int32, shape=(None,))
        self.keep_prob = tf.placeholder(tf.float32, shape=())

        # Product branch
        prod_inp = tf.math.log(self.X_prod+1)
        prod_net = tf.layers.dense(prod_inp, 1024, activation=tf.nn.elu)
        prod_net = tf.nn.dropout(prod_net, keep_prob=self.keep_prob)
        for _ in range(5):
            prod_net = highway_layer(prod_net, activation=tf.nn.elu)

        # Reaction branch
        react_net = tf.layers.dense(self.X_react, 1024, activation=tf.nn.elu)

        # Cosine similarity
        prod_norm = tf.nn.l2_normalize(prod_net, axis=-1)
        react_norm = tf.nn.l2_normalize(react_net, axis=-1)
        cosine_sim = tf.reduce_sum(tf.multiply(prod_norm, react_norm), axis=-1)

        # Paper's architecture passes the similarity through a sigmoid function
        # but that seems redundant?
        self.pred = tf.nn.sigmoid(cosine_sim)

        self.loss_op = tf.losses.log_loss(self.y, self.pred)
        self.train_op = tf.train.AdamOptimizer(learning_rate=0.001).minimize(self.loss_op)

        correct_pred = tf.equal(self.y, tf.cast(tf.round(self.pred), tf.int32))
        self.acc_op = tf.reduce_mean(tf.cast(correct_pred, tf.float32))

    def preprocess(self, X):
        # Compute fingerprints
        prod_mols, react_mols = zip(*X)
        prod_fps = fingerprint_mols(prod_mols, self.prod_fp_dim)
        react_fps = fingerprint_reactions(react_mols, self.react_fp_dim)
        return np.hstack([prod_fps, react_fps])


if __name__ == '__main__':
    print('Loading data...')
    prod_to_rules = defaultdict(set)
    with open('data/templates.dat', 'r') as f:
        for l in tqdm(f, desc='products'):
            rule, prod = l.strip().split('\t')
            prod_to_rules[prod].add(rule)

    rollout_rules = {}
    with open('data/rollout.dat', 'r') as f:
        for i, l in tqdm(enumerate(f), desc='rollout'):
            rule = l.strip()
            rollout_rules[rule] = i

    expansion_rules = {}
    with open('data/expansion.dat', 'r') as f:
        for i, l in tqdm(enumerate(f), desc='expansion'):
            rule = l.strip()
            expansion_rules[rule] = i

    prods = list(prod_to_rules.keys())
    exp_fp_dim = 1e4

    def chunker(seq, size):
        return (seq[pos:pos + size] for pos in range(0, len(seq), size))

    def reducefn(a, b):
        n_a, mean_a, var_a = a
        n_b, mean_b, var_b = b
        n_ab = n_a + n_b
        mean_ab = ((mean_a * n_a) + (mean_b * n_b)) / n_ab
        var_ab = (((n_a * var_a) + (n_b * var_b)) / n_ab) + ((n_a * n_b) * ((mean_b - mean_a) / n_ab)**2)
        return n_ab, mean_ab, var_ab

    def mapfn(chunk):
        chunk = fingerprint_mols(chunk, exp_fp_dim)
        arrs = np.log(fps_to_arr(chunk) + 1)
        return len(arrs), np.mean(arrs, axis=0), np.var(arrs, axis=0)

    save_path = 'model'
    ckpt_path = os.path.join(save_path, 'model.ckpt')
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    # Variance threshold
    # No specific threshold is mentioned in the paper,
    # just that it's used to "remove rare features"
    # idx = np.load(os.path.join(save_path, 'expansion.idx.npy'))
    from multiprocessing import Pool
    from functools import reduce
    chunk_size = 100
    prods = []
    for prod, rules in prod_to_rules.items():
        if any(r in expansion_rules for r in rules):
            prods.append(prod)
    # prods = random.sample(prods, 200000)
    with Pool() as p:
        _, _, var = reduce(reducefn, tqdm(p.imap(mapfn, chunker(prods, chunk_size)), total=len(prods)//chunk_size))
    idx = np.where(var > 0)[0]
    # np.save(os.path.join(save_path, 'expansion.idx'), idx, allow_pickle=False)

    rollout = RolloutPolicyNet(n_rules=len(rollout_rules))
    expansion = ExpansionPolicyNet(idx, n_rules=len(expansion_rules), fp_dim=exp_fp_dim)
    filter = InScopeFilterNet()

    sess = tf.Session()
    init = tf.global_variables_initializer()
    sess.run(init)

    saver = tf.train.Saver()
    # saver.restore(sess, ckpt_path)

    # Rollout training
    # print('Rollout training...')
    # X, y = [], []
    # for prod, rules in tqdm(prod_to_rules.items(), desc='data prep'):
    #     rules = [r for r in rules if r in rollout_rules]
    #     if not rules: continue

    #     # Ideally trained as multilabel,
    #     # but multiclass, single label is easier atm
    #     for r in rules:
    #         id = rollout_rules[r]
    #         y.append(id)
    #         X.append(prod)
    # print('Training size:', len(X))
    # train(sess, rollout, X, y, batch_size=256, epochs=100)
    # saver.save(sess, ckpt_path)

    # Check
    # X = rollout.preprocess(X[20:30])
    # y_pred = sess.run(rollout.pred_op, feed_dict={
    #     rollout.keep_prob: 1.,
    #     rollout.X: X
    # })
    # print(list(np.argmax(y_pred, 1)))
    # print(y[20:30])


    print('Expansion training...')
    X, y = [], []
    for prod, rules in tqdm(prod_to_rules.items(), desc='data prep'):
        rules = [r for r in rules if r in expansion_rules]
        if not rules: continue

        # Ideally trained as multilabel,
        # but multiclass, single label is easier atm
        for r in rules:
            id = expansion_rules[r]
            y.append(id)
            X.append(prod)
    print('Training size:', len(X))
    train(sess, expansion, X, y, batch_size=1024, epochs=2000)
    saver.save(sess, ckpt_path)

    # print('In-Scope Filter training...')
    # X, y = [], []
    # exists = set()
    # for prod, rules in tqdm(prod_to_rules.items(), desc='data prep'):
    #     rules = [r for r in rules if r in expansion_rules]
    #     if not rules: continue

    #     for r in rules:
    #         y.append(1.)
    #         X.append((prod, r))
    #         exists.add('{}_{}'.format(prod, r))

    # # Generate negative examples
    # target_size = len(X) * 2
    # pbar = tqdm(total=target_size//2, desc='data prep (negative)')
    # prods = list(prod_to_rules.keys())
    # exprules = list(expansion_rules.keys())
    # while len(X) < target_size:
    #     prod = random.choice(prods)
    #     rule = random.choice(exprules)

    #     key = '{}_{}'.format(prod, r)
    #     if key in exists:
    #         continue
    #     else:
    #         y.append(0.)
    #         X.append((prod, rule))
    #         pbar.update(1)
    # pbar.close()
    # print('Training size:', len(X))
    # train(sess, filter, X, y, batch_size=512, epochs=3)
    # saver.save(sess, ckpt_path)

    # for v in tf.trainable_variables():
    #     print(v)