add a predict module and bug fixes to performance metrics, non strand specific module

Author: kathyxchen

models/non_strand_specific_module.py

@@ -1,4 +1,3 @@
-import numpy as np
 import torch
 from torch.nn.modules import Module
 
@@ -23,7 +22,6 @@ class NonStrandSpecific(Module):
     def __init__(self, model, mode="mean"):
         super(NonStrandSpecific, self).__init__()
 
-        print(mode)
         self.model = model
 
         if mode != "mean" and mode != "max":
@@ -42,19 +40,5 @@ def forward(self, input):
         if self.mode == "mean":
             return (output + output_from_rev) / 2
         else:
-            max_output = torch.max(
-                output.abs(), output_from_rev.abs())
-            np_output = output.data.cpu().numpy()
-            print(np_output)
-
-            it = np.nditer(np_output, flags=["multi_index"])
-            while not it.finished:
-                index = it.multi_index
-                print(it[0])
-                if max_output.data[index] != abs(it[0]):
-                    max_output.data[index] = output_from_rev.data[index]
-                else:
-                    max_output.data[index] = it[0]
-                it.iternext()
-            return max_output
+            return torch.max(output, output_from_rev)
 

selene/__init__.py

@@ -1,2 +1,2 @@
-__all__ = ["sequences", "targets", "samplers", "utils"]
+__all__ = ["predict", "sequences", "targets", "samplers", "utils"]
 from .model_train import ModelController

selene/predict/__init__.py

@@ -0,0 +1,172 @@
+import itertools
+
+import numpy as np
+import torch
+from torch.autograd import Variable
+
+from .predict_handlers import DiffScoreHandler, LogitScoreHandler, \
+        WritePredictionsHandler
+from ..sequences import Genome
+from ..sequences import sequence_to_encoding
+
+
+def predict(model, batch_sequences, use_cuda=False):
+    inputs = torch.Tensor(batch_sequences)
+    if use_cuda:
+        inputs = inputs.cuda()
+    inputs = Variable(inputs, volatile=True)
+    outputs = model.forward(inputs.transpose(1, 2))
+    return outputs
+
+
+def predict_on_encoded_sequences(model,
+                                 sequences,
+                                 batch_size=64,
+                                 use_cuda=False):
+    predictions = []
+    n_examples, _, _ = sequences.shape
+
+    for i in range(0, n_examples, batch_size):
+        start = i
+        end = i + batch_size
+        batch_sequences = sequences[start:end, :, :]
+        outputs = predict(model, batch_sequences, use_cuda=use_cuda)
+        predictions.append(outputs.data.cpu().numpy())
+    return np.vstack(predictions)
+
+
+def in_silico_mutagenesis_sequences(input_sequence,
+                                    mutate_n_bases=1):
+    """Creates a list containing each mutation that occurs from in silico
+    mutagenesis across the whole sequence.
+
+    Parameters
+    ----------
+    input_sequence : str
+    mutate_n_bases : int
+
+    Returns
+    -------
+    list
+        A list of all possible mutations. Each element in the list is
+        itself a list of tuples, e.g. [(0, 'T')] if we are only mutating
+        1 base at a time. Each tuple is the position to mutate and the base
+        with which we are replacing the reference base.
+
+        For a sequence of length 1000, mutating 1 base at a time means that
+        we return a list of length 3000.
+    """
+    sequence_alts = []
+    for index, ref in enumerate(input_sequence):
+        alts = []
+        for base in Genome.BASES_ARR:
+            if base == ref:
+                continue
+            alts.append(base)
+        sequence_alts.append(alts)
+
+    all_mutated_sequences = []
+    for indices in itertools.combinations(
+            range(len(input_sequence)), mutate_n_bases):
+        pos_mutations = []
+        for i in indices:
+            pos_mutations.append(sequence_alts[i])
+        for mutations in itertools.product(*pos_mutations):
+            all_mutated_sequences.append(list(zip(indices, mutations)))
+    return all_mutated_sequences
+
+
+def _ism_sample_id(dna_sequence, mutation_information):
+    positions = []
+    refs = []
+    alts = []
+    for (position, alt) in mutation_information:
+        positions.append(str(position))
+        refs.append(dna_sequence[position])
+        alts.append(alt)
+    return (';'.join(positions), ';'.join(refs), ';'.join(alts))
+
+
+def in_silico_mutagenesis_predict(model,
+                                  batch_size,
+                                  sequence,
+                                  mutations_list,
+                                  use_cuda=False,
+                                  reporters=[]):
+    current_sequence_encoding = sequence_to_encoding(
+        sequence, Genome.BASE_TO_INDEX)
+    for i in range(0, len(mutations_list), batch_size):
+        start = i
+        end = i + batch_size
+
+        mutated_sequences = np.zeros(
+            (batch_size, *current_sequence_encoding.shape))
+
+        batch_ids = []
+        for ix, mutation_info in enumerate(mutations_list[start:end]):
+            mutated_seq = mutate_sequence(
+                current_sequence_encoding, mutation_info)
+            mutated_sequences[ix, :, :] = mutated_seq
+            batch_ids.append(_ism_sample_id(sequence, mutation_info))
+        outputs = predict(
+            model, mutated_sequences, use_cuda=use_cuda).data.cpu().numpy()
+
+        for r in reporters:
+            r.handle_batch_predictions(outputs, batch_ids)
+
+    for r in reporters:
+        r.write_to_file()
+
+
+def _reverse_strand(dna_sequence):
+    reverse_bases = [Genome.COMPLEMENTARY_BASE[b] for b in dna_sequence[::-1]]
+    return ''.join(reverse_bases)
+
+
+def mutate_sequence(dna_encoded_sequence, mutation_information):
+    mutated_seq = np.copy(dna_encoded_sequence)
+    for (position, alt) in mutation_information:
+        replace_base = Genome.BASE_TO_INDEX[alt]
+        mutated_seq[position, :] = 0
+        mutated_seq[position, replace_base] = 1
+    return mutated_seq
+
+
+def in_silico_mutagenesis(model,
+                          batch_size,
+                          input_sequence,
+                          features_list,
+                          save_diffs,
+                          mutate_n_bases=1,
+                          use_cuda=False,
+                          save_logits=None,
+                          save_predictions=None):
+    mutated_sequences = in_silico_mutagenesis_sequences(
+        input_sequence, mutate_n_bases=1)
+
+    current_sequence_encoding = sequence_to_encoding(
+        input_sequence, Genome.BASE_TO_INDEX)
+
+    base_encoding = current_sequence_encoding.reshape(
+        (1, *current_sequence_encoding.shape))
+    base_preds = predict(
+        model, base_encoding).data.cpu().numpy()
+
+    reporters = []
+    nonfeature_cols = ["pos", "ref", "alt"]
+    if save_diffs:
+        diff_handler = DiffScoreHandler(
+            base_preds, features_list, nonfeature_cols, save_diffs)
+        reporters.append(diff_handler)
+    if save_logits:
+        logit_handler = LogitScoreHandler(
+            base_preds, features_list, nonfeature_cols, save_logits)
+        reporters.append(logit_handler)
+    if save_predictions:
+        preds_handler = WritePredictionsHandler(
+            features_list, nonfeature_cols, save_predictions)
+        reporters.append(preds_handler)
+
+    in_silico_mutagenesis_predict(
+        model, batch_size, input_sequence, mutated_sequences,
+        use_cuda=use_cuda, reporters=reporters)

selene/predict/model_predict.py

@@ -0,0 +1,4 @@
+from .handler import PredictionsHandler
+from .diff_score_handler import DiffScoreHandler
+from .logit_score_handler import LogitScoreHandler
+from .write_predictions_handler import WritePredictionsHandler

selene/predict/predict_handlers/__init__.py

@@ -0,0 +1,33 @@
+import numpy as np
+
+from .handler import _write_to_file, PredictionsHandler
+
+
+class DiffScoreHandler(PredictionsHandler):
+
+    def __init__(self,
+                 baseline_prediction,
+                 features_list,
+                 nonfeature_columns,
+                 out_filename):
+        self.baseline_prediction = baseline_prediction
+        self.column_names = nonfeature_columns + features_list
+        self.results = []
+        self.samples = []
+        self.out_filename = out_filename
+
+    def handle_batch_predictions(self,
+                                 batch_predictions,
+                                 batch_ids):
+        absolute_diffs = np.abs(self.baseline_prediction - batch_predictions)
+        self.results.append(absolute_diffs)
+        self.samples.append(batch_ids)
+        return absolute_diffs
+
+    def write_to_file(self):
+        self.results = np.vstack(self.results)
+        self.samples = np.vstack(self.samples)
+        _write_to_file(self.results,
+                       self.samples,
+                       self.column_names,
+                       self.out_filename)

selene/predict/predict_handlers/diff_score_handler.py

@@ -0,0 +1,39 @@
+"""
+This class is the abstract base class for handling model predicions
+"""
+from abc import ABCMeta
+from abc import abstractmethod
+
+
+def _write_to_file(feature_predictions, info_cols, column_names, filename):
+    with open(filename, 'w+') as file_handle:
+        file_handle.write("{columns}\n".format(
+            columns='\t'.join(column_names)))
+        for info, preds in zip(info_cols, feature_predictions):
+            feature_cols = '\t'.join(
+                probabilities_to_string(preds))
+            info_cols = '\t'.join(info)
+            file_handle.write(f"{info_cols}\t{feature_cols}\n")
+
+
+def probabilities_to_string(probabilities):
+    return ["{:.2e}".format(p) for p in probabilities]
+
+
+class PredictionsHandler(metaclass=ABCMeta):
+    """
+    The base class for handling model predictions.
+    """
+    @abstractmethod
+    def handle_batch_predictions(self, *args, **kwargs):
+        """
+        Must be able to handle a batch of model predictions.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def write_to_file(self, *args, **kwargs):
+        """
+        Writes accumulated handler results to file.
+        """
+        raise NotImplementedError

selene/predict/predict_handlers/handler.py

@@ -0,0 +1,34 @@
+import numpy as np
+from scipy.special import logit
+
+from .handler import _write_to_file, PredictionsHandler
+
+
+class LogitScoreHandler(PredictionsHandler):
+
+    def __init__(self,
+                 baseline_prediction,
+                 features_list,
+                 nonfeature_columns,
+                 out_filename):
+        self.logit_baseline = logit(baseline_prediction)
+        self.column_names = nonfeature_columns + features_list
+        self.results = []
+        self.samples = []
+        self.out_filename = out_filename
+
+    def handle_batch_predictions(self,
+                                 batch_predictions,
+                                 batch_ids):
+        absolute_logits = np.abs(self.logit_baseline - logit(batch_predictions))
+        self.results.append(absolute_logits)
+        self.samples.append(batch_ids)
+        return absolute_logits
+
+    def write_to_file(self):
+        self.results = np.vstack(self.results)
+        self.samples = np.vstack(self.samples)
+        _write_to_file(self.results,
+                       self.samples,
+                       self.column_names,
+                       self.out_filename)

selene/predict/predict_handlers/logit_score_handler.py

@@ -0,0 +1,26 @@
+import numpy as np
+
+from .handler import _write_to_file, PredictionsHandler
+
+class WritePredictionsHandler(PredictionsHandler):
+
+    def __init__(self, features_list, nonfeature_columns, out_filename):
+        self.column_names = nonfeature_columns + features_list
+        self.results = []
+        self.samples = []
+        self.out_filename = out_filename
+
+    def handle_batch_predictions(self,
+                                 batch_predictions,
+                                 batch_ids):
+        self.results.append(batch_predictions)
+        self.samples.append(batch_ids)
+        return batch_predictions
+
+    def write_to_file(self):
+        self.results = np.vstack(self.results)
+        self.samples = np.vstack(self.samples)
+        _write_to_file(self.results,
+                       self.samples,
+                       self.column_names,
+                       self.out_filename)

selene/predict/predict_handlers/write_predictions_handler.py

@@ -0,0 +1,47 @@
+import unittest
+
+import numpy as np
+
+from selene.predict.model_predict import in_silico_mutagenesis_sequences, \
+        in_silico_mutagenesis_predict
+
+
+class TestModelPredict(unittest.TestCase):
+
+    def setUp(self):
+        self.bases_arr = ['A', 'C', 'G', 'T']
+        self.bases_encoding = {'A': 0, 'C': 1, 'G': 2, 'T': 3}
+        self.input_sequence = "ATCCG"
+
+    def test_in_silico_muta_sequences_single(self):
+        observed = in_silico_mutagenesis_sequences("ATCCG")
+        expected = [
+            (0, 'C'), (0, 'G'), (0, 'T'),
+            (1, 'A'), (1, 'C'), (1, 'G'),
+            (2, 'A'), (2, 'G'), (2, 'T'),
+            (3, 'A'), (3, 'G'), (3, 'T'),
+            (4, 'A'), (4, 'C'), (4, 'T')]
+
+        expected_lists = [[e] for e in expected]
+        self.assertListEqual(observed, expected_lists)
+
+    def test_in_silico_muta_sequences_double(self):
+        observed = in_silico_mutagenesis_sequences(
+            "ATC", mutate_n_bases=2)
+        expected = [
+            [(0, 'C'), (1, 'A')], [(0, 'G'), (1, 'A')], [(0, 'T'), (1, 'A')],
+            [(0, 'C'), (1, 'C')], [(0, 'G'), (1, 'C')], [(0, 'T'), (1, 'C')],
+            [(0, 'C'), (1, 'G')], [(0, 'G'), (1, 'G')], [(0, 'T'), (1, 'G')],
+
+            [(0, 'C'), (2, 'A')], [(0, 'G'), (2, 'A')], [(0, 'T'), (2, 'A')],
+            [(0, 'C'), (2, 'G')], [(0, 'G'), (2, 'G')], [(0, 'T'), (2, 'G')],
+            [(0, 'C'), (2, 'T')], [(0, 'G'), (2, 'T')], [(0, 'T'), (2, 'T')],
+
+            [(1, 'A'), (2, 'A')], [(1, 'C'), (2, 'A')], [(1, 'G'), (2, 'A')],
+            [(1, 'A'), (2, 'G')], [(1, 'C'), (2, 'G')], [(1, 'G'), (2, 'G')],
+            [(1, 'A'), (2, 'T')], [(1, 'C'), (2, 'T')], [(1, 'G'), (2, 'T')],
+        ]
+        self.assertCountEqual(observed, expected)
+
+if __name__ == "__main__":
+    unittest.main()