Created OntologyModel

"""A Model whose structure is defined by an ontology."""

from deepchem.models import TensorGraph

from deepchem.models.tensorgraph import layers

import tensorflow as tf

class OntologyNode(object):

  def __init__(self, id, n_outputs, feature_ids=[], children=[]):

    self.id = id

    self.n_outputs = n_outputs

    self.feature_ids = feature_ids

    self.children = children

class OntologyModel(TensorGraph):

  """Implements ontology based models.

  The model is based on Ma et al., "Using deep learning to model the hierarchical

  structure and function of a cell" (https://doi.org/10.1038/nmeth.4627).  The

  model structure is defined by an ontology: a set of features grouped into

  categories, which in turn are arranged hierarchically to form a directed

  acyclic graph.  An example is the Gene Ontology (GO) classifications which

  groups genes into a set of hierarchical categories based on their biological

  role.  Using a known ontology to define the model structure has two benefits.

  First, incorporating prior knowledge can sometimes lead to much more accurate

  predictions for a fixed model size.  Second, it makes the model's results

  much easier to interpret.

  """

  def __init__(self, n_tasks, feature_ids, root_node, mode="regression", n_classes=2, intermediate_loss_weight=0.3, weight_decay_penalty=0.0, **kwargs):

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

    self.n_tasks = n_tasks

    self.feature_ids = feature_ids

    self.mode = mode

    self.n_classes = n_classes

    self._feature_index = dict((f, i) for i,f in enumerate(feature_ids))

    self._features = layers.Transpose((1, 0), in_layers=layers.Feature(shape=(None, len(feature_ids))))

    self.output_for_node = {}

    self.prediction_for_node = {}

    if mode not in ('regression', 'classification'):

      raise ValueError('Mode must be "regression" or "classification"')

    # Construct layers for all nodes.

    logits_for_node = {}

    self._build_layers(root_node)

    for id in self.output_for_node:

      if mode == 'regression':

        prediction = layers.Dense(in_layers=self.output_for_node[id], out_channels=n_tasks)

      else:

        logits = layers.Reshape(shape=(-1, n_tasks, n_classes), in_layers=layers.Dense(in_layers=self.output_for_node[id], out_channels=n_tasks*n_classes))

        prediction = layers.SoftMax(logits)

        logits_for_node[id] = logits

      self.prediction_for_node[id] = prediction

    self.add_output(self.prediction_for_node[root_node.id])

    # Create the loss function.

    losses = []

    loss_weights = []

    weights = layers.Weights(shape=(None, n_tasks))

    if mode == 'regression':

      labels = layers.Label(shape=(None, n_tasks))

      for id in self.prediction_for_node:

        losses.append(layers.ReduceSum(layers.L2Loss([labels, self.prediction_for_node[id], weights])))

        loss_weights.append(1.0 if id == root_node.id else intermediate_loss_weight)

    else:

      labels = layers.Label(shape=(None, n_tasks, n_classes))

      for id in self.prediction_for_node:

        losses.append(layers.WeightedError([layers.SoftMaxCrossEntropy([labels, logits_for_node[id]]), weights]))

        loss_weights.append(1.0 if id == root_node.id else intermediate_loss_weight)

    loss = layers.Add(in_layers=losses, weights=loss_weights)

    if weight_decay_penalty != 0.0:

      loss = layers.WeightDecay(weight_decay_penalty, 'l2', in_layers=loss)

    self.set_loss(loss)

  def _build_layers(self, node):

    inputs = []

    # Create inputs for the features.

    if len(node.feature_ids) > 0:

      indices = []

      for f in node.feature_ids:

        if f in self._feature_index:

          indices.append([self._feature_index[f]])

        else:

          raise ValueError('Unknown feature "%s"' % f)

      inputs.append(layers.Transpose((1, 0), in_layers=layers.Gather(in_layers=self._features, indices=indices)))

    # Create inputs for the children.

    if len(node.children) > 0:

      for child in node.children:

        if child.id not in self.output_for_node:

          self._build_layers(child)

        inputs.append(self.output_for_node[child.id])

    # Concatenate all inputs together.

    if len(inputs) == 0:

      raise ValueError('OntologyNode must have at least one child or feature')

    if len(inputs) == 1:

      inputs = inputs[0]

    else:

      inputs = layers.Concat(inputs)

    # Create the output.

    dense = layers.Dense(node.n_outputs, in_layers=inputs, activation_fn=tf.tanh)

    output = layers.BatchNorm(dense)

    self.output_for_node[node.id] = output

import deepchem as dc

import numpy as np

from tensorflow.python.framework import test_util

class TestOntology(test_util.TensorFlowTestCase):

  """Test OntologyModel."""

  def test_ontology_regressor(self):

    """Test training an OntologyModel for regression."""

    # Create a dataset with three tasks.  The first two tasks each depend only

    # on half the features.  The third task depends on all of them.

    n_features = 8

    n_tasks = 3

    n_samples = 30

    x = np.random.random((n_samples, n_features))

    y = np.zeros((n_samples, n_tasks))

    y[:, 0] = np.sum(x[:, :(n_features//2)], axis=1)

    y[:, 1] = np.sum(x[:, (n_features//2):], axis=1)

    y[:, 2] = 0.5*np.sum(x, axis=1)

    dataset = dc.data.NumpyDataset(x, y)

    # Create an OntologyModel.  Two leaf nodes contain half the features.

    feature_ids = [str(i) for i in range(n_features)]

    node1 = dc.models.OntologyNode('leaf1', 5, feature_ids[:(n_features//2)])

    node2 = dc.models.OntologyNode('leaf2', 5, feature_ids[(n_features//2):])

    node3 = dc.models.OntologyNode('root', 5, children=[node1, node2])

    model = dc.models.OntologyModel(n_tasks, feature_ids, node3, learning_rate=0.02)

    # Train the model on the datase.

    model.fit(dataset, nb_epoch=1000)

    # It should have learned to predict all of the tasks accurately.

    pred_error = np.mean((model.predict_on_batch(x)-y)**2, axis=0)

    assert np.all(pred_error < 0.01)

    # In addition, it should be able to predict the first task based only on the

    # first leaf node, and the second task based only on the second leaf node.

    leaf1_error = np.mean((model.predict_on_batch(x, outputs=model.prediction_for_node['leaf1'])-y)**2, axis=0)

    leaf2_error = np.mean((model.predict_on_batch(x, outputs=model.prediction_for_node['leaf2'])-y)**2, axis=0)

    assert leaf1_error[0] < 0.01

    assert leaf2_error[1] < 0.01

    assert np.mean([leaf1_error[1], leaf1_error[2], leaf2_error[0], leaf2_error[2]]) > 0.01