Commit c08d9e1d authored by miaecle's avatar miaecle
Browse files

refractor DTNN

parent dd6f32b0

deepchem/models/tf_new_models/DTNN_regressor.py

deleted100644 → 0
+0 −30
Original line number Diff line number Diff line 
import tensorflow as tf

from deepchem.models.tf_new_models.multitask_regressor import MultitaskGraphRegressor





class DTNNGraphRegressor(MultitaskGraphRegressor):



  def build(self):

    # Create target inputs

    self.label_placeholder = tf.placeholder(

        dtype='float32', shape=(None, self.n_tasks), name="label_placeholder")

    self.weight_placeholder = tf.placeholder(

        dtype='float32', shape=(None, self.n_tasks), name="weight_placholder")



    feat = self.model.return_outputs()

    feat_size = self.feat_dim

    # dimension of `feat` becomes Unknown after tf.tensordot operation

    # need to define dimension of W and b explicitly

    outputs = []

    W_list = []

    b_list = []

    for task in range(self.n_tasks):

      W_list.append(

          tf.Variable(

              tf.truncated_normal([feat_size, 1], stddev=0.01),

              name='w',

              dtype=tf.float32))

      b_list.append(tf.Variable(tf.zeros([1]), name='b', dtype=tf.float32))

      outputs.append(

          tf.squeeze(tf.nn.xw_plus_b(feat, W_list[task], b_list[task])))

    return outputs

deepchem/models/tf_new_models/graph_models.py

+3 −10
Original line number Diff line number Diff line
@@ -86,11 +86,7 @@ class SequentialDTNNGraph(SequentialGraph): 
  automatically generates and passes topology placeholders to each layer. 

  """



  def __init__(self,

               max_n_atoms,

               n_distance=100,

               distance_min=-1.,

               distance_max=18.):

  def __init__(self, n_distance=100, distance_min=-1., distance_max=18.):

    """

    Parameters

    ----------
@@ -107,10 +103,7 @@ class SequentialDTNNGraph(SequentialGraph): 
    self.graph = tf.Graph()

    with self.graph.as_default():

      self.graph_topology = DTNNGraphTopology(

          max_n_atoms,

          n_distance,

          distance_min=distance_min,

          distance_max=distance_max)

          n_distance, distance_min=distance_min, distance_max=distance_max)

      self.output = self.graph_topology.get_atom_number_placeholder()

    # Keep track of the layers

    self.layers = []
@@ -123,7 +116,7 @@ class SequentialDTNNGraph(SequentialGraph): 
                            self.graph_topology.get_topology_placeholders())

      elif type(layer).__name__ in ['DTNNGather']:

        self.output = layer(

            [self.output, self.graph_topology.atom_mask_placeholder])

            [self.output, self.graph_topology.atom_membership_placeholder])

      else:

        self.output = layer(self.output)

      self.layers.append(layer)

deepchem/models/tf_new_models/graph_topology.py

+52 −46
Original line number Diff line number Diff line
@@ -146,7 +146,6 @@ class DTNNGraphTopology(GraphTopology): 
  """Manages placeholders associated with batch of graphs and their topology"""



  def __init__(self,

               max_n_atoms,

               n_distance=100,

               distance_min=-1.,

               distance_max=18.,
@@ -167,31 +166,29 @@ class DTNNGraphTopology(GraphTopology): 


    #self.n_atoms = n_atoms

    self.name = name

    self.max_n_atoms = max_n_atoms

    self.n_distance = n_distance

    self.distance_min = distance_min

    self.distance_max = distance_max



    self.atom_number_placeholder = tf.placeholder(

        dtype='int32',

        shape=(None, self.max_n_atoms),

        name=self.name + '_atom_number')

    self.atom_mask_placeholder = tf.placeholder(

        dtype='float32',

        shape=(None, self.max_n_atoms),

        name=self.name + '_atom_mask')

    self.distance_matrix_placeholder = tf.placeholder(

        dtype='float32',

        shape=(None, self.max_n_atoms, self.max_n_atoms, self.n_distance),

        name=self.name + '_distance_matrix')

    self.distance_matrix_mask_placeholder = tf.placeholder(

        dtype='int32', shape=(None,), name=self.name + '_atom_number')

    self.distance_placeholder = tf.placeholder(

        dtype='float32',

        shape=(None, self.max_n_atoms, self.max_n_atoms),

        name=self.name + '_distance_matrix_mask')

        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')

    self.distance_membership_j_placeholder = tf.placeholder(

        dtype='int32', shape=(None,), name=self.name + '_distance_membership')



    # Define the list of tensors to be used as topology

    self.topology = [

        self.distance_matrix_placeholder, self.distance_matrix_mask_placeholder

        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
@@ -199,8 +196,8 @@ class DTNNGraphTopology(GraphTopology): 
  def get_atom_number_placeholder(self):

    return self.atom_number_placeholder



  def get_distance_matrix_placeholder(self):

    return self.distance_matrix_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.
@@ -219,36 +216,44 @@ class DTNNGraphTopology(GraphTopology): 
      Can be merged with other feed_dicts for input into tensorflow

    """

    # Extract atom numbers

    atom_number = np.asarray(list(map(np.diag, batch)))

    atom_mask = np.sign(atom_number)

    atom_number = np.asarray(

        np.round(np.power(2 * atom_number, 1 / 2.4)), dtype=int)

    ZiZj = []

    for molecule in atom_number:

      ZiZj.append(np.outer(molecule, molecule))

    ZiZj = np.asarray(ZiZj)

    distance_matrix = np.expand_dims(batch[:], axis=3)

    distance_matrix = np.concatenate(

        [distance_matrix] * self.n_distance, axis=3)

    distance_matrix_mask = batch[:]

    for im, molecule in enumerate(batch):

      for ir, row in enumerate(molecule):

        for ie, element in enumerate(row):

          if element > 0 and ir != ie:

            # expand a float value distance to a distance vector

            distance_matrix[im, ir, ie, :] = self.gauss_expand(

                ZiZj[im, ir, ie] / element, self.n_distance, self.distance_min,

                self.distance_max)

            distance_matrix_mask[im, ir, ie] = 1

          else:

            distance_matrix[im, ir, ie, :] = 0

            distance_matrix_mask[im, ir, ie] = 0

    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 = batch[im, :num_atoms[im], :num_atoms[im]] / np.outer(

          molecule, molecule)

      np.fill_diagonal(distance_matrix, 0)

      distance_matrix = np.expand_dims(distance_matrix.flatten(), 1)

      distance_matrix = self.gauss_expand(distance_matrix, self.n_distance,

                                          self.distance_min, self.distance_max)

      distance.append(distance_matrix)

      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)

    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.atom_mask_placeholder: atom_mask,

        self.distance_matrix_placeholder: distance_matrix,

        self.distance_matrix_mask_placeholder: distance_matrix_mask

        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
@@ -256,6 +261,7 @@ class DTNNGraphTopology(GraphTopology): 
  def gauss_expand(distance, n_distance, distance_min, distance_max):

    step_size = (distance_max - distance_min) / n_distance

    steps = np.array([distance_min + i * step_size for i in range(n_distance)])

    steps = np.expand_dims(steps, 0)

    distance_vector = np.exp(-np.square(distance - steps) / (2 * step_size**2))

    return distance_vector

deepchem/nn/layers.py

+17 −15
Original line number Diff line number Diff line
@@ -914,20 +914,23 @@ class DTNNStep(Layer): 
    """

    self.build()

    atom_features = x[0]

    distance_matrix = x[1]

    distance_matrix_mask = x[2]

    outputs = tf.multiply(

        (tf.tensordot(distance_matrix, self.W_df, [[3], [0]]) + self.b_df),

        tf.expand_dims(

            tf.tensordot(atom_features, self.W_cf, [[2], [0]]) + self.b_cf,

            axis=1))

    distance = x[1]

    atom_membership = x[2]

    distance_membership_i = x[3]

    distance_membership_j = x[3]

    distance = tf.matmul(distance, self.W_df) + self.b_df

    distance = self.activation(distance)

    atom_features = tf.matmul(atom_features, self.W_cf) + self.b_cf

    atom_features = self.activation(atom_features)

    outputs = tf.multiply(distance,

                          tf.gather(atom_features, distance_membership_j))



    # for atom i in a molecule m, this step multiplies together distance info of atom pair(i,j)

    # and embeddings of atom j(both gone through a hidden layer)

    outputs = tf.tensordot(outputs, self.W_fc, [[3], [0]])

    outputs = tf.multiply(outputs, tf.expand_dims(distance_matrix_mask, axis=3))

    # masking the outputs tensor for pair(i,i) and all paddings

    outputs = tf.matmul(outputs, self.W_fc)

    outputs = self.activation(outputs)

    outputs = tf.reduce_sum(outputs, axis=2) + atom_features



    outputs = tf.segment_sum(outputs, distance_membership_i) + atom_features

    # for atom i, sum the influence from all other atom j in the molecule



    return outputs
@@ -983,13 +986,12 @@ class DTNNGather(Layer): 
    """

    self.build()

    output = x[0]

    atom_mask = x[1]

    atom_membership = x[1]

    for idw, W in enumerate(self.W_list):

      output = tf.tensordot(output, W, [[2], [0]]) + self.b_list[idw]

      output = tf.matmul(output, W) + self.b_list[idw]

      output = self.activation(output)



    output = tf.reduce_sum(

        tf.multiply(output, tf.expand_dims(atom_mask, axis=2)), axis=1)

    output = tf.segment_sum(output, atom_membership)



    return output

examples/qm7/qm7_DTNN.py

+2 −2
Original line number Diff line number Diff line
@@ -24,14 +24,14 @@ metric = [ 
# Batch size of models

batch_size = 50

n_embedding = 20

graph_model = dc.nn.SequentialDTNNGraph(max_n_atoms=23, n_distance=100)

graph_model = dc.nn.SequentialDTNNGraph(n_distance=100)

graph_model.add(dc.nn.DTNNEmbedding(n_embedding=n_embedding))

graph_model.add(dc.nn.DTNNStep(n_embedding=n_embedding, n_distance=100))

graph_model.add(dc.nn.DTNNStep(n_embedding=n_embedding, n_distance=100))

graph_model.add(dc.nn.DTNNGather(n_embedding=n_embedding))

n_feat = n_embedding



model = dc.models.DTNNGraphRegressor(

model = dc.models.MultitaskGraphRegressor(

    graph_model,

    len(tasks),

    n_feat,

CRA Git | Maintained and supported by SUSTech CRA and CCSE