<filename>deepchem/models/tf_new_models/graph_topology.py
"""Manages Placeholders for Graph convolution networks.
"""
from __future__ import print_function
from __future__ import division
from __future__ import unicode_literals

__author__ = "<NAME> and <NAME>"
__copyright__ = "Copyright 2016, Stanford University"
__license__ = "MIT"

import numpy as np
import tensorflow as tf
from deepchem.nn.copy import Input
from deepchem.feat.mol_graphs import ConvMol


def merge_two_dicts(x, y):
  z = x.copy()
  z.update(y)
  return z


def merge_dicts(l):
  """Convenience function to merge list of dictionaries."""
  merged = {}
  for dict in l:
    merged = merge_two_dicts(merged, dict)
  return merged


class GraphTopology(object):
  """Manages placeholders associated with batch of graphs and their topology"""

  def __init__(self, n_feat, name='topology', max_deg=10, min_deg=0):
    """
    Note that batch size is not specified in a GraphTopology object. A batch
    of molecules must be combined into a disconnected graph and fed to topology
    directly to handle batches.

    Parameters
    ----------
    n_feat: int
      Number of features per atom.
    name: str, optional
      Name of this manager.
    max_deg: int, optional
      Maximum #bonds for atoms in molecules.
    min_deg: int, optional
      Minimum #bonds for atoms in molecules.
    """

    #self.n_atoms = n_atoms
    self.n_feat = n_feat

    self.name = name
    self.max_deg = max_deg
    self.min_deg = min_deg

    self.atom_features_placeholder = tensor = tf.placeholder(
        dtype='float32',
        shape=(None, self.n_feat),
        name=self.name + '_atom_features')
    self.deg_adj_lists_placeholders = [
        tf.placeholder(
            dtype='int32',
            shape=(None, deg),
            name=self.name + '_deg_adj' + str(deg))
        for deg in range(1, self.max_deg + 1)
    ]
    self.deg_slice_placeholder = tf.placeholder(
        dtype='int32',
        shape=(self.max_deg - self.min_deg + 1, 2),
        name=self.name + '_deg_slice')
    self.membership_placeholder = tf.placeholder(
        dtype='int32', shape=(None,), name=self.name + '_membership')

    # Define the list of tensors to be used as topology
    self.topology = [self.deg_slice_placeholder, self.membership_placeholder]
    self.topology += self.deg_adj_lists_placeholders

    self.inputs = [self.atom_features_placeholder]
    self.inputs += self.topology

  def get_input_placeholders(self):
    """All placeholders.

    Contains atom_features placeholder and topology placeholders.
    """
    return self.inputs

  def get_topology_placeholders(self):
    """Returns topology placeholders

    Consists of deg_slice_placeholder, membership_placeholder, and the
    deg_adj_list_placeholders.
    """
    return self.topology

  def get_atom_features_placeholder(self):
    return self.atom_features_placeholder

  def get_deg_adjacency_lists_placeholders(self):
    return self.deg_adj_lists_placeholders

  def get_deg_slice_placeholder(self):
    return self.deg_slice_placeholder

  def get_membership_placeholder(self):
    return self.membership_placeholder

  def batch_to_feed_dict(self, batch):
    """Converts the current batch of mol_graphs into tensorflow feed_dict.

    Assigns the graph information in array of ConvMol objects to the
    placeholders tensors

    params
    ------
    batch : np.ndarray
      Array of ConvMol objects

    returns
    -------
    feed_dict : dict
      Can be merged with other feed_dicts for input into tensorflow
    """
    # Merge mol conv objects
    batch = ConvMol.agglomerate_mols(batch)
    atoms = batch.get_atom_features()
    deg_adj_lists = [
        batch.deg_adj_lists[deg] for deg in range(1, self.max_deg + 1)
    ]

    # Generate dicts
    deg_adj_dict = dict(
        list(zip(self.deg_adj_lists_placeholders, deg_adj_lists)))
    atoms_dict = {
        self.atom_features_placeholder: atoms,
        self.deg_slice_placeholder: batch.deg_slice,
        self.membership_placeholder: batch.membership
    }
    return merge_dicts([atoms_dict, deg_adj_dict])


class DTNNGraphTopology(GraphTopology):
  """Manages placeholders associated with batch of graphs and their topology"""

  def __init__(self,
               n_distance=100,
               distance_min=-1.,
               distance_max=18.,
               name='DTNN_topology'):
    """
    Parameters
    ----------
    n_distance: int, optional
      granularity of distance matrix
      step size will be (distance_max-distance_min)/n_distance
    distance_min: float, optional
      minimum distance of atom pairs, default = -1 Angstorm
    distance_max: float, optional
      maximum distance of atom pairs, default = 18 Angstorm
    """

    #self.n_atoms = n_atoms
    self.name = name
    self.n_distance = n_distance
    self.distance_min = distance_min
    self.distance_max = distance_max
    self.step_size = (distance_max - distance_min) / n_distance
    self.steps = np.array(
        [distance_min + i * self.step_size for i in range(n_distance)])
    self.steps = np.expand_dims(self.steps, 0)

    self.atom_number_placeholder = tf.placeholder(
        dtype='int32', shape=(None,), name=self.name + '_atom_number')
    self.distance_placeholder = tf.placeholder(
        dtype='float32',
        shape=(None, self.n_distance),
        name=self.name + '_distance')
    self.atom_membership_placeholder = tf.placeholder(
        dtype='int32', shape=(None,), name=self.name + '_atom_membership')
    self.distance_membership_i_placeholder = tf.placeholder(
        dtype='int32', shape=(None,), name=self.name + '_distance_membership_i')
    self.distance_membership_j_placeholder = tf.placeholder(
        dtype='int32', shape=(None,), name=self.name + '_distance_membership_j')

    # Define the list of tensors to be used as topology
    self.topology = [
        self.distance_placeholder,
        self.atom_membership_placeholder,
        self.distance_membership_i_placeholder,
        self.distance_membership_j_placeholder,
    ]
    self.inputs = [self.atom_number_placeholder]
    self.inputs += self.topology

  def get_atom_number_placeholder(self):
    return self.atom_number_placeholder

  def get_distance_placeholder(self):
    return self.distance_placeholder

  def batch_to_feed_dict(self, batch):
    """Converts the current batch of Coulomb Matrix into tensorflow feed_dict.

    Assigns the atom number and distance info to the
    placeholders tensors

    params
    ------
    batch : np.ndarray
      Array of Coulomb Matrix

    returns
    -------
    feed_dict : dict
      Can be merged with other feed_dicts for input into tensorflow
    """
    # Extract atom numbers
    num_atoms = list(map(sum, batch.astype(bool)[:, :, 0]))
    atom_number = [
        np.round(
            np.power(2 * np.diag(batch[i, :num_atoms[i], :num_atoms[i]]), 1 /
                     2.4)).astype(int) for i in range(len(num_atoms))
    ]

    distance = []
    atom_membership = []
    distance_membership_i = []
    distance_membership_j = []
    start = 0
    for im, molecule in enumerate(atom_number):
      distance_matrix = np.outer(
          molecule, molecule) / batch[im, :num_atoms[im], :num_atoms[im]]
      np.fill_diagonal(distance_matrix, -100)
      distance.append(np.expand_dims(distance_matrix.flatten(), 1))
      atom_membership.append([im] * num_atoms[im])
      membership = np.array([np.arange(num_atoms[im])] * num_atoms[im])
      membership_i = membership.flatten(order='F')
      membership_j = membership.flatten()
      distance_membership_i.append(membership_i + start)
      distance_membership_j.append(membership_j + start)
      start = start + num_atoms[im]
    atom_number = np.concatenate(atom_number)
    distance = np.concatenate(distance, 0)
    distance = np.exp(-np.square(distance - self.steps) /
                      (2 * self.step_size**2))
    distance_membership_i = np.concatenate(distance_membership_i)
    distance_membership_j = np.concatenate(distance_membership_j)
    atom_membership = np.concatenate(atom_membership)
    # Generate dicts
    dict_DTNN = {
        self.atom_number_placeholder: atom_number,
        self.distance_placeholder: distance,
        self.atom_membership_placeholder: atom_membership,
        self.distance_membership_i_placeholder: distance_membership_i,
        self.distance_membership_j_placeholder: distance_membership_j
    }
    return dict_DTNN


class DAGGraphTopology(GraphTopology):
  """GraphTopology for DAG models
  """

  def __init__(self, n_atom_feat=75, max_atoms=50, name='topology'):
    """
    Parameters
    ----------
    n_atom_feat: int, optional
      Number of features per atom.
    max_atoms: int, optional
      Maximum number of atoms in a molecule, should be defined based on dataset
    """
    self.n_atom_feat = n_atom_feat
    self.max_atoms = max_atoms
    self.name = name
    self.atom_features_placeholder = tf.placeholder(
        dtype='float32',
        shape=(None, self.n_atom_feat),
        name=self.name + '_atom_features')

    self.parents_placeholder = tf.placeholder(
        dtype='int32',
        shape=(None, self.max_atoms, self.max_atoms),
        # molecule * atom(graph) => step => features
        name=self.name + '_parents')

    self.calculation_orders_placeholder = tf.placeholder(
        dtype='int32',
        shape=(None, self.max_atoms),
        # molecule * atom(graph) => step
        name=self.name + '_orders')

    self.calculation_masks_placeholder = tf.placeholder(
        dtype='bool',
        shape=(None, self.max_atoms),
        # molecule * atom(graph) => step
        name=self.name + '_masks')

    self.membership_placeholder = tf.placeholder(
        dtype='int32', shape=(None,), name=self.name + '_membership')

    self.n_atoms_placeholder = tf.placeholder(
        dtype='int32', shape=(), name=self.name + '_n_atoms')

    # Define the list of tensors to be used as topology
    self.topology = [
        self.parents_placeholder, self.calculation_orders_placeholder,
        self.calculation_masks_placeholder, self.membership_placeholder,
        self.n_atoms_placeholder
    ]

    self.inputs = [self.atom_features_placeholder]
    self.inputs += self.topology

  def get_parents_placeholder(self):
    return self.parents_placeholder

  def get_calculation_orders_placeholder(self):
    return self.calculation_orders_placeholder

  def batch_to_feed_dict(self, batch):
    """Converts the current batch of mol_graphs into tensorflow feed_dict.

    Assigns the graph information in array of ConvMol objects to the
    placeholders tensors for DAG models

    params
    ------
    batch : np.ndarray
      Array of ConvMol objects

    returns
    -------
    feed_dict : dict
      Can be merged with other feed_dicts for input into tensorflow
    """

    atoms_per_mol = [mol.get_num_atoms() for mol in batch]
    n_atoms = sum(atoms_per_mol)
    start_index = [0] + list(np.cumsum(atoms_per_mol)[:-1])

    atoms_all = []
    # calculation orders for a batch of molecules
    parents_all = []
    calculation_orders = []
    calculation_masks = []
    membership = []
    for idm, mol in enumerate(batch):
      # padding atom features vector of each molecule with 0
      atoms_all.append(mol.get_atom_features())
      parents = mol.parents
      parents_all.extend(parents)
      calculation_index = np.array(parents)[:, :, 0]
      mask = np.array(calculation_index - self.max_atoms, dtype=bool)
      calculation_orders.append(calculation_index + start_index[idm])
      calculation_masks.append(mask)
      membership.extend([idm] * atoms_per_mol[idm])

    atoms_all = np.concatenate(atoms_all, axis=0)
    parents_all = np.stack(parents_all, axis=0)
    calculation_orders = np.concatenate(calculation_orders, axis=0)
    calculation_masks = np.concatenate(calculation_masks, axis=0)
    membership = np.array(membership)

    atoms_dict = {
        self.atom_features_placeholder: atoms_all,
        self.parents_placeholder: parents_all,
        self.calculation_orders_placeholder: calculation_orders,
        self.calculation_masks_placeholder: calculation_masks,
        self.membership_placeholder: membership,
        self.n_atoms_placeholder: n_atoms
    }

    return atoms_dict


class WeaveGraphTopology(GraphTopology):
  """Manages placeholders associated with batch of graphs and their topology"""

  def __init__(self,
               max_atoms=50,
               n_atom_feat=75,
               n_pair_feat=14,
               name='Weave_topology'):
    """
    Parameters
    ----------
    max_atoms: int, optional
      maximum number of atoms in a molecule
    n_atom_feat: int, optional
      number of basic features of each atom
    n_pair_feat: int, optional
      number of basic features of each pair
    """

    #self.n_atoms = n_atoms
    self.name = name
    self.max_atoms = max_atoms
    self.n_atom_feat = n_atom_feat
    self.n_pair_feat = n_pair_feat

    self.atom_features_placeholder = tf.placeholder(
        dtype='float32',
        shape=(None, self.max_atoms, self.n_atom_feat),
        name=self.name + '_atom_features')
    self.atom_mask_placeholder = tf.placeholder(
        dtype='float32',
        shape=(None, self.max_atoms),
        name=self.name + '_atom_mask')
    self.pair_features_placeholder = tf.placeholder(
        dtype='float32',
        shape=(None, self.max_atoms, self.max_atoms, self.n_pair_feat),
        name=self.name + '_pair_features')
    self.pair_mask_placeholder = tf.placeholder(
        dtype='float32',
        shape=(None, self.max_atoms, self.max_atoms),
        name=self.name + '_pair_mask')
    self.membership_placeholder = tf.placeholder(
        dtype='int32', shape=(None,), name=self.name + '_membership')
    # Define the list of tensors to be used as topology
    self.topology = [self.atom_mask_placeholder, self.pair_mask_placeholder]
    self.inputs = [self.atom_features_placeholder]
    self.inputs += self.topology

  def get_pair_features_placeholder(self):
    return self.pair_features_placeholder

  def batch_to_feed_dict(self, batch):
    """Converts the current batch of WeaveMol into tensorflow feed_dict.

    Assigns the atom features and pair features to the
    placeholders tensors

    params
    ------
    batch : np.ndarray
      Array of WeaveMol

    returns
    -------
    feed_dict : dict
      Can be merged with other feed_dicts for input into tensorflow
    """
    # Extract atom numbers
    atom_feat = []
    pair_feat = []
    atom_mask = []
    pair_mask = []
    membership = []
    max_atoms = self.max_atoms
    for im, mol in enumerate(batch):
      n_atoms = mol.get_num_atoms()
      atom_feat.append(
          np.pad(mol.get_atom_features(), ((0, max_atoms - n_atoms), (0, 0)),
                 'constant'))
      atom_mask.append(
          np.array([1] * n_atoms + [0] * (max_atoms - n_atoms), dtype=float))
      pair_feat.append(
          np.pad(mol.get_pair_features(), ((0, max_atoms - n_atoms), (
              0, max_atoms - n_atoms), (0, 0)), 'constant'))
      pair_mask.append(np.array([[1]*n_atoms + [0]*(max_atoms-n_atoms)]*n_atoms + \
                       [[0]*max_atoms]*(max_atoms-n_atoms), dtype=float))
      membership.extend([im] * n_atoms)
    atom_feat = np.stack(atom_feat)
    pair_feat = np.stack(pair_feat)
    atom_mask = np.stack(atom_mask)
    pair_mask = np.stack(pair_mask)
    membership = np.array(membership)
    # Generate dicts
    dict_DTNN = {
        self.atom_features_placeholder: atom_feat,
        self.pair_features_placeholder: pair_feat,
        self.atom_mask_placeholder: atom_mask,
        self.pair_mask_placeholder: pair_mask,
        self.membership_placeholder: membership
    }
    return dict_DTNN


class AlternateWeaveGraphTopology(GraphTopology):
  """Manages placeholders associated with batch of graphs and their topology"""

  def __init__(self,
               batch_size,
               max_atoms=50,
               n_atom_feat=75,
               n_pair_feat=14,
               name='Weave_topology'):
    """
    Parameters
    ----------
    batch_size: int
      number of molecules in a batch
    max_atoms: int, optional
      maximum number of atoms in a molecule
    n_atom_feat: int, optional
      number of basic features of each atom
    n_pair_feat: int, optional
      number of basic features of each pair
    """

    #self.n_atoms = n_atoms
    self.name = name
    self.batch_size = batch_size
    self.max_atoms = max_atoms * batch_size
    self.n_atom_feat = n_atom_feat
    self.n_pair_feat = n_pair_feat

    self.atom_features_placeholder = tf.placeholder(
        dtype='float32',
        shape=(None, self.n_atom_feat),
        name=self.name + '_atom_features')
    self.pair_features_placeholder = tf.placeholder(
        dtype='float32',
        shape=(None, self.n_pair_feat),
        name=self.name + '_pair_features')
    self.pair_split_placeholder = tf.placeholder(
        dtype='int32', shape=(None,), name=self.name + '_pair_split')
    self.atom_split_placeholder = tf.placeholder(
        dtype='int32', shape=(None,), name=self.name + '_atom_split')
    self.atom_to_pair_placeholder = tf.placeholder(
        dtype='int32', shape=(None, 2), name=self.name + '_atom_to_pair')

    # Define the list of tensors to be used as topology
    self.topology = [
        self.pair_split_placeholder, self.atom_split_placeholder,
        self.atom_to_pair_placeholder
    ]
    self.inputs = [self.atom_features_placeholder]
    self.inputs += self.topology

  def get_pair_features_placeholder(self):
    return self.pair_features_placeholder

  def batch_to_feed_dict(self, batch):
    """Converts the current batch of WeaveMol into tensorflow feed_dict.

    Assigns the atom features and pair features to the
    placeholders tensors

    params
    ------
    batch : np.ndarray
      Array of WeaveMol

    returns
    -------
    feed_dict : dict
      Can be merged with other feed_dicts for input into tensorflow
    """
    # Extract atom numbers
    atom_feat = []
    pair_feat = []
    atom_split = []
    atom_to_pair = []
    pair_split = []
    max_atoms = self.max_atoms
    start = 0
    for im, mol in enumerate(batch):
      n_atoms = mol.get_num_atoms()
      # number of atoms in each molecule
      atom_split.extend([im] * n_atoms)
      # index of pair features
      C0, C1 = np.meshgrid(np.arange(n_atoms), np.arange(n_atoms))
      atom_to_pair.append(
          np.transpose(np.array([C1.flatten() + start, C0.flatten() + start])))
      # number of pairs for each atom
      pair_split.extend(C1.flatten() + start)
      start = start + n_atoms

      # atom features
      atom_feat.append(mol.get_atom_features())
      # pair features
      pair_feat.append(
          np.reshape(mol.get_pair_features(), (n_atoms * n_atoms,
                                               self.n_pair_feat)))

    atom_feat = np.concatenate(atom_feat, axis=0)
    pair_feat = np.concatenate(pair_feat, axis=0)
    atom_to_pair = np.concatenate(atom_to_pair, axis=0)
    atom_split = np.array(atom_split)
    # Generate dicts
    dict_DTNN = {
        self.atom_features_placeholder: atom_feat,
        self.pair_features_placeholder: pair_feat,
        self.pair_split_placeholder: pair_split,
        self.atom_split_placeholder: atom_split,
        self.atom_to_pair_placeholder: atom_to_pair
    }
    return dict_DTNN

<filename>deepchem/models/tf_new_models/graph_topology.py
"""Manages Placeholders for Graph convolution networks.
"""
from __future__ import print_function
from __future__ import division
from __future__ import unicode_literals

__author__ = "<NAME> and <NAME>"
__copyright__ = "Copyright 2016, Stanford University"
__license__ = "MIT"

import numpy as np
import tensorflow as tf
from deepchem.nn.copy import Input
from deepchem.feat.mol_graphs import ConvMol


def merge_two_dicts(x, y):
  z = x.copy()
  z.update(y)
  return z


def merge_dicts(l):
  """Convenience function to merge list of dictionaries."""
  merged = {}
  for dict in l:
    merged = merge_two_dicts(merged, dict)
  return merged


class GraphTopology(object):
  """Manages placeholders associated with batch of graphs and their topology"""

  def __init__(self, n_feat, name='topology', max_deg=10, min_deg=0):
    """
    Note that batch size is not specified in a GraphTopology object. A batch
    of molecules must be combined into a disconnected graph and fed to topology
    directly to handle batches.

    Parameters
    ----------
    n_feat: int
      Number of features per atom.
    name: str, optional
      Name of this manager.
    max_deg: int, optional
      Maximum #bonds for atoms in molecules.
    min_deg: int, optional
      Minimum #bonds for atoms in molecules.
    """

    #self.n_atoms = n_atoms
    self.n_feat = n_feat

    self.name = name
    self.max_deg = max_deg
    self.min_deg = min_deg

    self.atom_features_placeholder = tensor = tf.placeholder(
        dtype='float32',
        shape=(None, self.n_feat),
        name=self.name + '_atom_features')
    self.deg_adj_lists_placeholders = [
        tf.placeholder(
            dtype='int32',
            shape=(None, deg),
            name=self.name + '_deg_adj' + str(deg))
        for deg in range(1, self.max_deg + 1)
    ]
    self.deg_slice_placeholder = tf.placeholder(
        dtype='int32',
        shape=(self.max_deg - self.min_deg + 1, 2),
        name=self.name + '_deg_slice')
    self.membership_placeholder = tf.placeholder(
        dtype='int32', shape=(None,), name=self.name + '_membership')

    # Define the list of tensors to be used as topology
    self.topology = [self.deg_slice_placeholder, self.membership_placeholder]
    self.topology += self.deg_adj_lists_placeholders

    self.inputs = [self.atom_features_placeholder]
    self.inputs += self.topology

  def get_input_placeholders(self):
    """All placeholders.

    Contains atom_features placeholder and topology placeholders.
    """
    return self.inputs

  def get_topology_placeholders(self):
    """Returns topology placeholders

    Consists of deg_slice_placeholder, membership_placeholder, and the
    deg_adj_list_placeholders.
    """
    return self.topology

  def get_atom_features_placeholder(self):
    return self.atom_features_placeholder

  def get_deg_adjacency_lists_placeholders(self):
    return self.deg_adj_lists_placeholders

  def get_deg_slice_placeholder(self):
    return self.deg_slice_placeholder

  def get_membership_placeholder(self):
    return self.membership_placeholder

  def batch_to_feed_dict(self, batch):
    """Converts the current batch of mol_graphs into tensorflow feed_dict.

    Assigns the graph information in array of ConvMol objects to the
    placeholders tensors

    params
    ------
    batch : np.ndarray
      Array of ConvMol objects

    returns
    -------
    feed_dict : dict
      Can be merged with other feed_dicts for input into tensorflow
    """
    # Merge mol conv objects
    batch = ConvMol.agglomerate_mols(batch)
    atoms = batch.get_atom_features()
    deg_adj_lists = [
        batch.deg_adj_lists[deg] for deg in range(1, self.max_deg + 1)
    ]

    # Generate dicts
    deg_adj_dict = dict(
        list(zip(self.deg_adj_lists_placeholders, deg_adj_lists)))
    atoms_dict = {
        self.atom_features_placeholder: atoms,
        self.deg_slice_placeholder: batch.deg_slice,
        self.membership_placeholder: batch.membership
    }
    return merge_dicts([atoms_dict, deg_adj_dict])


class DTNNGraphTopology(GraphTopology):
  """Manages placeholders associated with batch of graphs and their topology"""

  def __init__(self,
               n_distance=100,
               distance_min=-1.,
               distance_max=18.,
               name='DTNN_topology'):
    """
    Parameters
    ----------
    n_distance: int, optional
      granularity of distance matrix
      step size will be (distance_max-distance_min)/n_distance
    distance_min: float, optional
      minimum distance of atom pairs, default = -1 Angstorm
    distance_max: float, optional
      maximum distance of atom pairs, default = 18 Angstorm
    """

    #self.n_atoms = n_atoms
    self.name = name
    self.n_distance = n_distance
    self.distance_min = distance_min
    self.distance_max = distance_max
    self.step_size = (distance_max - distance_min) / n_distance
    self.steps = np.array(
        [distance_min + i * self.step_size for i in range(n_distance)])
    self.steps = np.expand_dims(self.steps, 0)

    self.atom_number_placeholder = tf.placeholder(
        dtype='int32', shape=(None,), name=self.name + '_atom_number')
    self.distance_placeholder = tf.placeholder(
        dtype='float32',
        shape=(None, self.n_distance),
        name=self.name + '_distance')
    self.atom_membership_placeholder = tf.placeholder(
        dtype='int32', shape=(None,), name=self.name + '_atom_membership')
    self.distance_membership_i_placeholder = tf.placeholder(
        dtype='int32', shape=(None,), name=self.name + '_distance_membership_i')
    self.distance_membership_j_placeholder = tf.placeholder(
        dtype='int32', shape=(None,), name=self.name + '_distance_membership_j')

    # Define the list of tensors to be used as topology
    self.topology = [
        self.distance_placeholder,
        self.atom_membership_placeholder,
        self.distance_membership_i_placeholder,
        self.distance_membership_j_placeholder,
    ]
    self.inputs = [self.atom_number_placeholder]
    self.inputs += self.topology

  def get_atom_number_placeholder(self):
    return self.atom_number_placeholder

  def get_distance_placeholder(self):
    return self.distance_placeholder

  def batch_to_feed_dict(self, batch):
    """Converts the current batch of Coulomb Matrix into tensorflow feed_dict.

    Assigns the atom number and distance info to the
    placeholders tensors

    params
    ------
    batch : np.ndarray
      Array of Coulomb Matrix

    returns
    -------
    feed_dict : dict
      Can be merged with other feed_dicts for input into tensorflow
    """
    # Extract atom numbers
    num_atoms = list(map(sum, batch.astype(bool)[:, :, 0]))
    atom_number = [
        np.round(
            np.power(2 * np.diag(batch[i, :num_atoms[i], :num_atoms[i]]), 1 /
                     2.4)).astype(int) for i in range(len(num_atoms))
    ]

    distance = []
    atom_membership = []
    distance_membership_i = []
    distance_membership_j = []
    start = 0
    for im, molecule in enumerate(atom_number):
      distance_matrix = np.outer(
          molecule, molecule) / batch[im, :num_atoms[im], :num_atoms[im]]
      np.fill_diagonal(distance_matrix, -100)
      distance.append(np.expand_dims(distance_matrix.flatten(), 1))
      atom_membership.append([im] * num_atoms[im])
      membership = np.array([np.arange(num_atoms[im])] * num_atoms[im])
      membership_i = membership.flatten(order='F')
      membership_j = membership.flatten()
      distance_membership_i.append(membership_i + start)
      distance_membership_j.append(membership_j + start)
      start = start + num_atoms[im]
    atom_number = np.concatenate(atom_number)
    distance = np.concatenate(distance, 0)
    distance = np.exp(-np.square(distance - self.steps) /
                      (2 * self.step_size**2))
    distance_membership_i = np.concatenate(distance_membership_i)
    distance_membership_j = np.concatenate(distance_membership_j)
    atom_membership = np.concatenate(atom_membership)
    # Generate dicts
    dict_DTNN = {
        self.atom_number_placeholder: atom_number,
        self.distance_placeholder: distance,
        self.atom_membership_placeholder: atom_membership,
        self.distance_membership_i_placeholder: distance_membership_i,
        self.distance_membership_j_placeholder: distance_membership_j
    }
    return dict_DTNN


class DAGGraphTopology(GraphTopology):
  """GraphTopology for DAG models
  """

  def __init__(self, n_atom_feat=75, max_atoms=50, name='topology'):
    """
    Parameters
    ----------
    n_atom_feat: int, optional
      Number of features per atom.
    max_atoms: int, optional
      Maximum number of atoms in a molecule, should be defined based on dataset
    """
    self.n_atom_feat = n_atom_feat
    self.max_atoms = max_atoms
    self.name = name
    self.atom_features_placeholder = tf.placeholder(
        dtype='float32',
        shape=(None, self.n_atom_feat),
        name=self.name + '_atom_features')

    self.parents_placeholder = tf.placeholder(
        dtype='int32',
        shape=(None, self.max_atoms, self.max_atoms),
        # molecule * atom(graph) => step => features
        name=self.name + '_parents')

    self.calculation_orders_placeholder = tf.placeholder(
        dtype='int32',
        shape=(None, self.max_atoms),
        # molecule * atom(graph) => step
        name=self.name + '_orders')

    self.calculation_masks_placeholder = tf.placeholder(
        dtype='bool',
        shape=(None, self.max_atoms),
        # molecule * atom(graph) => step
        name=self.name + '_masks')

    self.membership_placeholder = tf.placeholder(
        dtype='int32', shape=(None,), name=self.name + '_membership')

    self.n_atoms_placeholder = tf.placeholder(
        dtype='int32', shape=(), name=self.name + '_n_atoms')

    # Define the list of tensors to be used as topology
    self.topology = [
        self.parents_placeholder, self.calculation_orders_placeholder,
        self.calculation_masks_placeholder, self.membership_placeholder,
        self.n_atoms_placeholder
    ]

    self.inputs = [self.atom_features_placeholder]
    self.inputs += self.topology

  def get_parents_placeholder(self):
    return self.parents_placeholder

  def get_calculation_orders_placeholder(self):
    return self.calculation_orders_placeholder

  def batch_to_feed_dict(self, batch):
    """Converts the current batch of mol_graphs into tensorflow feed_dict.

    Assigns the graph information in array of ConvMol objects to the
    placeholders tensors for DAG models

    params
    ------
    batch : np.ndarray
      Array of ConvMol objects

    returns
    -------
    feed_dict : dict
      Can be merged with other feed_dicts for input into tensorflow
    """

    atoms_per_mol = [mol.get_num_atoms() for mol in batch]
    n_atoms = sum(atoms_per_mol)
    start_index = [0] + list(np.cumsum(atoms_per_mol)[:-1])

    atoms_all = []
    # calculation orders for a batch of molecules
    parents_all = []
    calculation_orders = []
    calculation_masks = []
    membership = []
    for idm, mol in enumerate(batch):
      # padding atom features vector of each molecule with 0
      atoms_all.append(mol.get_atom_features())
      parents = mol.parents
      parents_all.extend(parents)
      calculation_index = np.array(parents)[:, :, 0]
      mask = np.array(calculation_index - self.max_atoms, dtype=bool)
      calculation_orders.append(calculation_index + start_index[idm])
      calculation_masks.append(mask)
      membership.extend([idm] * atoms_per_mol[idm])

    atoms_all = np.concatenate(atoms_all, axis=0)
    parents_all = np.stack(parents_all, axis=0)
    calculation_orders = np.concatenate(calculation_orders, axis=0)
    calculation_masks = np.concatenate(calculation_masks, axis=0)
    membership = np.array(membership)

    atoms_dict = {
        self.atom_features_placeholder: atoms_all,
        self.parents_placeholder: parents_all,
        self.calculation_orders_placeholder: calculation_orders,
        self.calculation_masks_placeholder: calculation_masks,
        self.membership_placeholder: membership,
        self.n_atoms_placeholder: n_atoms
    }

    return atoms_dict


class WeaveGraphTopology(GraphTopology):
  """Manages placeholders associated with batch of graphs and their topology"""

  def __init__(self,
               max_atoms=50,
               n_atom_feat=75,
               n_pair_feat=14,
               name='Weave_topology'):
    """
    Parameters
    ----------
    max_atoms: int, optional
      maximum number of atoms in a molecule
    n_atom_feat: int, optional
      number of basic features of each atom
    n_pair_feat: int, optional
      number of basic features of each pair
    """

    #self.n_atoms = n_atoms
    self.name = name
    self.max_atoms = max_atoms
    self.n_atom_feat = n_atom_feat
    self.n_pair_feat = n_pair_feat

    self.atom_features_placeholder = tf.placeholder(
        dtype='float32',
        shape=(None, self.max_atoms, self.n_atom_feat),
        name=self.name + '_atom_features')
    self.atom_mask_placeholder = tf.placeholder(
        dtype='float32',
        shape=(None, self.max_atoms),
        name=self.name + '_atom_mask')
    self.pair_features_placeholder = tf.placeholder(
        dtype='float32',
        shape=(None, self.max_atoms, self.max_atoms, self.n_pair_feat),
        name=self.name + '_pair_features')
    self.pair_mask_placeholder = tf.placeholder(
        dtype='float32',
        shape=(None, self.max_atoms, self.max_atoms),
        name=self.name + '_pair_mask')
    self.membership_placeholder = tf.placeholder(
        dtype='int32', shape=(None,), name=self.name + '_membership')
    # Define the list of tensors to be used as topology
    self.topology = [self.atom_mask_placeholder, self.pair_mask_placeholder]
    self.inputs = [self.atom_features_placeholder]
    self.inputs += self.topology

  def get_pair_features_placeholder(self):
    return self.pair_features_placeholder

  def batch_to_feed_dict(self, batch):
    """Converts the current batch of WeaveMol into tensorflow feed_dict.

    Assigns the atom features and pair features to the
    placeholders tensors

    params
    ------
    batch : np.ndarray
      Array of WeaveMol

    returns
    -------
    feed_dict : dict
      Can be merged with other feed_dicts for input into tensorflow
    """
    # Extract atom numbers
    atom_feat = []
    pair_feat = []
    atom_mask = []
    pair_mask = []
    membership = []
    max_atoms = self.max_atoms
    for im, mol in enumerate(batch):
      n_atoms = mol.get_num_atoms()
      atom_feat.append(
          np.pad(mol.get_atom_features(), ((0, max_atoms - n_atoms), (0, 0)),
                 'constant'))
      atom_mask.append(
          np.array([1] * n_atoms + [0] * (max_atoms - n_atoms), dtype=float))
      pair_feat.append(
          np.pad(mol.get_pair_features(), ((0, max_atoms - n_atoms), (
              0, max_atoms - n_atoms), (0, 0)), 'constant'))
      pair_mask.append(np.array([[1]*n_atoms + [0]*(max_atoms-n_atoms)]*n_atoms + \
                       [[0]*max_atoms]*(max_atoms-n_atoms), dtype=float))
      membership.extend([im] * n_atoms)
    atom_feat = np.stack(atom_feat)
    pair_feat = np.stack(pair_feat)
    atom_mask = np.stack(atom_mask)
    pair_mask = np.stack(pair_mask)
    membership = np.array(membership)
    # Generate dicts
    dict_DTNN = {
        self.atom_features_placeholder: atom_feat,
        self.pair_features_placeholder: pair_feat,
        self.atom_mask_placeholder: atom_mask,
        self.pair_mask_placeholder: pair_mask,
        self.membership_placeholder: membership
    }
    return dict_DTNN


class AlternateWeaveGraphTopology(GraphTopology):
  """Manages placeholders associated with batch of graphs and their topology"""

  def __init__(self,
               batch_size,
               max_atoms=50,
               n_atom_feat=75,
               n_pair_feat=14,
               name='Weave_topology'):
    """
    Parameters
    ----------
    batch_size: int
      number of molecules in a batch
    max_atoms: int, optional
      maximum number of atoms in a molecule
    n_atom_feat: int, optional
      number of basic features of each atom
    n_pair_feat: int, optional
      number of basic features of each pair
    """

    #self.n_atoms = n_atoms
    self.name = name
    self.batch_size = batch_size
    self.max_atoms = max_atoms * batch_size
    self.n_atom_feat = n_atom_feat
    self.n_pair_feat = n_pair_feat

    self.atom_features_placeholder = tf.placeholder(
        dtype='float32',
        shape=(None, self.n_atom_feat),
        name=self.name + '_atom_features')
    self.pair_features_placeholder = tf.placeholder(
        dtype='float32',
        shape=(None, self.n_pair_feat),
        name=self.name + '_pair_features')
    self.pair_split_placeholder = tf.placeholder(
        dtype='int32', shape=(None,), name=self.name + '_pair_split')
    self.atom_split_placeholder = tf.placeholder(
        dtype='int32', shape=(None,), name=self.name + '_atom_split')
    self.atom_to_pair_placeholder = tf.placeholder(
        dtype='int32', shape=(None, 2), name=self.name + '_atom_to_pair')

    # Define the list of tensors to be used as topology
    self.topology = [
        self.pair_split_placeholder, self.atom_split_placeholder,
        self.atom_to_pair_placeholder
    ]
    self.inputs = [self.atom_features_placeholder]
    self.inputs += self.topology

  def get_pair_features_placeholder(self):
    return self.pair_features_placeholder

  def batch_to_feed_dict(self, batch):
    """Converts the current batch of WeaveMol into tensorflow feed_dict.

    Assigns the atom features and pair features to the
    placeholders tensors

    params
    ------
    batch : np.ndarray
      Array of WeaveMol

    returns
    -------
    feed_dict : dict
      Can be merged with other feed_dicts for input into tensorflow
    """
    # Extract atom numbers
    atom_feat = []
    pair_feat = []
    atom_split = []
    atom_to_pair = []
    pair_split = []
    max_atoms = self.max_atoms
    start = 0
    for im, mol in enumerate(batch):
      n_atoms = mol.get_num_atoms()
      # number of atoms in each molecule
      atom_split.extend([im] * n_atoms)
      # index of pair features
      C0, C1 = np.meshgrid(np.arange(n_atoms), np.arange(n_atoms))
      atom_to_pair.append(
          np.transpose(np.array([C1.flatten() + start, C0.flatten() + start])))
      # number of pairs for each atom
      pair_split.extend(C1.flatten() + start)
      start = start + n_atoms

      # atom features
      atom_feat.append(mol.get_atom_features())
      # pair features
      pair_feat.append(
          np.reshape(mol.get_pair_features(), (n_atoms * n_atoms,
                                               self.n_pair_feat)))

    atom_feat = np.concatenate(atom_feat, axis=0)
    pair_feat = np.concatenate(pair_feat, axis=0)
    atom_to_pair = np.concatenate(atom_to_pair, axis=0)
    atom_split = np.array(atom_split)
    # Generate dicts
    dict_DTNN = {
        self.atom_features_placeholder: atom_feat,
        self.pair_features_placeholder: pair_feat,
        self.pair_split_placeholder: pair_split,
        self.atom_split_placeholder: atom_split,
        self.atom_to_pair_placeholder: atom_to_pair
    }
    return dict_DTNN

Manages Placeholders for Graph convolution networks. Convenience function to merge list of dictionaries. Manages placeholders associated with batch of graphs and their topology Note that batch size is not specified in a GraphTopology object. A batch     of molecules must be combined into a disconnected graph and fed to topology     directly to handle batches.      Parameters     ----------     n_feat: int       Number of features per atom.     name: str, optional       Name of this manager.     max_deg: int, optional       Maximum #bonds for atoms in molecules.     min_deg: int, optional       Minimum #bonds for atoms in molecules. #self.n_atoms = n_atoms # Define the list of tensors to be used as topology All placeholders.      Contains atom_features placeholder and topology placeholders. Returns topology placeholders      Consists of deg_slice_placeholder, membership_placeholder, and the     deg_adj_list_placeholders. Converts the current batch of mol_graphs into tensorflow feed_dict.      Assigns the graph information in array of ConvMol objects to the     placeholders tensors      params     ------     batch : np.ndarray       Array of ConvMol objects      returns     -------     feed_dict : dict       Can be merged with other feed_dicts for input into tensorflow # Merge mol conv objects # Generate dicts Manages placeholders associated with batch of graphs and their topology Parameters     ----------     n_distance: int, optional       granularity of distance matrix       step size will be (distance_max-distance_min)/n_distance     distance_min: float, optional       minimum distance of atom pairs, default = -1 Angstorm     distance_max: float, optional       maximum distance of atom pairs, default = 18 Angstorm #self.n_atoms = n_atoms # Define the list of tensors to be used as topology Converts the current batch of Coulomb Matrix into tensorflow feed_dict.      Assigns the atom number and distance info to the     placeholders tensors      params     ------     batch : np.ndarray       Array of Coulomb Matrix      returns     -------     feed_dict : dict       Can be merged with other feed_dicts for input into tensorflow # Extract atom numbers # Generate dicts GraphTopology for DAG models Parameters     ----------     n_atom_feat: int, optional       Number of features per atom.     max_atoms: int, optional       Maximum number of atoms in a molecule, should be defined based on dataset # molecule * atom(graph) => step => features # molecule * atom(graph) => step # molecule * atom(graph) => step # Define the list of tensors to be used as topology Converts the current batch of mol_graphs into tensorflow feed_dict.      Assigns the graph information in array of ConvMol objects to the     placeholders tensors for DAG models      params     ------     batch : np.ndarray       Array of ConvMol objects      returns     -------     feed_dict : dict       Can be merged with other feed_dicts for input into tensorflow # calculation orders for a batch of molecules # padding atom features vector of each molecule with 0 Manages placeholders associated with batch of graphs and their topology Parameters     ----------     max_atoms: int, optional       maximum number of atoms in a molecule     n_atom_feat: int, optional       number of basic features of each atom     n_pair_feat: int, optional       number of basic features of each pair #self.n_atoms = n_atoms # Define the list of tensors to be used as topology Converts the current batch of WeaveMol into tensorflow feed_dict.      Assigns the atom features and pair features to the     placeholders tensors      params     ------     batch : np.ndarray       Array of WeaveMol      returns     -------     feed_dict : dict       Can be merged with other feed_dicts for input into tensorflow # Extract atom numbers # Generate dicts Manages placeholders associated with batch of graphs and their topology Parameters     ----------     batch_size: int       number of molecules in a batch     max_atoms: int, optional       maximum number of atoms in a molecule     n_atom_feat: int, optional       number of basic features of each atom     n_pair_feat: int, optional       number of basic features of each pair #self.n_atoms = n_atoms # Define the list of tensors to be used as topology Converts the current batch of WeaveMol into tensorflow feed_dict.      Assigns the atom features and pair features to the     placeholders tensors      params     ------     batch : np.ndarray       Array of WeaveMol      returns     -------     feed_dict : dict       Can be merged with other feed_dicts for input into tensorflow # Extract atom numbers # number of atoms in each molecule # index of pair features # number of pairs for each atom # atom features # pair features # Generate dicts