Beginning dragonn adaptation

from __future__ import absolute_import, division, print_function

import matplotlib

import numpy as np

import os

import subprocess

import sys

import tempfile

matplotlib.use('pdf')

import matplotlib.pyplot as plt

from abc import abstractmethod, ABCMeta

from dragonn.metrics import ClassificationResult

from sklearn.svm import SVC as scikit_SVC

from sklearn.tree import DecisionTreeClassifier as scikit_DecisionTree

from sklearn.ensemble import RandomForestClassifier

class Model(object):

  __metaclass__ = ABCMeta

  @abstractmethod

  def __init__(self, **hyperparameters):

    pass

  @abstractmethod

  def train(self, X, y, validation_data):

    pass

  @abstractmethod

  def predict(self, X):

    pass

  def test(self, X, y):

    return ClassificationResult(y, self.predict(X))

  def score(self, X, y, metric):

    return self.test(X, y)[metric]

class SequenceDNN(Model):

  """

  Sequence DNN models.

  Parameters

  ----------

  seq_length : int, optional

      length of input sequence.

  keras_model : instance of keras.models.Sequential, optional

      seq_length or keras_model must be specified.

  num_tasks : int, optional

      number of tasks. Default: 1.

  num_filters : list[int] | tuple[int]

      number of convolutional filters in each layer. Default: (15,).

  conv_width : list[int] | tuple[int]

      width of each layer's convolutional filters. Default: (15,).

  pool_width : int

      width of max pooling after the last layer. Default: 35.

  L1 : float

      strength of L1 penalty.

  dropout : float

      dropout probability in every convolutional layer. Default: 0.

  verbose: int

      Verbosity level during training. Valida values: 0, 1, 2.

  Returns

  -------

  Compiled DNN model.

  """

  def __init__(self,

               seq_length=None,

               keras_model=None,

               use_RNN=False,

               num_tasks=1,

               num_filters=(15, 15, 15),

               conv_width=(15, 15, 15),

               pool_width=35,

               GRU_size=35,

               TDD_size=15,

               L1=0,

               dropout=0.0,

               num_epochs=100,

               verbose=1):

    from keras.models import Sequential

    from keras.layers.core import (Activation, Dense, Dropout, Flatten, Permute,

                                   Reshape, TimeDistributedDense)

    from keras.layers.convolutional import Convolution2D, MaxPooling2D

    from keras.layers.recurrent import GRU

    from keras.regularizers import l1

    self.num_tasks = num_tasks

    self.num_epochs = num_epochs

    self.verbose = verbose

    self.train_metrics = []

    self.valid_metrics = []

    if keras_model is not None and seq_length is None:

      self.model = keras_model

      self.num_tasks = keras_model.layers[-1].output_shape[-1]

    elif seq_length is not None and keras_model is None:

      self.model = Sequential()

      assert len(num_filters) == len(conv_width)

      for i, (nb_filter, nb_col) in enumerate(zip(num_filters, conv_width)):

        conv_height = 4 if i == 0 else 1

        self.model.add(

            Convolution2D(

                nb_filter=nb_filter,

                nb_row=conv_height,

                nb_col=nb_col,

                activation='linear',

                init='he_normal',

                input_shape=(1, 4, seq_length),

                W_regularizer=l1(L1),

                b_regularizer=l1(L1)))

        self.model.add(Activation('relu'))

        self.model.add(Dropout(dropout))

      self.model.add(MaxPooling2D(pool_size=(1, pool_width)))

      if use_RNN:

        num_max_pool_outputs = self.model.layers[-1].output_shape[-1]

        self.model.add(Reshape((num_filters[-1], num_max_pool_outputs)))

        self.model.add(Permute((2, 1)))

        self.model.add(GRU(GRU_size, return_sequences=True))

        self.model.add(TimeDistributedDense(TDD_size, activation='relu'))

      self.model.add(Flatten())

      self.model.add(Dense(output_dim=self.num_tasks))

      self.model.add(Activation('sigmoid'))

      self.model.compile(optimizer='adam', loss='binary_crossentropy')

    else:

      raise ValueError(

          "Exactly one of seq_length or keras_model must be specified!")

  def train(self,

            X,

            y,

            validation_data,

            early_stopping_metric='Loss',

            early_stopping_patience=5,

            save_best_model_to_prefix=None):

    if y.dtype != bool:

      assert set(np.unique(y)) == {0, 1}

      y = y.astype(bool)

    multitask = y.shape[1] > 1

    if not multitask:

      num_positives = y.sum()

      num_sequences = len(y)

      num_negatives = num_sequences - num_positives

    if self.verbose >= 1:

      print('Training model (* indicates new best result)...')

    X_valid, y_valid = validation_data

    early_stopping_wait = 0

    best_metric = np.inf if early_stopping_metric == 'Loss' else -np.inf

    for epoch in range(1, self.num_epochs + 1):

      self.model.fit(

          X,

          y,

          batch_size=128,

          nb_epoch=1,

          class_weight={

              True: num_sequences / num_positives,

              False: num_sequences / num_negatives

          } if not multitask else None,

          verbose=self.verbose >= 2)

      epoch_train_metrics = self.test(X, y)

      epoch_valid_metrics = self.test(X_valid, y_valid)

      self.train_metrics.append(epoch_train_metrics)

      self.valid_metrics.append(epoch_valid_metrics)

      if self.verbose >= 1:

        print('Epoch {}:'.format(epoch))

        print('Train {}'.format(epoch_train_metrics))

        print('Valid {}'.format(epoch_valid_metrics), end='')

      current_metric = epoch_valid_metrics[early_stopping_metric].mean()

      if (early_stopping_metric == 'Loss') == (current_metric <= best_metric):

        if self.verbose >= 1:

          print(' *')

        best_metric = current_metric

        best_epoch = epoch

        early_stopping_wait = 0

        if save_best_model_to_prefix is not None:

          self.save(save_best_model_to_prefix)

      else:

        if self.verbose >= 1:

          print()

        if early_stopping_wait >= early_stopping_patience:

          break

        early_stopping_wait += 1

    if self.verbose >= 1:

      print('Finished training after {} epochs.'.format(epoch))

      if save_best_model_to_prefix is not None:

        print("The best model's architecture and weights (from epoch {0}) "

              'were saved to {1}.arch.json and {1}.weights.h5'.format(

                  best_epoch, save_best_model_to_prefix))

  def predict(self, X):

    return self.model.predict(X, batch_size=128, verbose=False)

  def get_sequence_filters(self):

    """

    Returns 3D array of 2D sequence filters.

    """

    return self.model.layers[0].get_weights()[0].squeeze(axis=1)

  def deeplift(self, X, batch_size=200):

    """

    Returns (num_task, num_samples, 1, num_bases, sequence_length) deeplift score array.

    """

    assert len(np.shape(X)) == 4 and np.shape(X)[1] == 1

    from deeplift.conversion import keras_conversion as kc

    # convert to deeplift model and get scoring function

    deeplift_model = kc.convert_sequential_model(self.model, verbose=False)

    score_func = deeplift_model.get_target_contribs_func(

        find_scores_layer_idx=0)

    # use a 40% GC reference

    input_references = [np.array([0.3, 0.2, 0.2, 0.3])[None, None, :, None]]

    # get deeplift scores

    deeplift_scores = np.zeros((self.num_tasks,) + X.shape)

    for i in range(self.num_tasks):

      deeplift_scores[i] = score_func(

          task_idx=i,

          input_data_list=[X],

          batch_size=batch_size,

          progress_update=None,

          input_references_list=input_references)

    return deeplift_scores

  def in_silico_mutagenesis(self, X):

    """

    Returns (num_task, num_samples, 1, num_bases, sequence_length) ISM score array.

    """

    mutagenesis_scores = np.empty(X.shape + (self.num_tasks,), dtype=np.float32)

    wild_type_predictions = self.predict(X)

    wild_type_predictions = wild_type_predictions[:, np.newaxis, np.newaxis,

                                                  np.newaxis]

    for sequence_index, (sequence, wild_type_prediction) in enumerate(

        zip(X, wild_type_predictions)):

      mutated_sequences = np.repeat(

          sequence[np.newaxis], np.prod(sequence.shape), axis=0)

      # remove wild-type

      arange = np.arange(len(mutated_sequences))

      horizontal_cycle = np.tile(

          np.arange(sequence.shape[-1]), sequence.shape[-2])

      mutated_sequences[arange, :, :, horizontal_cycle] = 0

      # add mutant

      vertical_repeat = np.repeat(

          np.arange(sequence.shape[-2]), sequence.shape[-1])

      mutated_sequences[arange, :, vertical_repeat, horizontal_cycle] = 1

      # make mutant predictions

      mutated_predictions = self.predict(mutated_sequences)

      mutated_predictions = mutated_predictions.reshape(sequence.shape +

                                                        (self.num_tasks,))

      mutagenesis_scores[

          sequence_index] = wild_type_prediction - mutated_predictions

    return np.rollaxis(mutagenesis_scores, -1)

  @staticmethod

  def _plot_scores(X, output_directory, peak_width, score_func, score_name):

    from dragonn.plot import plot_bases_on_ax

    scores = score_func(X).squeeze(

        axis=2)  # (num_task, num_samples, num_bases, sequence_length)

    try:

      os.makedirs(output_directory)

    except OSError:

      pass

    num_tasks = len(scores)

    for task_index, task_scores in enumerate(scores):

      for sequence_index, sequence_scores in enumerate(task_scores):

        # sequence_scores is num_bases x sequence_length

        basewise_max_sequence_scores = sequence_scores.max(axis=0)

        plt.clf()

        figure, (top_axis, bottom_axis) = plt.subplots(2)

        top_axis.plot(

            range(1,

                  len(basewise_max_sequence_scores) + 1),

            basewise_max_sequence_scores)

        top_axis.set_title('{} scores (motif highlighted)'.format(score_name))

        peak_position = basewise_max_sequence_scores.argmax()

        top_axis.axvspan(

            peak_position - peak_width,

            peak_position + peak_width,

            color='grey',

            alpha=0.1)

        peak_sequence_scores = sequence_scores[:, peak_position - peak_width:

                                               peak_position + peak_width].T

        # Set non-max letter_heights to zero

        letter_heights = np.zeros_like(peak_sequence_scores)

        letter_heights[np.arange(len(letter_heights)),

                       peak_sequence_scores.argmax(axis=1)] = \

            basewise_max_sequence_scores[peak_position - peak_width :

                                         peak_position + peak_width]

        plot_bases_on_ax(letter_heights, bottom_axis)

        bottom_axis.set_xticklabels(

            tuple(

                map(str,

                    np.arange(peak_position - peak_width,

                              peak_position + peak_width + 1))))

        bottom_axis.tick_params(axis='x', labelsize='small')

        plt.xlabel('Position')

        plt.ylabel('Score')

        plt.savefig(

            os.path.join(output_directory, 'sequence_{}{}'.format(

                sequence_index, '_task_{}'.format(task_index)

                if num_tasks > 1 else '')))

        plt.close()

  def plot_deeplift(self, X, output_directory, peak_width=10):

    self._plot_scores(

        X,

        output_directory,

        peak_width,

        score_func=self.deeplift,

        score_name='DeepLift')

  def plot_in_silico_mutagenesis(self, X, output_directory, peak_width=10):

    self._plot_scores(

        X,

        output_directory,

        peak_width,

        score_func=self.in_silico_mutagenesis,

        score_name='ISM')

  def plot_architecture(self, output_file):

    from dragonn.visualize_util import plot as plot_keras_model

    plot_keras_model(self.model, output_file, show_shape=True)

  def save(self, save_best_model_to_prefix):

    arch_fname = save_best_model_to_prefix + '.arch.json'

    weights_fname = save_best_model_to_prefix + '.weights.h5'

    open(arch_fname, 'w').write(self.model.to_json())

    self.model.save_weights(weights_fname, overwrite=True)

  @staticmethod

  def load(arch_fname, weights_fname=None):

    from keras.models import model_from_json

    model_json_string = open(arch_fname).read()

    sequence_dnn = SequenceDNN(keras_model=model_from_json(model_json_string))

    if weights_fname is not None:

      sequence_dnn.model.load_weights(weights_fname)

    return sequence_dnn

class MotifScoreRNN(Model):

  def __init__(self, input_shape, gru_size=10, tdd_size=4):

    from keras.models import Sequential

    from keras.layers.core import (Activation, Dense, Flatten,

                                   TimeDistributedDense)

    from keras.layers.recurrent import GRU

    self.model = Sequential()

    self.model.add(

        GRU(gru_size, return_sequences=True, input_shape=input_shape))

    if tdd_size is not None:

      self.model.add(TimeDistributedDense(tdd_size))

    self.model.add(Flatten())

    self.model.add(Dense(1))

    self.model.add(Activation('sigmoid'))

    print('Compiling model...')

    self.model.compile(optimizer='adam', loss='binary_crossentropy')

  def train(self, X, y, validation_data):

    from keras.callbacks import EarlyStopping

    print('Training model...')

    multitask = y.shape[1] > 1

    if not multitask:

      num_positives = y.sum()

      num_sequences = len(y)

      num_negatives = num_sequences - num_positives

    self.model.fit(

        X,

        y,

        batch_size=128,

        nb_epoch=100,

        validation_data=validation_data,

        class_weight={

            True: num_sequences / num_positives,

            False: num_sequences / num_negatives

        } if not multitask else None,

        callbacks=[EarlyStopping(monitor='val_loss', patience=10)],

        verbose=True)

  def predict(self, X):

    return self.model.predict(X, batch_size=128, verbose=False)

class gkmSVM(Model):

  def __init__(self,

               prefix='./gkmSVM',

               word_length=11,

               mismatches=3,

               C=1,

               threads=1,

               cache_memory=100,

               verbosity=4):

    self.word_length = word_length

    self.mismatches = mismatches

    self.C = C

    self.threads = threads

    self.prefix = '_'.join(map(str, (prefix, word_length, mismatches, C)))

    options_list = zip(

        ['-l', '-d', '-c', '-T', '-m', '-v'],

        map(str,

            (word_length, mismatches, C, threads, cache_memory, verbosity)))

    self.options = ' '.join([' '.join(option) for option in options_list])

  @property

  def model_file(self):

    model_fname = '{}.model.txt'.format(self.prefix)

    return model_fname if os.path.isfile(model_fname) else None

  @staticmethod

  def encode_sequence_into_fasta_file(sequence_iterator, ofname):

    """writes sequences into fasta file

        """

    with open(ofname, "w") as wf:

      for i, seq in enumerate(sequence_iterator):

        print('>{}'.format(i), file=wf)

        print(seq, file=wf)

  def train(self, X, y, validation_data=None):

    """

        Trains gkm-svm, saves model file.

        """

    y = y.squeeze()

    pos_sequence = X[y]

    neg_sequence = X[~y]

    pos_fname = "%s.pos_seq.fa" % self.prefix

    neg_fname = "%s.neg_seq.fa" % self.prefix

    # create temporary fasta files

    self.encode_sequence_into_fasta_file(pos_sequence, pos_fname)

    self.encode_sequence_into_fasta_file(neg_sequence, neg_fname)

    # run command

    command = ' '.join(('gkmtrain', self.options, pos_fname, neg_fname,

                        self.prefix))

    process = subprocess.Popen(command, stdout=subprocess.PIPE, shell=True)

    process.wait()  # wait for it to finish

    # remove fasta files

    os.system("rm %s" % pos_fname)

    os.system("rm %s" % neg_fname)

  def predict(self, X):

    if self.model_file is None:

      raise RuntimeError("GkmSvm hasn't been trained!")

    # write test fasta file

    test_fname = "%s.test.fa" % self.prefix

    self.encode_sequence_into_fasta_file(X, test_fname)

    # test gkmsvm

    temp_ofp = tempfile.NamedTemporaryFile()

    threads_option = '-T %s' % (str(self.threads))

    command = ' '.join([

        'gkmpredict', test_fname, self.model_file, temp_ofp.name, threads_option

    ])

    process = subprocess.Popen(command, shell=True)

    process.wait()  # wait for it to finish

    os.system("rm %s" % test_fname)  # remove fasta file

    # get classification results

    temp_ofp.seek(0)

    y = np.array([line.split()[-1] for line in temp_ofp], dtype=float)

    temp_ofp.close()

    return np.expand_dims(y, 1)

class SVC(Model):

  def __init__(self):

    self.classifier = scikit_SVC(probability=True, kernel='linear')

  def train(self, X, y, validation_data=None):

    self.classifier.fit(X, y)

  def predict(self, X):

    return self.classifier.predict_proba(X)[:, 1:]

class DecisionTree(Model):

  def __init__(self):

    self.classifier = scikit_DecisionTree()

  def train(self, X, y, validation_data=None):

    self.classifier.fit(X, y)

  def predict(self, X):

    predictions = np.asarray(self.classifier.predict_proba(X))[..., 1]

    if len(predictions.shape) == 2:  # multitask

      predictions = predictions.T

    else:  # single-task

      predictions = np.expand_dims(predictions, 1)

    return predictions

class RandomForest(DecisionTree):

  def __init__(self):

    self.classifier = RandomForestClassifier(n_estimators=100)