find top guides for all transcripts and then scan off-targets simultaneously

.gitignore

@@ -0,0 +1,2 @@
+off_target.csv
+on_target.csv

tiger.py

@@ -1,7 +1,6 @@
 import argparse
 import os
 import gzip
-import numpy as np
 import pandas as pd
 import tensorflow as tf
 from Bio import SeqIO
@@ -15,6 +14,7 @@ NUCLEOTIDE_COMPLEMENT = dict(zip(['A', 'C', 'G', 'T'], ['T', 'G', 'C', 'A']))
 NUM_TOP_GUIDES = 10
 NUM_MISMATCHES = 3
 REFERENCE_TRANSCRIPTS = ('gencode.v19.pc_transcripts.fa.gz', 'gencode.v19.lncRNA_transcripts.fa.gz')
+BATCH_SIZE = 500
 
 # configure GPUs
 for gpu in tf.config.list_physical_devices('GPU'):
@@ -105,41 +105,42 @@ def predict_on_target(transcript_seq: str, model: tf.keras.Model):
     # get predictions
     normalized_lfc = model.predict_step(model_inputs)
     predictions = pd.DataFrame({'Guide': guide_seq, 'Normalized LFC': tf.squeeze(normalized_lfc).numpy()})
-    predictions = predictions.set_index('Guide').sort_values('Normalized LFC')
+    predictions = predictions.sort_values('Normalized LFC')
 
     return predictions
 
 
-def find_off_targets(guides, batch_size=500):
+def find_off_targets(top_guides: pd.DataFrame):
 
     # load reference transcripts
     reference_transcripts = load_transcripts([os.path.join('transcripts', f) for f in REFERENCE_TRANSCRIPTS])
 
     # one-hot encode guides to form a filter
-    guide_filter = one_hot_encode_sequence(sequence_complement(guides), add_context_padding=False)
+    guide_filter = one_hot_encode_sequence(sequence_complement(top_guides['Guide']), add_context_padding=False)
     guide_filter = tf.transpose(guide_filter, [1, 2, 0])
     guide_filter = tf.cast(guide_filter, tf.float16)
 
     # loop over transcripts in batches
     i = 0
     print('Scanning for off-targets')
-    df_off_targets = pd.DataFrame()
+    off_targets = pd.DataFrame()
     while i < len(reference_transcripts):
         # select batch
-        df_batch = reference_transcripts.iloc[i:min(i + batch_size, len(reference_transcripts))]
-        i += batch_size
+        df_batch = reference_transcripts.iloc[i:min(i + BATCH_SIZE, len(reference_transcripts))]
+        i += BATCH_SIZE
 
         # find and log off-targets
         transcripts = one_hot_encode_sequence(df_batch['seq'].values.tolist(), add_context_padding=False)
         transcripts = tf.cast(transcripts, guide_filter.dtype)
         num_mismatches = GUIDE_LEN - tf.nn.conv1d(transcripts, guide_filter, stride=1, padding='SAME')
         loc_off_targets = tf.where(tf.round(num_mismatches) <= NUM_MISMATCHES).numpy()
-        df_off_targets = pd.concat([df_off_targets, pd.DataFrame({
-            'Guide': np.array(guides)[loc_off_targets[:, 2]],
-            'Isoform': df_batch.index.values[loc_off_targets[:, 0]],
+        off_targets = pd.concat([off_targets, pd.DataFrame({
+            'On-target ID': top_guides.iloc[loc_off_targets[:, 2]]['On-target ID'],
+            'Guide': top_guides.iloc[loc_off_targets[:, 2]]['Guide'],
+            'Off-target ID': df_batch.index.values[loc_off_targets[:, 0]],
+            'Target': df_batch['seq'].values[loc_off_targets[:, 0]],
             'Mismatches': tf.gather_nd(num_mismatches, loc_off_targets).numpy().astype(int),
             'Midpoint': loc_off_targets[:, 1],
-            'Target': df_batch['seq'].values[loc_off_targets[:, 0]],
         })])
 
         # progress update
@@ -147,7 +148,7 @@ def find_off_targets(guides, batch_size=500):
     print('')
 
     # trim transcripts to targets
-    dict_off_targets = df_off_targets.to_dict('records')
+    dict_off_targets = off_targets.to_dict('records')
     for row in dict_off_targets:
         start_location = row['Midpoint'] - (GUIDE_LEN // 2)
         if start_location < CONTEXT_5P:
@@ -160,9 +161,9 @@ def find_off_targets(guides, batch_size=500):
             row['Target'] = row['Target'][start_location - CONTEXT_5P:start_location + GUIDE_LEN + CONTEXT_3P]
         if row['Mismatches'] == 0 and 'N' not in row['Target']:
             assert row['Guide'] == sequence_complement([row['Target'][CONTEXT_5P:TARGET_LEN-CONTEXT_3P]])[0]
-    df_off_targets = pd.DataFrame(dict_off_targets)
+    off_targets = pd.DataFrame(dict_off_targets)
 
-    return df_off_targets
+    return off_targets
 
 
 def predict_off_target(off_targets: pd.DataFrame, model: tf.keras.Model):
@@ -174,12 +175,12 @@ def predict_off_target(off_targets: pd.DataFrame, model: tf.keras.Model):
         tf.reshape(one_hot_encode_sequence(off_targets['Target'], add_context_padding=False), [len(off_targets), -1]),
         tf.reshape(one_hot_encode_sequence(off_targets['Guide'], add_context_padding=True), [len(off_targets), -1]),
         ], axis=-1)
-    off_targets['Normalized LFC'] = model.predict_step(model_inputs)
+    off_targets['Normalized LFC'] = model.predict(model_inputs, batch_size=BATCH_SIZE, verbose=False)
 
-    return off_targets.set_index('Guide').sort_values('Normalized LFC')
+    return off_targets.sort_values('Normalized LFC')
 
 
-def tiger_exhibit(transcript):
+def tiger_exhibit(transcripts: pd.DataFrame):
 
     # load model
     if os.path.exists('model'):
@@ -188,20 +189,47 @@ def tiger_exhibit(transcript):
         print('no saved model!')
         exit()
 
-    # on-target predictions
-    on_target_predictions = predict_on_target(transcript, model=tiger)
-
-    # keep only top guides
-    on_target_predictions = on_target_predictions.iloc[:NUM_TOP_GUIDES]
+    # find top guides for each transcript
+    on_target_predictions = pd.DataFrame(columns=['On-target ID', 'Guide', 'Normalized LFC'])
+    for index, row in transcripts.iterrows():
+        df = predict_on_target(row['seq'], model=tiger)
+        df['On-target ID'] = index
+        on_target_predictions = pd.concat([on_target_predictions, df.iloc[:NUM_TOP_GUIDES]])
 
     # predict off-target effects for top guides
-    off_targets = find_off_targets(on_target_predictions.index.values.tolist())
+    off_targets = find_off_targets(on_target_predictions)
     off_target_predictions = predict_off_target(off_targets, model=tiger)
 
-    return on_target_predictions, off_target_predictions
+    return on_target_predictions.reset_index(drop=True), off_target_predictions.reset_index(drop=True)
 
 
 if __name__ == '__main__':
 
+    # common arguments
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--fasta_path', type=str, default=None)
+    parser.add_argument('--simple_test', action='store_true', default=False)
+    args = parser.parse_args()
+
     # simple test case
-    print(tiger_exhibit('ATGCAGGACGCGGAGAACGTGGCGGTGCCCGAGGCGGCCGAGGAGCGCGC'.lower()))  # first 50 from EIF3B-003's CDS
+    if args.simple_test:
+        # first 50 from EIF3B-003's CDS
+        simple_test = pd.DataFrame(dict(id=['user entry'], seq=['ATGCAGGACGCGGAGAACGTGGCGGTGCCCGAGGCGGCCGAGGAGCGCGC']))
+        simple_test.set_index('id', inplace=True)
+        df_on_target, df_off_target = tiger_exhibit(simple_test)
+        df_on_target.to_csv('on_target.csv')
+        df_off_target.to_csv('off_target.csv')
+
+    # # directory of fasta files
+    # elif args.dir_in is not None and os.path.exists(args.fasta_path):
+    #     transcripts = pd.DataFrame()
+    #     for fasta in os.listdir(args.fasta_path):
+    #         df = pd.DataFrame([(t.id, str(t.seq)) for t in SeqIO.parse(fasta, 'fasta')], columns=['id', 'seq'])
+    #
+    #         try:
+    #             for tran in SeqIO.parse(os.path.join(in_path, f), 'fasta'):
+    #                 on_targets, off_targets = tiger_exhibit(str(tran.seq))
+    #                 on_targets.to_csv(os.path.join(out_path, tran.id + '-top-guides.csv'))
+    #                 off_targets.to_csv(os.path.join(out_path, tran.id + '-off-targets.csv'))
+    #         except Exception:
+    #             warnings.warn(f)