<commit_before>import itertools as it
import autograd.numpy as np
from features import N_atom_features, N_bond_features
from util import memoize, WeightsParser
from mol_graph import graph_from_smiles_tuple
from autograd import grad
from build_vanilla_net import build_vanilla_net

from autograd.scipy.misc import logsumexp

def fast_array_from_list(xs):
    return np.concatenate([np.expand_dims(x, axis=0) for x in xs], axis=0)

def apply_and_stack(idxs, features, op):
    return fast_array_from_list([op(features[idx_list, :]) for idx_list in idxs])

def softened_max(X, axis=0):
    """Takes the row-wise max, but gently."""
    exp_X = np.exp(X)
    return np.sum(exp_X * X, axis) / np.sum(exp_X, axis)

def weighted_softened_max(X, axis=0):
    # X is (permutations x atoms x features).
    # The first feature softly scales how much each
    # permutation is used to weight the final activations.
    scale_feature = X[:, :, 0:1]
    # logsumexp with keepdims requires scipy >= 0.15.
    return np.sum(X * np.exp(scale_feature - logsumexp(scale_feature, axis, keepdims=True)), axis)

def matmult_neighbors(mol_nodes, other_ntypes, feature_sets, get_weights, permutations):
    def neighbor_list(degree, other_ntype):
        return mol_nodes[(other_ntype + '_neighbors', degree)]
    activations_by_degree = []
    for degree in [1, 2, 3, 4]:
        # TODO: make this (plus the stacking) into an autograd primitive
        neighbor_features = [features[neighbor_list(degree, other_ntype)]
                             for other_ntype, features in zip(other_ntypes, feature_sets)]
        if any([len(feat) > 0 for feat in neighbor_features]):
            # dims of stacked_neighbors are [atoms, neighbors, features]
            stacked_neighbors = np.concatenate(neighbor_features, axis=2)
            if permutations:
                weightses = [get_weights(degree)[n, :, :] for n in range(degree)]
                neighbors = [stacked_neighbors[:, d, :] for d in range(degree)]
                products = [[np.dot(n, w) for w in weightses] for n in neighbors]
                # (Atoms x neighbors1 x features), (neighbors2 x features x hiddens)
                # -> (atoms x neighbors x hiddens)
                #products = np.einsum("anf,mfh->nmah", stacked_neighbors, get_weights(degree))
                candidates = [sum([products[i][j] for i, j in enumerate(p)])
                              for p in it.permutations(range(degree))]
                # dims candidates is (permutations x atoms x features)
                activations_by_degree.append(weighted_softened_max(fast_array_from_list(candidates)))
            else:
                # (Atoms x neighbors x features), (features x hiddens) -> (atoms x hiddens)
                # activations = np.einsum("anf,fh->ah", stacked_neighbors, get_weights(degree))
                activations = np.dot(np.sum(stacked_neighbors, axis=1), get_weights(degree))
                activations_by_degree.append(activations)
    # This is brittle! Relies on atoms being sorted by degree in the first place,
    # in Node.graph_from_smiles_tuple()
    return np.concatenate(activations_by_degree, axis=0)

def weights_name(layer, degree, neighbor=None):
    if neighbor:
        neighborstr = " neighbor " + str(neighbor)
    else:
        neighborstr = ""
    return "layer " + str(layer) + " degree " + str(degree) + neighborstr + " filter"

def build_convnet(bond_vec_dim=1, num_hidden_features=[20, 50, 50],
                        permutations=False, l2_penalty=0.0, pool_funcs=['softened_max']):
    """Sets up functions to compute convnets over all molecules in a minibatch together.
       The number of hidden layers is the length of num_hidden_features - 1."""
    parser = WeightsParser()
    parser.add_weights('atom2vec', (N_atom_features, num_hidden_features[0]))
    parser.add_weights('bond2vec', (N_bond_features, bond_vec_dim))

    in_and_out_sizes = zip(num_hidden_features[:-1], num_hidden_features[1:])
    for layer, (N_prev, N_cur) in enumerate(in_and_out_sizes):
        parser.add_weights("layer " + str(layer) + " biases", (1, N_cur))
        parser.add_weights("layer " + str(layer) + " self filter", (N_prev, N_cur))
        base_shape = (N_prev + bond_vec_dim, N_cur)
        for degree in [1, 2, 3, 4]:
            if permutations:
                shape = (degree,) + base_shape
            else:
                shape = base_shape
            parser.add_weights(weights_name(layer, degree), shape)

    parser.add_weights('output bias', (1, ))
    parser.add_weights('output weights', (num_hidden_features[-1] * 2, ))

    def output_layer_fun(weights, smiles):
        """Computes layer-wise convolution, and returns a fixed-size output."""
        mol_nodes = arrayrep_from_smiles(tuple(smiles))

        atom_features = np.dot(mol_nodes['atom_features'], parser.get(weights, 'atom2vec'))
        bond_features = np.dot(mol_nodes['bond_features'], parser.get(weights, 'bond2vec'))

        for layer in xrange(len(num_hidden_features) - 1):
            def get_weights_func(degree):#, neighbor=None):
                return parser.get(weights, weights_name(layer, degree))#, neighbor))
            layer_bias = parser.get(weights, "layer " + str(layer) + " biases")
            layer_self_weights = parser.get(weights, "layer " + str(layer) + " self filter")
            self_activations = np.dot(atom_features, layer_self_weights)
            neighbour_activations = matmult_neighbors(mol_nodes, ('atom', 'bond'),
                (atom_features, bond_features), get_weights_func, permutations)
            atom_features = np.tanh(layer_bias + self_activations + neighbour_activations)

        # Include both a softened-max and a sum node to pool all atom features together.
        atom_neighbor_idxs = mol_nodes['mol_atom_neighbors']
        pooled_features = []
        if 'softened_max' in pool_funcs:
            pooled_features.append(apply_and_stack(atom_neighbor_idxs, atom_features,
                                                   softened_max))
        if 'sum' in pool_funcs:
            pooled_features.append(apply_and_stack(atom_neighbor_idxs, atom_features,
                                                   lambda x: np.sum(x, axis=0)))
        return np.concatenate(pooled_features, axis=1)

    def prediction_fun(weights, smiles):
        output_weights = parser.get(weights, 'output weights')
        output_bias = parser.get(weights, 'output bias')
        hidden_units = output_layer_fun(weights, smiles)
        return np.dot(hidden_units, output_weights) + output_bias

    def loss_fun(weights, smiles, targets):
        preds = prediction_fun(weights, smiles)
        acc_loss = np.sum((preds - targets)**2)
        l2reg = l2_penalty * np.sum(weights**2)
        return acc_loss + l2reg<commit_msg>Added mean pooling<commit_after>import itertools as it
import autograd.numpy as np
from features import N_atom_features, N_bond_features
from util import memoize, WeightsParser
from mol_graph import graph_from_smiles_tuple
from autograd import grad
from build_vanilla_net import build_vanilla_net

from autograd.scipy.misc import logsumexp

def fast_array_from_list(xs):
    return np.concatenate([np.expand_dims(x, axis=0) for x in xs], axis=0)

def apply_and_stack(idxs, features, op):
    return fast_array_from_list([op(features[idx_list, :]) for idx_list in idxs])

def softened_max(X, axis=0):
    """Takes the row-wise max, but gently."""
    exp_X = np.exp(X)
    return np.sum(exp_X * X, axis) / np.sum(exp_X, axis)

def weighted_softened_max(X, axis=0):
    # X is (permutations x atoms x features).
    # The first feature softly scales how much each
    # permutation is used to weight the final activations.
    scale_feature = X[:, :, 0:1]
    # logsumexp with keepdims requires scipy >= 0.15.
    return np.sum(X * np.exp(scale_feature - logsumexp(scale_feature, axis, keepdims=True)), axis)

def matmult_neighbors(mol_nodes, other_ntypes, feature_sets, get_weights, permutations):
    def neighbor_list(degree, other_ntype):
        return mol_nodes[(other_ntype + '_neighbors', degree)]
    activations_by_degree = []
    for degree in [1, 2, 3, 4]:
        # TODO: make this (plus the stacking) into an autograd primitive
        neighbor_features = [features[neighbor_list(degree, other_ntype)]
                             for other_ntype, features in zip(other_ntypes, feature_sets)]
        if any([len(feat) > 0 for feat in neighbor_features]):
            # dims of stacked_neighbors are [atoms, neighbors, features]
            stacked_neighbors = np.concatenate(neighbor_features, axis=2)
            if permutations:
                weightses = [get_weights(degree)[n, :, :] for n in range(degree)]
                neighbors = [stacked_neighbors[:, d, :] for d in range(degree)]
                products = [[np.dot(n, w) for w in weightses] for n in neighbors]
                # (Atoms x neighbors1 x features), (neighbors2 x features x hiddens)
                # -> (atoms x neighbors x hiddens)
                #products = np.einsum("anf,mfh->nmah", stacked_neighbors, get_weights(degree))
                candidates = [sum([products[i][j] for i, j in enumerate(p)])
                              for p in it.permutations(range(degree))]
                # dims candidates is (permutations x atoms x features)
                activations_by_degree.append(weighted_softened_max(fast_array_from_list(candidates)))
            else:
                # (Atoms x neighbors x features), (features x hiddens) -> (atoms x hiddens)
                # activations = np.einsum("anf,fh->ah", stacked_neighbors, get_weights(degree))
                activations = np.dot(np.sum(stacked_neighbors, axis=1), get_weights(degree))
                activations_by_degree.append(activations)
    # This is brittle! Relies on atoms being sorted by degree in the first place,
    # in Node.graph_from_smiles_tuple()
    return np.concatenate(activations_by_degree, axis=0)

def weights_name(layer, degree, neighbor=None):
    if neighbor:
        neighborstr = " neighbor " + str(neighbor)
    else:
        neighborstr = ""
    return "layer " + str(layer) + " degree " + str(degree) + neighborstr + " filter"

def build_convnet(bond_vec_dim=1, num_hidden_features=[20, 50, 50],
                        permutations=False, l2_penalty=0.0, pool_funcs=['softened_max']):
    """Sets up functions to compute convnets over all molecules in a minibatch together.
       The number of hidden layers is the length of num_hidden_features - 1."""
    parser = WeightsParser()
    parser.add_weights('atom2vec', (N_atom_features, num_hidden_features[0]))
    parser.add_weights('bond2vec', (N_bond_features, bond_vec_dim))

    in_and_out_sizes = zip(num_hidden_features[:-1], num_hidden_features[1:])
    for layer, (N_prev, N_cur) in enumerate(in_and_out_sizes):
        parser.add_weights("layer " + str(layer) + " biases", (1, N_cur))
        parser.add_weights("layer " + str(layer) + " self filter", (N_prev, N_cur))
        base_shape = (N_prev + bond_vec_dim, N_cur)
        for degree in [1, 2, 3, 4]:
            if permutations:
                shape = (degree,) + base_shape
            else:
                shape = base_shape
            parser.add_weights(weights_name(layer, degree), shape)

    parser.add_weights('output bias', (1, ))
    parser.add_weights('output weights', (num_hidden_features[-1] * 2, ))

    def output_layer_fun(weights, smiles):
        """Computes layer-wise convolution, and returns a fixed-size output."""
        mol_nodes = arrayrep_from_smiles(tuple(smiles))

        atom_features = np.dot(mol_nodes['atom_features'], parser.get(weights, 'atom2vec'))
        bond_features = np.dot(mol_nodes['bond_features'], parser.get(weights, 'bond2vec'))

        for layer in xrange(len(num_hidden_features) - 1):
            def get_weights_func(degree):#, neighbor=None):
                return parser.get(weights, weights_name(layer, degree))#, neighbor))
            layer_bias = parser.get(weights, "layer " + str(layer) + " biases")
            layer_self_weights = parser.get(weights, "layer " + str(layer) + " self filter")
            self_activations = np.dot(atom_features, layer_self_weights)
            neighbour_activations = matmult_neighbors(mol_nodes, ('atom', 'bond'),
                (atom_features, bond_features), get_weights_func, permutations)
            atom_features = np.tanh(layer_bias + self_activations + neighbour_activations)

        # Include both a softened-max and a sum node to pool all atom features together.
        atom_neighbor_idxs = mol_nodes['mol_atom_neighbors']
        pooled_features = []
        if 'softened_max' in pool_funcs:
            pooled_features.append(apply_and_stack(atom_neighbor_idxs, atom_features,
                                                   softened_max))
        if 'mean' in pool_funcs:
            pooled_features.append(apply_and_stack(atom_neighbor_idxs, atom_features,
                                                   lambda x: np.mean(x, axis=0)))
        if 'sum' in pool_funcs:
            pooled_features.append(apply_and_stack(atom_neighbor_idxs, atom_features,
                                                   lambda x: np.sum(x, axis=0)))
        return np.concatenate(pooled_features, axis=1)

    def prediction_fun(weights, smiles):
        output_weights = parser.get(weights, 'output weights')
        output_bias = parser.get(weights, 'output bias')
        hidden_units = output_layer_fun(weights, smiles)
        return np.dot(hidden_units, output_weights) + output_bias

    def loss_fun(weights, smiles, targets):
        preds = prediction_fun(weights, smiles)
        acc_loss = np.sum((preds - targets)**2)
        l2reg = l2_penalty * np.sum(weights**2)
        return acc_loss + l2reg
{'fraction_non_alphanumeric': 0.10403120936280884, 'fraction_numerical': 0.007730096806819823, 'mean_word_length': 2.9574042309891366, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 14, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '217218', 'n_tokens_mistral': 4381, 'n_tokens_neox': 4199, 'n_words': 1037}