DTNN regressor

from deepchem.models.sklearn_models import SklearnModel

from deepchem.models.tf_new_models.multitask_classifier import MultitaskGraphClassifier

from deepchem.models.tf_new_models.multitask_regressor import MultitaskGraphRegressor

from deepchem.models.tf_new_models.DTNN_regressor import DTNNRegressor

from deepchem.models.tf_new_models.support_classifier import SupportGraphClassifier

from deepchem.models.multitask import SingletaskToMultitask

from deepchem.models.sequential import Sequential

#!/usr/bin/env python2

# -*- coding: utf-8 -*-

"""

Created on Mon Mar 13 14:42:40 2017

@author: zqwu

"""

import os

import sys

import numpy as np

import tensorflow as tf

import sklearn.metrics

import tempfile

from deepchem.utils.save import log

from deepchem.models import Model

from deepchem.nn.copy import Input

from deepchem.nn.copy import Dense

from deepchem.data import pad_features

from deepchem.nn import model_ops

# TODO(rbharath): Find a way to get rid of this import?

from deepchem.models.tf_new_models.graph_topology import merge_dicts

from deepchem.models.tf_new_models.multitask_classifier import get_loss_fn

from deepchem.models.tf_new_models.multitask_regressor import MultitaskGraphRegressor

class MultitaskGraphRegressor(Model):

class DTNNRegressor(MultitaskGraphRegressor):

  def __init__(self,

               model,

               n_tasks,

               n_feat,

               n_tasks=1,

               logdir=None,

               batch_size=50,

               final_loss='weighted_L2',

      self.pad_batches = pad_batches

      # Get graph topology for x

      self.graph_topology = self.model.get_graph_topology()

      self.feat_dim = n_feat

      # Building outputs

      self.outputs = self.build()

      self.add_optimizer()

      # Initialize

      self.init_fn = tf.initialize_all_variables()

      self.init_fn = tf.global_variables_initializer()

      self.sess.run(self.init_fn)

      # Path to save checkpoint files, which matches the

    self.weight_placeholder = tf.placeholder(

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

    feat = self.model.return_outputs()

    return feat

  def add_optimizer(self):

    if self.optimizer_type == "adam":

      self.optimizer = tf.train.AdamOptimizer(

          self.learning_rate,

          beta1=self.optimizer_beta1,

          beta2=self.optimizer_beta2,

          epsilon=self.epsilon)

    else:

      raise ValueError("Optimizer type not recognized.")

    # Get train function

    self.train_op = self.optimizer.minimize(self.loss_op)

  def construct_feed_dict(self, X_b, y_b=None, w_b=None, training=True):

    """Get initial information about task normalization"""

    # TODO(rbharath): I believe this is total amount of data

    n_samples = len(X_b)

    if y_b is None:

      y_b = np.zeros((n_samples, self.n_tasks))

    if w_b is None:

      w_b = np.zeros((n_samples, self.n_tasks))

    targets_dict = {self.label_placeholder: y_b, self.weight_placeholder: w_b}

    # Get graph information

    atoms_dict = self.graph_topology.batch_to_feed_dict(X_b)

    # TODO (hraut->rhbarath): num_datapoints should be a vector, with ith element being

    # the number of labeled data points in target_i. This is to normalize each task

    # num_dat_dict = {self.num_datapoints_placeholder : self.}

    # Get other optimizer information

    # TODO(rbharath): Figure out how to handle phase appropriately

    feed_dict = merge_dicts([targets_dict, atoms_dict])

    return feed_dict

  def add_training_loss(self, final_loss, outputs):

    """Computes loss using logits."""

    loss_fn = get_loss_fn(final_loss)  # Get loss function

    task_losses = []

    # label_placeholder of shape (batch_size, n_tasks). Split into n_tasks

    # tensors of shape (batch_size,)

    task_labels = tf.split(1, self.n_tasks, self.label_placeholder)

    task_weights = tf.split(1, self.n_tasks, self.weight_placeholder)

    for task in range(self.n_tasks):

      task_label_vector = task_labels[task]

      task_weight_vector = task_weights[task]

      task_loss = loss_fn(outputs[task],

                          tf.squeeze(task_label_vector),

                          tf.squeeze(task_weight_vector))

      task_losses.append(task_loss)

    # It's ok to divide by just the batch_size rather than the number of nonzero

    # examples (effect averages out)

    total_loss = tf.add_n(task_losses)

    total_loss = tf.div(total_loss, self.batch_size)

    return total_loss

  def fit(self,

          dataset,

          nb_epoch=10,

          max_checkpoints_to_keep=5,

          log_every_N_batches=50,

          checkpoint_interval=10,

          **kwargs):

    # Perform the optimization

    log("Training for %d epochs" % nb_epoch, self.verbose)

    # TODO(rbharath): Disabling saving for now to try to debug.

    for epoch in range(nb_epoch):

      log("Starting epoch %d" % epoch, self.verbose)

      for batch_num, (X_b, y_b, w_b, ids_b) in enumerate(

          dataset.iterbatches(self.batch_size, pad_batches=self.pad_batches)):

        if batch_num % log_every_N_batches == 0:

          log("On batch %d" % batch_num, self.verbose)

        self.sess.run(

            self.train_op, feed_dict=self.construct_feed_dict(X_b, y_b, w_b))

  def save(self):

    """

    No-op since this model doesn't currently support saving... 

    """

    pass

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

    """Wraps predict to set batch_size/padding."""

    return super(MultitaskGraphRegressor, self).predict(

        dataset, transformers, batch_size=self.batch_size)

  def predict_on_batch(self, X):

    """Return model output for the provided input.

    """

    if self.pad_batches:

      X = pad_features(self.batch_size, X)

    # run eval data through the model

    n_tasks = self.n_tasks

    with self.sess.as_default():

      feed_dict = self.construct_feed_dict(X)

      # Shape (n_samples, n_tasks)

      batch_outputs = self.sess.run(self.outputs, feed_dict=feed_dict)

    n_samples = len(X)

    outputs = np.zeros((n_samples, self.n_tasks))

    for task, output in enumerate(batch_outputs):

      outputs[:, task] = output

    return outputs

  def get_num_tasks(self):

    """Needed to use Model.predict() from superclass."""

    return self.n_tasks

    outputs = self.model.return_outputs()

    return [outputs]

  def get_layer(self, layer_id):

    return self.layers[layer_id]

class SequentialDTNNGraph(object):

class SequentialDTNNGraph(SequentialGraph):

  """An analog of Keras Sequential class for Graph data.

  Like the Sequential class from Keras, but automatically passes topology

  to the network. Non graph layers don't get the extra placeholders. 

  """

  def __init__(self, max_n_atoms, n_distance):

  def __init__(self, max_n_atoms=30, n_distance=100):

    """

    Parameters

    ----------

    """

    self.graph = tf.Graph()

    with self.graph.as_default():

      self.graph_DTNN_topology = DTNNGraphTopology(max_n_atoms, n_distance)

      self.graph_topology = DTNNGraphTopology(max_n_atoms, n_distance)

      self.output = self.graph_topology.get_atom_number_placeholder()

    # Keep track of the layers

    self.layers = []

      # For graphical layers, add connectivity placeholders 

      if type(layer).__name__ in ['DTNNStep']:

        self.output = layer([self.output] +

                            self.graph_DTNN_topology.get_topology_placeholders())

                            self.graph_topology.get_topology_placeholders())

      else:

        self.output = layer(self.output)

      ############################################# DEBUG

      # Add layer to the layer list

      self.layers.append(layer)

  def get_graph_topology(self):

    return self.graph_topology

  def get_num_output_features(self):

    """Gets the output shape of the featurization layers of the network"""

    return self.layers[-1].output_shape[1]

  def return_outputs(self):

    return self.output

  def return_inputs(self):

    return self.graph_topology.get_input_placeholders()

    return self.graph_topology.get_atom_number_placeholders()

  def get_layer(self, layer_id):

    return self.layers[layer_id]

    }

    return merge_dicts([atoms_dict, deg_adj_dict])

class DTNNGraphTopology(object):

class DTNNGraphTopology(GraphTopology):

  """Manages placeholders associated with batch of graphs and their topology"""

  def __init__(self, max_n_atoms=30, n_distance=100, name='DTNN_topology'):

    self.inputs = [self.atom_number_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_number_placeholder(self):

    return self.atom_number_placeholder

  if final_loss == 'L2':

    def loss_fn(x, t):

      diff = tf.sub(x, t)

      diff = tf.substract(x, t)

      return tf.reduce_sum(tf.square(diff), 0)

  elif final_loss == 'weighted_L2':

    def loss_fn(x, t, w):

      diff = tf.sub(x, t)

      weighted_diff = tf.mul(diff, w)

      diff = tf.substract(x, t)

      weighted_diff = tf.multiply(diff, w)

      return tf.reduce_sum(tf.square(weighted_diff), 0)

  elif final_loss == 'L1':

    def loss_fn(x, t):

      diff = tf.sub(x, t)

      diff = tf.substract(x, t)

      return tf.reduce_sum(tf.abs(diff), 0)

  elif final_loss == 'huber':

    def loss_fn(x, t):

      diff = tf.sub(x, t)

      diff = tf.substract(x, t)

      return tf.reduce_sum(

          tf.minimum(0.5 * tf.square(diff),

                     huber_d * (tf.abs(diff) - 0.5 * huber_d)), 0)

  elif final_loss == 'cross_entropy':

    def loss_fn(x, t, w):

      costs = tf.nn.sigmoid_cross_entropy_with_logits(x, t)

      weighted_costs = tf.mul(costs, w)

      costs = tf.nn.sigmoid_cross_entropy_with_logits(logits=x, labels=t)

      weighted_costs = tf.multiply(costs, w)

      return tf.reduce_sum(weighted_costs)

  elif final_loss == 'hinge':

    def loss_fn(x, t, w):

      t = tf.mul(2.0, t) - 1

      costs = tf.maximum(0.0, 1.0 - tf.mul(t, x))

      weighted_costs = tf.mul(costs, w)

      t = tf.multiply(2.0, t) - 1

      costs = tf.maximum(0.0, 1.0 - tf.multiply(t, x))

      weighted_costs = tf.multiply(costs, w)

      return tf.reduce_sum(weighted_costs)

  return loss_fn

    task_losses = []

    # label_placeholder of shape (batch_size, n_tasks). Split into n_tasks

    # tensors of shape (batch_size,)

    task_labels = tf.split(1, self.n_tasks, self.label_placeholder)

    task_weights = tf.split(1, self.n_tasks, self.weight_placeholder)

    task_labels = tf.split(

        axis=1, num_or_size_splits=self.n_tasks, value=self.label_placeholder)

    task_weights = tf.split(

        axis=1, num_or_size_splits=self.n_tasks, value=self.weight_placeholder)

    for task in range(self.n_tasks):

      task_label_vector = task_labels[task]

      task_weight_vector = task_weights[task]