Cleanup

  missing = np.zeros_like(y).astype(int)

  feature_shape = None

  #for ind in range(n_samples):

  #  for task in range(n_tasks):

  #    if y[ind, task] == "":

  #      missing[ind, task] = 1

  #x_list = list(df[feature_type].values)

  #valid_inds = np.array([1 if elt.size > 0 else 0 for elt in x_list], dtype=bool)

  #x_list = [elt for (is_valid, elt) in zip(valid_inds, x_list) if is_valid]

  #x = np.squeeze(np.array(x_list))

  sorted_ids = df[id_field].values

  # Set missing data to have weight zero

  for ind in range(n_samples):

        y[ind, task] = 0.

        w[ind, task] = 0.

  #sorted_ids = sorted_ids[valid_inds]

  #y = y[valid_inds]

  #w = w[valid_inds]

  # Adding this assertion in to avoid ill-formed outputs.

  #assert len(sorted_ids) == len(x) == len(y) == len(w)

  assert len(sorted_ids) == len(y) == len(w)

  #return sorted_ids, x, y.astype(float), w.astype(float)

  return sorted_ids, y.astype(float), w.astype(float)

#def write_dataframe(basename, df, data_dir, feature_type, tasks=None,

#                    raw_data=None, mol_id_field="mol_id",

#                    verbose=False):

#  """Writes data from dataframe to disk."""

#  if featurizer is not None and tasks is not None:

#    feature_type = featurizer.__class__.__name__

#    time1 = time.time()

#    ids, X, y, w = convert_df_to_numpy(

#        df, feature_type, tasks, mol_id_field)

#    time2 = time.time()

#    log("TIMING: convert_df_to_numpy took %0.3f s" % (time2-time1), verbose)

#  else:

#    ids, X, y, w = raw_data

#    assert X.shape[0] == y.shape[0]

#    assert y.shape == w.shape

#    assert len(ids) == X.shape[0]

#  return DiskDataset.write_data_to_disk(

#      data_dir, basename, tasks, X, y, w, ids)

def featurize_smiles_df(df, featurizer, field, log_every_N=1000, verbose=True):

  """Featurize individual compounds in dataframe.

        time2 = time.time()

        log("TIMING: featurizing shard %d took %0.3f s" % (shard_num, time2-time1),

            self.verbose)

        yield ids, X, y, w

        yield X, y, w, ids

    return DiskDataset(shard_generator(), data_dir, self.tasks)

  def get_shards(self, input_files, shard_size):

    """Featurizes a shard of an input dataframe."""

    raise NotImplementedError

class CSVLoader(DataLoader):

  """

  Handles loading of CSV files.

        % self.featurizer.__class__.__name__, self.verbose)

    return featurize_smiles_df(shard, self.featurizer,

                               field=self.smiles_field)

class UserCSVLoader(DataLoader):

  """

  Handles loading of CSV files with user-defined featurizers.

    metadata_rows = []

    time1 = time.time()

    for shard_num, (ids, X, y, w) in enumerate(shard_generator):

    for shard_num, (X, y, w, ids) in enumerate(shard_generator):

      basename = "shard-%d" % shard_num 

      metadata_rows.append(

          DiskDataset.write_data_to_disk(

      for shard_num, row in self.metadata_df.iterrows():

        X, y, w, ids = self.get_shard(shard_num)

        newx, newy, neww = fn(X, y, w)

        yield (ids, X, y, w)

        yield (X, y, w, ids)

    return DiskDataset(generator(), data_dir=out_dir)

  @staticmethod

      w = np.ones_like(y)

    if tasks is None:

      tasks = np.arange(n_tasks)

    raw_data = (ids, X, y, w)

    raw_data = (X, y, w, ids)

    return DiskDataset(data_dir=data_dir, tasks=tasks, raw_data=raw_data)

  @staticmethod

    def generator():

      for ind, dataset in enumerate(datasets):

        X, y, w, ids = (dataset.X, dataset.y, dataset.w, dataset.ids)

        yield (ids, X, y, w)

        yield (X, y, w, ids)

    return DiskDataset(generator(), data_dir=merge_dir)

  def subset(self, shard_nums, subset_dir=None):

        if shard_num not in shard_nums:

          continue

        X, y, w, ids = self.get_shard(shard_num)

        yield (ids, X, y, w)

        yield (X, y, w, ids)

    return DiskDataset(generator(), data_dir=subset_dir)

  def sparse_shuffle(self):

      return self.eval_graph.session

  def fit_task(self, sess, dataset, task, task_train_op, nb_epoch=10, pad_batches=False,

               log_every_N_batches=50):

    """Fit the model.

    Fit one task.

    TODO(rbharath): Figure out if the logging will work correctly with the

    global_step set as it is.

    Parameters

    ---------- 

    dataset: dc.data.Dataset

      Dataset object holding training data 

    task: int

      The index of the task to train on.

    nb_epoch: 10

      Number of training epochs.

    pad_batches: bool

      Whether or not to pad each batch to exactly be of size batch_size.

    max_checkpoints_to_keep: int

      Maximum number of checkpoints to keep; older checkpoints will be deleted.

    log_every_N_batches: int

      Report every N batches. Useful for training on very large datasets,

      where epochs can take long time to finish.

    Raises

    ------

    AssertionError

      If model is not in training mode.

    """

    ############################################################## TIMING

    time1 = time.time()

    ############################################################## TIMING

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

    for epoch in range(nb_epoch):

      avg_loss, n_batches = 0., 0

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

          # Turns out there are valid cases where we don't want pad-batches

          # on by default.

          #dataset.iterbatches(batch_size, pad_batches=True)):

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

        if ind % log_every_N_batches == 0:

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

        feed_dict = self.construct_task_feed_dict(task, X_b, y_b, w_b, ids_b)

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

        fetches = self.train_graph.output + [

            task_train_op, self.train_graph.loss[task]]

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

        fetched_values = sess.run(fetches, feed_dict=feed_dict)

        output = fetched_values[:len(self.train_graph.output)]

        loss = fetched_values[-1]

        avg_loss += loss

        y_pred = np.squeeze(np.array(output))

        y_b = y_b.flatten()

        n_batches += 1

      #saver.save(sess, self._save_path, global_step=epoch)

      avg_loss = float(avg_loss)/n_batches

      log('Ending epoch %d: Average loss %g' % (epoch, avg_loss), self.verbose)

    ############################################################## TIMING

    time2 = time.time()

    print("TIMING: model fitting took %0.3f s" % (time2-time1),

          self.verbose)

    ############################################################## TIMING

class ProgressiveMultitaskClassifier(TensorflowMultiTaskClassifier):

  """Implements a progressive multitask neural network.

  Progressive Networks: https://arxiv.org/pdf/1606.04671v3.pdf

  Progressive networks allow for multitask learning where each task

  gets a new column of weights. As a result, there is no exponential

  forgetting where previous tasks are ignored.

  TODO(rbharath): This code has a lot of overlap with regressor code above.

  Refactor code together.

  """

  def __init__(self, n_tasks, n_features, alpha_init_stddevs=[.02], **kwargs):

    """Creates a progressive network.

    Only listing parameters specific to progressive networks here.

    Parameters

    ----------

    n_tasks: int

      Number of tasks

    n_features: int

      Number of input features

    alpha_init_stddevs: list

      List of standard-deviations for alpha in adapter layers.

    """

    self.alpha_init_stddevs = alpha_init_stddevs

    super(ProgressiveMultitaskClassifier, self).__init__(

        n_tasks, n_features, **kwargs)

    # Consistency check

    lengths_set = {

        len(self.layer_sizes),

        len(self.weight_init_stddevs),

        len(self.alpha_init_stddevs),

        len(self.bias_init_consts),

        len(self.dropouts),

        }

    assert len(lengths_set) == 1, "All layer params must have same length."

  def construct_graph(self, training, seed):

    """Returns a TensorflowGraph object."""

    graph = tf.Graph() 

    # Lazily created by _get_shared_session().

    shared_session = None

    # Cache of TensorFlow scopes, to prevent '_1' appended scope names

    # when subclass-overridden methods use the same scopes.

    name_scopes = {}

    # Setup graph

    with graph.as_default():

      if seed is not None:

        tf.set_random_seed(seed)

      features, labels, weights = self.add_placeholders(graph, name_scopes)

      outputs = self.add_progressive_lattice(graph, name_scopes, training)

      if training:

        loss = self.add_task_training_costs(graph, name_scopes, outputs, labels,

                                           weights)

      else:

        loss = None

    return TensorflowGraph(graph=graph,

                           session=shared_session,

                           name_scopes=name_scopes,

                           output=outputs,

                           labels=labels,

                           weights=weights,

                           loss=loss)

  def add_placeholders(self, graph, name_scopes):

    """Adds all placeholders for this model."""

    # Create placeholders

    placeholder_scope = TensorflowGraph.get_placeholder_scope(

        graph, name_scopes)

    labels, weights = [], []

    n_features = self.n_features

    with placeholder_scope:

      self.mol_features = tf.placeholder(

          tf.float32,

          shape=[None, n_features],

          name='mol_features')

      for task in range(self.n_tasks):

        weights.append(tf.identity(

            tf.placeholder(tf.float32, shape=[None,],

                           name='weights_%d' % task)))

        labels.append(tf.identity(

            tf.placeholder(tf.float32, shape=[None, self.n_classes],

                           name='labels_%d' % task)))

    return self.mol_features, labels, weights

  def add_progressive_lattice(self, graph, name_scopes, training):

    """Constructs the graph architecture as specified in its config.

    This method creates the following Placeholders:

      mol_features: Molecule descriptor (e.g. fingerprint) tensor with shape

        batch_size x n_features.

    """

    n_features = self.n_features

    placeholder_scope = TensorflowGraph.get_placeholder_scope(

        graph, name_scopes)

    with graph.as_default():

      layer_sizes = self.layer_sizes

      weight_init_stddevs = self.weight_init_stddevs

      bias_init_consts = self.bias_init_consts

      dropouts = self.dropouts

      lengths_set = {

          len(layer_sizes),

          len(weight_init_stddevs),

          len(bias_init_consts),

          len(dropouts),

          }

      assert len(lengths_set) == 1, 'All layer params must have same length.'

      n_layers = lengths_set.pop()

      assert n_layers > 0, 'Must have some layers defined.'

      prev_layer = self.mol_features

      prev_layer_size = n_features 

      all_layers = {}

      for i in range(n_layers):

        for task in range(self.n_tasks):

          task_scope = TensorflowGraph.shared_name_scope(

              "task%d_ops" % task, graph, name_scopes)

          print("Adding weights for task %d, layer %d" % (task, i))

          with task_scope as scope:

            if i == 0:

              prev_layer = self.mol_features

              prev_layer_size = self.n_features

            else:

              prev_layer = all_layers[(i-1, task)]

              prev_layer_size = layer_sizes[i-1]

              if task > 0:

                lateral_contrib = self.add_adapter(all_layers, task, i)

            print("Creating W_layer_%d_task%d of shape %s" %

                  (i, task, str([prev_layer_size, layer_sizes[i]])))

            W = tf.Variable(

                tf.truncated_normal(

                    shape=[prev_layer_size, layer_sizes[i]],

                    stddev=self.weight_init_stddevs[i]),

                name='W_layer_%d_task%d' % (i, task), dtype=tf.float32)

            print("Creating b_layer_%d_task%d of shape %s" %

                  (i, task, str([layer_sizes[i]])))

            b = tf.Variable(tf.constant(value=self.bias_init_consts[i],

                            shape=[layer_sizes[i]]),

                            name='b_layer_%d_task%d' % (i, task), dtype=tf.float32)

            layer = tf.matmul(prev_layer, W) + b

            if i > 0 and task > 0:

              layer = layer + lateral_contrib

            layer = tf.nn.relu(layer)

            layer = model_ops.dropout(layer, dropouts[i], training)

            all_layers[(i, task)] = layer

      output = []

      for task in range(self.n_tasks):

        prev_layer = all_layers[(i, task)]

        prev_layer_size = layer_sizes[i]

        task_scope = TensorflowGraph.shared_name_scope(

            "task%d" % task, graph, name_scopes)

        with task_scope as scope:

          if task > 0:

            lateral_contrib = tf.squeeze(self.add_adapter(all_layers, task, i+1))

          weight_init = tf.truncated_normal(

              shape=[prev_layer_size, 1],

              stddev=weight_init_stddevs[i])

          bias_init = tf.constant(value=bias_init_consts[i],

                                  shape=[1])

          print("Creating W_output_task%d of shape %s"

                % (task, str([prev_layer_size, 1])))

          w = tf.Variable(weight_init, name='W_output_task%d'%task,

                          dtype=tf.float32)

          print("Creating b_output_task%d of shape %s" % (task, str([1])))

          b = tf.Variable(bias_init, name='b_output_task%d'%task,

                          dtype=tf.float32)

          layer = tf.squeeze(tf.matmul(prev_layer, w) + b)

          if i > 0 and task > 0:

            layer = layer + lateral_contrib

          output.append(layer)

      # Add softmax ops at end 

      softmax = []

      with tf.name_scope('inference'):

        for i, logits in enumerate(output):

          softmax.append(tf.nn.softmax(logits, name='softmax_%d' % i))

      output = softmax

      return output

  def add_adapter(self, all_layers, task, layer_num):

    """Add an adapter connection for given task/layer combo"""

    i = layer_num

    prev_layers = []

    # Handle output layer

    if i < len(self.layer_sizes):

      layer_sizes = self.layer_sizes

      alpha_init_stddev = self.alpha_init_stddevs[i]

      weight_init_stddev = self.weight_init_stddevs[i]

      bias_init_const = self.bias_init_consts[i]

    elif i == len(self.layer_sizes):

      layer_sizes = self.layer_sizes + [1]

      alpha_init_stddev = self.alpha_init_stddevs[-1]

      weight_init_stddev = self.weight_init_stddevs[-1]

      bias_init_const = self.bias_init_consts[-1]

    else:

      raise ValueError("layer_num too large for add_adapter.")

    # Iterate over all previous tasks.

    for prev_task in range(task):

      prev_layers.append(all_layers[(i-1, prev_task)])

    # prev_layers is a list with elements of size

    # (batch_size, layer_sizes[i-1])

    prev_layer = tf.concat(1, prev_layers)

    alpha = tf.Variable(

        tf.truncated_normal([1,], stddev=alpha_init_stddev),

        name="alpha_layer_%d_task%d" % (i, task))

    prev_layer = tf.mul(alpha, prev_layer)

    prev_layer_size = task*layer_sizes[i-1]

    print("Creating V_layer_%d_task%d of shape %s" %

          (i, task, str([prev_layer_size, layer_sizes[i-1]])))

    V = tf.Variable(

        tf.truncated_normal(

            shape=[prev_layer_size, layer_sizes[i-1]],

            stddev=weight_init_stddev),

        name="V_layer_%d_task%d" % (i, task), dtype=tf.float32)

    print("Creating b_lat_layer_%d_task%d of shape %s" %

          (i, task, str([layer_sizes[i-1]])))

    b_lat = tf.Variable(

        tf.constant(value=bias_init_const, shape=[layer_sizes[i-1]]),

        name='b_lat_layer_%d_task%d' % (i, task),

        dtype=tf.float32)

    prev_layer = tf.matmul(prev_layer, V) + b_lat

    print("Creating U_layer_%d_task%d of shape %s" %

          (i, task, str([layer_sizes[i-1], layer_sizes[i]])))

    U = tf.Variable(

        tf.truncated_normal(

            shape=[layer_sizes[i-1], layer_sizes[i]],

            stddev=weight_init_stddev),

        name="U_layer_%d_task%d" % (i, task), dtype=tf.float32)

    return tf.matmul(prev_layer, U)

  def get_training_op(self, graph, loss):

    """Get training op for applying gradients to variables.

    Subclasses that need to do anything fancy with gradients should override

    this method.

    Returns:

    A training op.

    """

    with graph.as_default():

      opt = model_ops.optimizer(self.optimizer, self.learning_rate, self.momentum)

      return opt.minimize(loss, name='train')

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

    """Construct a feed dictionary from minibatch data.

    TODO(rbharath): ids_b is not used here. Can we remove it?

    Args:

      X_b: np.ndarray of shape (batch_size, n_features)

      y_b: np.ndarray of shape (batch_size, n_tasks)

      w_b: np.ndarray of shape (batch_size, n_tasks)

      ids_b: List of length (batch_size) with datapoint identifiers.

    """ 

    orig_dict = {}

    orig_dict["mol_features"] = X_b

    for task in range(self.n_tasks):

      if y_b is not None:

        orig_dict["labels_%d" % task] = y_b[:, task]

      else:

        # Dummy placeholders

        orig_dict["labels_%d" % task] = np.squeeze(

            np.zeros((self.batch_size,)))

      if w_b is not None:

        orig_dict["weights_%d" % task] = w_b[:, task]

      else:

        # Dummy placeholders

        orig_dict["weights_%d" % task] = np.ones(

            (self.batch_size,)) 

    return TensorflowGraph.get_feed_dict(orig_dict)

  def fit(self, dataset, tasks=None, close_session=True,

          max_checkpoints_to_keep=5, **kwargs):

    """Fit the model.

    Progressive networks are fit by training one task at a time. Iteratively

    fits one task at a time with other weights frozen.

    Parameters

    ---------- 

    dataset: dc.data.Dataset

      Dataset object holding training data 

    Raises

    ------

    AssertionError

      If model is not in training mode.

    """

    if tasks is None:

      tasks = range(self.n_tasks)

    with self.train_graph.graph.as_default():

      task_train_ops = {}

      for task in tasks:

        task_train_ops[task] = self.get_task_training_op(

            self.train_graph.graph, self.train_graph.loss, task)

      sess = self._get_shared_session(train=True)

      #with self._get_shared_session(train=True) as sess:

      sess.run(tf.initialize_all_variables())

      # Save an initial checkpoint.

      saver = tf.train.Saver(max_to_keep=max_checkpoints_to_keep)

      saver.save(sess, self._save_path, global_step=0)

      for task in tasks:

        print("Fitting on task %d" % task)

        self.fit_task(sess, dataset, task, task_train_ops[task], **kwargs)

        saver.save(sess, self._save_path, global_step=task)

      # Always save a final checkpoint when complete.

      saver.save(sess, self._save_path, global_step=self.n_tasks)

      if close_session:

        sess.close()

  def get_task_training_op(self, graph, losses, task):

    """Get training op for applying gradients to variables.

    Subclasses that need to do anything fancy with gradients should override

    this method.

    Parameters

    ----------

    graph: tf.Graph

      Graph for this op

    losses: dict

      Dictionary mapping task to losses

    Returns

    -------

    A training op.

    """

    with graph.as_default():

      task_loss = losses[task]

      task_root = "task%d_ops" % task

      task_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, task_root)

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

      print("get_task_training_op for task %d" % task)

      print("len(task_vars)")

      print(len(task_vars))

      print("[task_var.name for task_var in task_vars]")

      print([task_var.name for task_var in task_vars])

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

      opt = model_ops.optimizer(self.optimizer, self.learning_rate, self.momentum)

      return opt.minimize(task_loss, name='train', var_list=task_vars)

  def add_task_training_costs(self, graph, name_scopes, outputs, labels, weights):

    """Adds the training costs for each task.

    Since each task is trained separately, each task is optimized w.r.t a separate

    task.

    TODO(rbharath): Figure out how to support weight decay for this model.

    Since each task is trained separately, weight decay should only be used

    on weights in column for that task.

    Parameters

    ----------

    graph: tf.Graph

      Graph for the model.

    name_scopes: dict

      Contains all the scopes for model

    outputs: list

      List of output tensors from model.

    weights: list

      List of weight placeholders for model.

    """

    task_costs = {}

    with TensorflowGraph.shared_name_scope('costs', graph, name_scopes):

      for task in range(self.n_tasks):

        with TensorflowGraph.shared_name_scope(

            'cost_%d' % task, graph, name_scopes):

          weighted_cost = self.cost(outputs[task], labels[task],

                                    weights[task])

          # Note that we divide by the batch size and not the number of

          # non-zero weight examples in the batch.  Also, instead of using

          # tf.reduce_mean (which can put ops on the CPU) we explicitly

          # calculate with div/sum so it stays on the GPU.

          task_cost = tf.div(tf.reduce_sum(weighted_cost), self.batch_size)

          task_costs[task] = task_cost 

    return task_costs

  def construct_task_feed_dict(self, this_task, X_b, y_b=None, w_b=None, ids_b=None):

    """Construct a feed dictionary from minibatch data.

    TODO(rbharath): ids_b is not used here. Can we remove it?

    Args:

      X_b: np.ndarray of shape (batch_size, n_features)

      y_b: np.ndarray of shape (batch_size, n_tasks)

      w_b: np.ndarray of shape (batch_size, n_tasks)

      ids_b: List of length (batch_size) with datapoint identifiers.

    """ 

    orig_dict = {}

    orig_dict["mol_features"] = X_b

    n_samples = len(X_b)

    for task in range(self.n_tasks):

      if (this_task == task) and y_b is not None:

        orig_dict["labels_%d" % task] = to_one_hot(y_b[:, task])

      else:

        # Dummy placeholders

        orig_dict["labels_%d" % task] = np.squeeze(to_one_hot(

            np.zeros((n_samples,))))

      if (this_task == task) and w_b is not None:

        orig_dict["weights_%d" % task] = np.reshape(w_b[:, task], (n_samples,))

      else:

        # Dummy placeholders

        orig_dict["weights_%d" % task] = np.zeros((n_samples,)) 

    return TensorflowGraph.get_feed_dict(orig_dict)

  def _get_shared_session(self, train):

    # allow_soft_placement=True allows ops without a GPU implementation

    # to run on the CPU instead.

    if train:

      if not self.train_graph.session:

        config = tf.ConfigProto(allow_soft_placement=True)

        self.train_graph.session = tf.Session(config=config)

      return self.train_graph.session

    else:

      if not self.eval_graph.session: