Merge pull request #1594 from peastman/graphkeras

    must have the same shape as the corresponding prediction output, and each

    element is an estimate of the variance in the corresponding prediction.

    Also be aware that if a model supports uncertainty, it MUST use dropout on

    every layer.  Otherwise, the uncertainties it computes will be inaccurate.

    every layer, and dropout most be enabled during uncertainty prediction.

    Otherwise, the uncertainties it computes will be inaccurate.

  """

  def __init__(self,

    self._inputs_built = False

    self._training_ops_built = False

    self._initialized_vars = set()

    self._output_functions = {}

  def _ensure_built(self):

    """The first time this is called, create internal data structures."""

    if len(self._input_placeholders) == 1:

      self._output_tensors = self.model(

          self._input_placeholders[0], training=False)

      self._uncertainty_tensors = self.model(

          self._input_placeholders[0], training=True)

    else:

      self._output_tensors = self.model(

          self._input_placeholders, training=False)

      self._uncertainty_tensors = self.model(

          self._input_placeholders, training=True)

    if isinstance(self._output_tensors, tf.Tensor):

      self._output_tensors = [self._output_tensors]

    if self._prediction_outputs is None:

        # In eager mode we execute the loss function, accumulating the gradients.

        with tf.GradientTape() as tape:

          outputs = self.model(inputs[0])

          if len(inputs) == 1:

            inputs = inputs[0]

          outputs = self.model(inputs)

          if isinstance(outputs, tf.Tensor):

            outputs = [outputs]

          if self._loss_outputs is not None:

    dataset = NumpyDataset(X, y, w)

    return self.fit(dataset, nb_epoch=1)

  def _predict(self, generator, transformers, uncertainty):

  def _predict(self, generator, transformers, outputs, uncertainty):

    """

    Predict outputs for data provided by a generator.

    transformers: list of dc.trans.Transformers

      Transformers that the input data has been transformed by.  The output

      is passed through these transformers to undo the transformations.

    outputs: Tensor or list of Tensors

      The outputs to return.  If this is None, the model's standard prediction

      outputs will be returned.  Alternatively one or more Tensors within the

      model may be specified, in which case the output of those Tensors will be

      returned.

    uncertainty: bool

      specifies whether this is being called as part of estimating uncertainty.

      If True, it sets the training flag so that dropout will be enabled, and

    results = None

    variances = None

    if uncertainty:

      assert outputs is None

      if self._variance_outputs is None or len(self._variance_outputs) == 0:

        raise ValueError('This model cannot compute uncertainties')

      if len(self._variance_outputs) != len(self._prediction_outputs):

        raise ValueError(

            'The number of variances must exactly match the number of outputs')

    if tf.executing_eagerly() and outputs is not None and len(

        self.model.inputs) == 0:

      raise ValueError(

          "Cannot use 'outputs' argument in eager mode with a model that does not specify its inputs"

      )

    if isinstance(outputs, tf.Tensor):

      outputs = [outputs]

    for batch in generator:

      inputs, labels, weights = batch

      self._create_inputs(inputs)

        if len(inputs) == 1:

          inputs = inputs[0]

        outputs = self.model(inputs, training=uncertainty)

        outputs = [t.numpy() for t in outputs]

        if outputs is not None:

          outputs = tuple(outputs)

          if outputs not in self._output_functions:

            self._output_functions[outputs] = tf.keras.backend.function(

                self.model.inputs, outputs)

          output_values = self._output_functions[outputs](inputs)

        else:

          output_values = self.model(inputs, training=False)

          output_values = [t.numpy() for t in output_values]

      else:

        # In graph mode we execute the output tensors.

        if uncertainty:

          fetches = self._uncertainty_tensors

        if outputs is not None:

          fetches = outputs

        else:

          fetches = self._output_tensors

        feed_dict = dict(zip(self._input_placeholders, inputs))

        outputs = self.session.run(fetches, feed_dict=feed_dict)

        output_values = self.session.run(fetches, feed_dict=feed_dict)

      # Apply tranformers and record results.

      if uncertainty:

        var = [outputs[i] for i in self._variance_outputs]

        var = [output_values[i] for i in self._variance_outputs]

        if variances is None:

          variances = var

          variances = [var]

        else:

          for i, t in enumerate(var):

            variances[i].append(t)

      if self._prediction_outputs is not None:

        outputs = [outputs[i] for i in self._prediction_outputs]

        output_values = [output_values[i] for i in self._prediction_outputs]

      if len(transformers) > 0:

        if len(outputs) > 1:

        if len(output_values) > 1:

          raise ValueError(

              "predict() does not support Transformers for models with multiple outputs."

          )

        elif len(outputs) == 1:

          outputs = [undo_transforms(outputs[0], transformers)]

        elif len(output_values) == 1:

          output_values = [undo_transforms(output_values[0], transformers)]

      if results is None:

        results = [outputs]

        results = [output_values]

      else:

        for i, t in enumerate(outputs):

        for i, t in enumerate(output_values):

          results[i].append(t)

    # Concatenate arrays to create the final results.

    else:

      return final_results

  def predict_on_generator(self, generator, transformers=[]):

  def predict_on_generator(self, generator, transformers=[], outputs=None):

    """

    Parameters

    ----------

    transformers: list of dc.trans.Transformers

      Transformers that the input data has been transformed by.  The output

      is passed through these transformers to undo the transformations.

    outputs: Tensor or list of Tensors

      The outputs to return.  If this is None, the model's standard prediction

      outputs will be returned.  Alternatively one or more Tensors within the

      model may be specified, in which case the output of those Tensors will be

      returned.

    Returns:

      a NumPy array of the model produces a single output, or a list of arrays

      if it produces multiple outputs

    """

    return self._predict(generator, transformers, False)

    return self._predict(generator, transformers, outputs, False)

  def predict_on_batch(self, X, transformers=[]):

  def predict_on_batch(self, X, transformers=[], outputs=None):

    """Generates predictions for input samples, processing samples in a batch.

    Parameters

    transformers: list of dc.trans.Transformers

      Transformers that the input data has been transformed by.  The output

      is passed through these transformers to undo the transformations.

    outputs: Tensor or list of Tensors

      The outputs to return.  If this is None, the model's standard prediction

      outputs will be returned.  Alternatively one or more Tensors within the

      model may be specified, in which case the output of those Tensors will be

      returned.

    Returns

    -------

    if it produces multiple outputs

    """

    dataset = NumpyDataset(X=X, y=None)

    return self.predict(dataset, transformers)

    return self.predict(dataset, transformers, outputs)

  def predict_uncertainty_on_batch(self, X, masks=50):

    """

    dataset = NumpyDataset(X=X, y=None)

    return self.predict_uncertainty(dataset, masks)

  def predict(self, dataset, transformers=[]):

  def predict(self, dataset, transformers=[], outputs=None):

    """

    Uses self to make predictions on provided Dataset object.

    transformers: list of dc.trans.Transformers

      Transformers that the input data has been transformed by.  The output

      is passed through these transformers to undo the transformations.

    outputs: Tensor or list of Tensors

      The outputs to return.  If this is None, the model's standard prediction

      outputs will be returned.  Alternatively one or more Tensors within the

      model may be specified, in which case the output of those Tensors will be

      returned.

    Returns

    -------

    a NumPy array of the model produces a single output, or a list of arrays

    if it produces multiple outputs

    """

    generator = self.default_generator(dataset, predict=True, pad_batches=False)

    return self.predict_on_generator(generator, transformers)

    generator = self.default_generator(

        dataset, mode='predict', pad_batches=False)

    return self.predict_on_generator(generator, transformers, outputs)

  def predict_uncertainty(self, dataset, masks=50):

    """

    sum_var = []

    for i in range(masks):

      generator = self.default_generator(

          dataset, predict=True, pad_batches=False)

      results = self._predict(generator, [], True)

          dataset, mode='uncertainty', pad_batches=False)

      results = self._predict(generator, [], None, True)

      if len(sum_pred) == 0:

        for p, v in results:

          sum_pred.append(p)

  def default_generator(self,

                        dataset,

                        epochs=1,

                        predict=False,

                        mode='fit',

                        deterministic=True,

                        pad_batches=True):

    """Create a generator that iterates batches for a dataset.

    Subclasses may override this method to customize how model inputs are

    generated from the data.

    Parameters

    ----------

    dataset: Dataset

      the data to iterate

    epochs: int

      the number of times to iterate over the full dataset

    mode: str

      allowed values are 'fit' (called during training), 'predict' (called

      during prediction), and 'uncertainty' (called during uncertainty

      prediction)

    deterministic: bool

      whether to iterate over the dataset in order, or randomly shuffle the

      data for each epoch

    pad_batches: bool

      whether to pad each batch up to this model's preferred batch size

    Returns

    -------

    a generator that iterates batches, each represented as a tuple of lists:

    ([inputs], [outputs], [weights])

    """

    for epoch in range(epochs):

      for (X_b, y_b, w_b, ids_b) in dataset.iterbatches(

          batch_size=self.batch_size,

    return [x + p, xp + q]

class SwitchedDropout(tf.keras.layers.Layer):

  """Apply dropout based on an input.

  This is required for uncertainty prediction.  The standard Keras Dropout

  layer only performs dropout during training, but we sometimes need to do it

  during prediction.  The second input to this layer should be a scalar equal to

  0 or 1, indicating whether to perform dropout.

  """

  def __init__(self, rate, **kwargs):

    self.rate = rate

    super(SwitchedDropout, self).__init__(**kwargs)

  def call(self, inputs):

    rate = self.rate * tf.squeeze(inputs[1])

    return tf.nn.dropout(inputs[0], rate=rate)

class WeightedLinearCombo(tf.keras.layers.Layer):

  """Computes a weighted linear combination of input layers, with the weights defined by trainable variables."""

  def call(self, inputs):

    """In layers should be of shape dtype tf.float32, (None, self.max_atoms, 4)"""

    atom_numbers = tf.cast(inputs[0][:, :, 0], tf.int32)

    atom_numbers = tf.cast(inputs[:, :, 0], tf.int32)

    flags = tf.sign(atom_numbers)

    flags = tf.cast(

        tf.expand_dims(flags, 1) * tf.expand_dims(flags, 2), tf.float32)

    coordinates = inputs[0][:, :, 1:]

    coordinates = inputs[:, :, 1:]

    if self.coordinates_in_bohr:

      coordinates = coordinates * 0.52917721092

    self.weights_initializer = weights_initializer

  def build(self, input_shape):

    out_channels = input_shape[0][1]

    if isinstance(input_shape, collections.Sequence):

      input_shape = input_shape[0]

    out_channels = input_shape[1]

    if self.biases_initializer is None:

      biases_initializer = None

    else:

    self.built = True

  def call(self, inputs):

    if isinstance(inputs, collections.Sequence):

      parent = inputs[0]

    else:

      parent = inputs

    dense_H = self.dense_H(parent)

    dense_T = self.dense_T(parent)

    return tf.multiply(dense_H, dense_T) + tf.multiply(parent, 1 - dense_T)

    """

    parent layers: atom_number

    """

    atom_number = inputs[0]

    atom_number = inputs

    return tf.nn.embedding_lookup(self.embedding_list, atom_number)

    return tf.segment_sum(output, atom_membership)

def _DAGgraph_step(batch_inputs, W_list, b_list, activation, dropout, training):

def _DAGgraph_step(batch_inputs, W_list, b_list, activation, dropout,

                   dropout_switch):

  outputs = batch_inputs

  for idw, W in enumerate(W_list):

    outputs = tf.nn.xw_plus_b(outputs, W, b_list[idw])

    outputs = activation(outputs)

    rate_scale = 1.0 if training else 0.0

    if not dropout is None:

      outputs = tf.nn.dropout(outputs, rate=dropout * rate_scale)

      outputs = tf.nn.dropout(outputs, rate=dropout * dropout_switch)

  return outputs

    ]))

    self.built = True

  def call(self, inputs, training=None):

  def call(self, inputs):

    """

    parent layers: atom_features, parents, calculation_orders, calculation_masks, n_atoms

    """

    calculation_orders = inputs[2]

    calculation_masks = inputs[3]

    n_atoms = inputs[4]

    n_atoms = tf.squeeze(inputs[4])

    dropout_switch = tf.squeeze(inputs[5])

    # initialize graph features for each graph

    graph_features_initial = tf.zeros((self.max_atoms * self.batch_size,

                                       self.max_atoms + 1, self.n_graph_feat))

      # of shape: (batch_size*max_atoms) * n_graph_features

      # representing the graph features of target atoms in each graph

      batch_outputs = _DAGgraph_step(batch_inputs, self.W_list, self.b_list,

                                     self.activation, self.dropout, training)

                                     self.activation, self.dropout,

                                     dropout_switch)

      # index for targe atoms

      target_index = tf.stack([tf.range(n_atoms), parents[:, count, 0]], axis=1)

    ]))

    self.built = True

  def call(self, inputs, training=None):

  def call(self, inputs):

    """

    parent layers: atom_features, membership

    """

    atom_features = inputs[0]

    membership = inputs[1]

    dropout_switch = tf.squeeze(inputs[2])

    # Extract atom_features

    graph_features = tf.segment_sum(atom_features, membership)

    # sum all graph outputs

    return _DAGgraph_step(graph_features, self.W_list, self.b_list,

                          self.activation, self.dropout, training)

                          self.activation, self.dropout, dropout_switch)

class MessagePassing(tf.keras.layers.Layer):

import deepchem as dc

from deepchem.models import KerasModel

from deepchem.models.layers import SwitchedDropout

from deepchem.utils.save import log

from deepchem.metrics import to_one_hot, from_one_hot

from tensorflow.keras.layers import Input, Dense, Reshape, Softmax, Dropout, Activation

logger = logging.getLogger(__name__)

    # Add the input features.

    mol_features = tf.keras.Input(shape=(n_features,))

    mol_features = Input(shape=(n_features,))

    prev_layer = mol_features

    # Add the dense layers

    for size, weight_stddev, bias_const, dropout, activation_fn in zip(

        layer_sizes, weight_init_stddevs, bias_init_consts, dropouts,

        activation_fns):

      layer = tf.keras.layers.Dense(

      layer = Dense(

          size,

          activation=activation_fn,

          kernel_initializer=tf.truncated_normal_initializer(

          bias_initializer=tf.constant_initializer(value=bias_const),

          kernel_regularizer=regularizer)(prev_layer)

      if dropout > 0.0:

        layer = tf.keras.layers.Dropout(rate=dropout)(layer)

        layer = Dropout(rate=dropout)(layer)

      prev_layer = layer

    self.neural_fingerprint = prev_layer

    logits = tf.keras.layers.Reshape((n_tasks, n_classes))(

        tf.keras.layers.Dense(n_tasks * n_classes)(prev_layer))

    output = tf.keras.layers.Softmax()(logits)

    logits = Reshape((n_tasks,

                      n_classes))(Dense(n_tasks * n_classes)(prev_layer))

    output = Softmax()(logits)

    model = tf.keras.Model(inputs=mol_features, outputs=[output, logits])

    super(MultitaskClassifier, self).__init__(

        model,

  def default_generator(self,

                        dataset,

                        epochs=1,

                        predict=False,

                        mode='fit',

                        deterministic=True,

                        pad_batches=True):

    for epoch in range(epochs):

    # Add the input features.

    mol_features = tf.keras.Input(shape=(n_features,))

    mol_features = Input(shape=(n_features,))

    dropout_switch = Input(shape=tuple())

    prev_layer = mol_features

    # Add the dense layers

    for size, weight_stddev, bias_const, dropout, activation_fn in zip(

        layer_sizes, weight_init_stddevs, bias_init_consts, dropouts,

        activation_fns):

      layer = tf.keras.layers.Dense(

      layer = Dense(

          size,

          activation=activation_fn,

          kernel_initializer=tf.truncated_normal_initializer(

          bias_initializer=tf.constant_initializer(value=bias_const),

          kernel_regularizer=regularizer)(prev_layer)

      if dropout > 0.0:

        layer = tf.keras.layers.Dropout(rate=dropout)(layer)

        layer = SwitchedDropout(rate=dropout)([layer, dropout_switch])

      prev_layer = layer

    self.neural_fingerprint = prev_layer

    output = tf.keras.layers.Reshape((n_tasks, 1))(tf.keras.layers.Dense(

    output = Reshape((n_tasks, 1))(Dense(

        n_tasks,

        kernel_initializer=tf.truncated_normal_initializer(

            stddev=weight_init_stddevs[-1]),

        bias_initializer=tf.constant_initializer(

            value=bias_init_consts[-1]))(prev_layer))

    if uncertainty:

      log_var = tf.keras.layers.Reshape((n_tasks, 1))(tf.keras.layers.Dense(

      log_var = Reshape((n_tasks, 1))(Dense(

          n_tasks,

          kernel_initializer=tf.truncated_normal_initializer(

              stddev=weight_init_stddevs[-1]),

          bias_initializer=tf.constant_initializer(value=0.0))(prev_layer))

      var = tf.keras.layers.Activation(tf.exp)(log_var)

      var = Activation(tf.exp)(log_var)

      outputs = [output, var, output, log_var]

      output_types = ['prediction', 'variance', 'loss', 'loss']

      outputs = [output]

      output_types = ['prediction']

      loss = dc.models.losses.L2Loss()

    model = tf.keras.Model(inputs=mol_features, outputs=outputs)

    model = tf.keras.Model(

        inputs=[mol_features, dropout_switch], outputs=outputs)

    super(MultitaskRegressor, self).__init__(

        model, loss, output_types=output_types, **kwargs)

  def default_generator(self,

                        dataset,

                        epochs=1,

                        mode='fit',

                        deterministic=True,

                        pad_batches=True):

    for epoch in range(epochs):

      for (X_b, y_b, w_b, ids_b) in dataset.iterbatches(

          batch_size=self.batch_size,

          deterministic=deterministic,

          pad_batches=pad_batches):

        if mode == 'predict':

          dropout = np.array(0.0)

        else:

          dropout = np.array(1.0)

        yield ([X_b, dropout], [y_b], [w_b])

class MultitaskFitTransformRegressor(MultitaskRegressor):

  """Implements a MultitaskRegressor that performs on-the-fly transformation during fit/predict.

  def default_generator(self,

                        dataset,

                        epochs=1,

                        predict=False,

                        mode='fit',

                        deterministic=True,

                        pad_batches=True):

    for epoch in range(epochs):

        if y_b is not None:

          y_b = y_b.reshape(-1, self.n_tasks, 1)

        if X_b is not None:

          if not predict:

          if mode == 'fit':

            for transformer in self.fit_transformers:

              X_b = transformer.X_transform(X_b)

        yield ([X_b], [y_b], [w_b])

  def predict_on_generator(self, generator, transformers=[]):

  def predict_on_generator(self, generator, transformers=[], outputs=None):

    def transform_generator():

      for inputs, labels, weights in generator:

          yield ([X_t], labels, weights)

    return super(MultitaskFitTransformRegressor, self).predict_on_generator(

        transform_generator(), transformers)