titration mode!

tiger.py

@@ -38,7 +38,7 @@ NUM_MISMATCHES = 3
 RUN_MODES = dict(
     all='All on-target guides per transcript',
     top_guides='Top {:d} guides per transcript'.format(NUM_TOP_GUIDES),
-    # titration='Top {:d} guides per transcript & their titration candidates'.format(NUM_TOP_GUIDES)  # TODO: do this!
+    titration='Top {:d} guides per transcript & their titration candidates'.format(NUM_TOP_GUIDES)
 )
 
 
@@ -175,7 +175,7 @@ def get_on_target_predictions(transcripts: pd.DataFrame, model: tf.keras.Model,
         scores = prediction_transform(tf.squeeze(lfc_estimate).numpy())
         predictions = pd.concat([predictions, pd.DataFrame({
             ID_COL: [index] * len(scores),
-            TARGET_COL: [seq[CONTEXT_5P:len(seq) - CONTEXT_3P] for seq in target_seq],
+            TARGET_COL: target_seq,
             GUIDE_COL: guide_seq,
             SCORE_COL: scores})])
 
@@ -202,6 +202,27 @@ def top_guides_per_transcript(predictions: pd.DataFrame):
     return top_guides.reset_index(drop=True)
 
 
+def get_titration_candidates(top_guide_predictions: pd.DataFrame):
+
+    # generate a table of all titration candidates
+    titration_candidates = pd.DataFrame()
+    for _, row in top_guide_predictions.iterrows():
+        for i in range(len(row[GUIDE_COL])):
+            nt = row[GUIDE_COL][i]
+            for mutation in set(NUCLEOTIDE_TOKENS.keys()) - {nt, 'N'}:
+                sm_guide = list(row[GUIDE_COL])
+                sm_guide[i] = mutation
+                sm_guide = ''.join(sm_guide)
+                assert row[GUIDE_COL] != sm_guide
+                titration_candidates = pd.concat([titration_candidates, pd.DataFrame({
+                    ID_COL: [row[ID_COL]],
+                    TARGET_COL: [row[TARGET_COL]],
+                    GUIDE_COL: [sm_guide]
+                })])
+
+    return titration_candidates
+
+
 def find_off_targets(top_guides: pd.DataFrame, status_update_fn=None):
 
     # load reference transcripts
@@ -284,7 +305,7 @@ def predict_off_target(off_targets: pd.DataFrame, model: tf.keras.Model):
     # trim context sequence
     off_targets[TARGET_COL] = off_targets[TARGET_COL].apply(lambda seq: seq[CONTEXT_5P:len(seq) - CONTEXT_3P])
 
-    return off_targets.sort_values(SCORE_COL, ascending=False).reset_index(drop=True)
+    return off_targets.reset_index(drop=True)
 
 
 def tiger_exhibit(transcripts: pd.DataFrame, mode: str, check_off_targets: bool, status_update_fn=None):
@@ -300,30 +321,36 @@ def tiger_exhibit(transcripts: pd.DataFrame, mode: str, check_off_targets: bool,
     on_target_predictions = get_on_target_predictions(transcripts, tiger, status_update_fn)
 
     # initialize other outputs
-    off_target_predictions = None
+    titration_predictions = off_target_predictions = None
 
     if mode == 'all' and not check_off_targets:
-        pass  # nothing to do!
+        off_target_candidates = None
 
     elif mode == 'top_guides':
         on_target_predictions = top_guides_per_transcript(on_target_predictions)
+        off_target_candidates = on_target_predictions
 
-    # TODO: add titration candidates
+    elif mode == 'titration':
+        on_target_predictions = top_guides_per_transcript(on_target_predictions)
+        titration_candidates = get_titration_candidates(on_target_predictions)
+        titration_predictions = predict_off_target(titration_candidates, model=tiger)
+        off_target_candidates = pd.concat([on_target_predictions, titration_predictions])
 
     else:
         raise NotImplementedError
 
     # check off-target effects for top guides
-    if check_off_targets:
-        off_targets = find_off_targets(on_target_predictions, status_update_fn)
-        off_target_predictions = predict_off_target(off_targets, model=tiger)
+    if check_off_targets and off_target_candidates is not None:
+        off_target_candidates = find_off_targets(off_target_candidates, status_update_fn)
+        off_target_predictions = predict_off_target(off_target_candidates, model=tiger)
+        off_target_predictions = off_target_predictions.sort_values(SCORE_COL, ascending=False)
 
     # reverse guide sequences
     on_target_predictions[GUIDE_COL] = on_target_predictions[GUIDE_COL].apply(lambda s: s[::-1])
     if check_off_targets and len(off_target_predictions) > 0:
         off_target_predictions[GUIDE_COL] = off_target_predictions[GUIDE_COL].apply(lambda s: s[::-1])
 
-    return on_target_predictions, off_target_predictions
+    return on_target_predictions, titration_predictions, off_target_predictions
 
 
 if __name__ == '__main__':
@@ -336,7 +363,7 @@ if __name__ == '__main__':
     args = parser.parse_args()
 
     # check for any existing results
-    if os.path.exists('on_target.csv') or os.path.exists('off_target.csv'):
+    if os.path.exists('on_target.csv') or os.path.exists('titration.csv') or os.path.exists('off_target.csv'):
         raise FileExistsError('please rename or delete existing results')
 
     # load transcripts from a directory of fasta files
@@ -360,11 +387,17 @@ if __name__ == '__main__':
 
         # run batch
         idx_stop = min(idx + BATCH_SIZE_TRANSCRIPTS, len(df_transcripts))
-        df_on_target, df_off_target = tiger_exhibit(df_transcripts[idx:idx_stop], args.mode, args.check_off_targets)
+        df_on_target, df_titration, df_off_target = tiger_exhibit(
+            transcripts=df_transcripts[idx:idx_stop],
+            mode=args.mode,
+            check_off_targets=args.check_off_targets
+        )
 
         # save batch results
         df_on_target.to_csv('on_target.csv', header=batch == 1, index=False, mode='a')
-        if args.check_off_targets:
+        if df_titration is not None:
+            df_titration.to_csv('titration.csv', header=batch == 1, index=False, mode='a')
+        if df_off_target is not None:
             df_off_target.to_csv('off_target.csv', header=batch == 1, index=False, mode='a')
 
         # clear session to prevent memory blow up