off-target predictions

tiger.py

@@ -65,20 +65,13 @@ def process_data(transcript_seq: str):
     return target_seq, guide_seq, model_inputs
 
 
-def predict_on_target(transcript_seq: str):
-
-    # load model
-    if os.path.exists('model'):
-        tiger = tf.keras.models.load_model('model')
-    else:
-        print('no saved model!')
-        exit()
+def predict_on_target(transcript_seq: str, model: tf.keras.Model):
 
     # parse transcript sequence
     target_seq, guide_seq, model_inputs = process_data(transcript_seq)
 
     # get predictions
-    normalized_lfc = tiger.predict_step(model_inputs)
+    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')
 
@@ -86,6 +79,8 @@ def predict_on_target(transcript_seq: str):
 
 
 def find_off_targets(guides, batch_size=1000):
+
+    # load reference transcripts
     with gzip.open(os.path.join('transcripts', 'gencode.v19.pc_transcripts.fa.gz'), 'rt') as file:
         df_transcripts = pd.DataFrame([(t.id, str(t.seq)) for t in SeqIO.parse(file, 'fasta')], columns=['id', 'seq'])
     df_transcripts['id'] = df_transcripts['id'].apply(lambda s: s.split('|')[4])
@@ -132,16 +127,36 @@ def find_off_targets(guides, batch_size=1000):
     return df_off_targets
 
 
+def predict_off_target(off_targets: pd.DataFrame, model: tf.keras.Model):
+
+    # append predictions off-target predictions
+    model_inputs = tf.concat([
+        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)
+
+    return off_targets
+
+
 def tiger_exhibit(transcript):
 
+    # load model
+    if os.path.exists('model'):
+        tiger = tf.keras.models.load_model('model')
+    else:
+        print('no saved model!')
+        exit()
+
     # on-target predictions
-    on_target_predictions = predict_on_target(transcript)
+    on_target_predictions = predict_on_target(transcript, model=tiger)
 
     # keep only top guides
     on_target_predictions = on_target_predictions.iloc[:NUM_TOP_GUIDES]
 
-    # scan for off-targets for top guides
+    # predict off-target effects for top guides
     off_targets = find_off_targets(on_target_predictions.index.values.tolist())
+    off_targets = predict_off_target(off_targets, model=tiger)
 
     return on_target_predictions, off_targets