massive cleanup with better table columns

tiger.py

@@ -1,6 +1,7 @@
 import argparse
 import os
 import gzip
+import pickle
 import numpy as np
 import pandas as pd
 import tensorflow as tf
@@ -14,6 +15,13 @@ NUCLEOTIDE_TOKENS = dict(zip(['A', 'C', 'G', 'T', 'N'], [0, 1, 2, 3, 255]))
 NUCLEOTIDE_COMPLEMENT = dict(zip(['A', 'C', 'G', 'T'], ['T', 'G', 'C', 'A']))
 NUM_TOP_GUIDES = 10
 NUM_MISMATCHES = 3
+ID_COL = 'Transcript ID'
+SEQ_COL = 'Sequence'
+TARGET_COL = 'Target Sequence'
+GUIDE_COL = 'Guide Sequence'
+SCORE_COL = 'Guide Score'
+RUN_MODE_ALL_PM = 'All on-target guides per transcript'
+RUN_MODE_TITRATION = 'Top guides per transcript'
 REFERENCE_TRANSCRIPTS = ('gencode.v19.pc_transcripts.fa.gz', 'gencode.v19.lncRNA_transcripts.fa.gz')
 BATCH_SIZE_COMPUTE = 500
 BATCH_SIZE_SCAN = 20
@@ -35,18 +43,18 @@ def load_transcripts(fasta_files):
         try:
             if os.path.splitext(file)[1] == '.gz':
                 with gzip.open(file, 'rt') as f:
-                    df = pd.DataFrame([(t.id, str(t.seq)) for t in SeqIO.parse(f, 'fasta')], columns=['id', 'seq'])
+                    df = pd.DataFrame([(t.id, str(t.seq)) for t in SeqIO.parse(f, 'fasta')], columns=[ID_COL, SEQ_COL])
             else:
-                df = pd.DataFrame([(t.id, str(t.seq)) for t in SeqIO.parse(file, 'fasta')], columns=['id', 'seq'])
+                df = pd.DataFrame([(t.id, str(t.seq)) for t in SeqIO.parse(file, 'fasta')], columns=[ID_COL, SEQ_COL])
         except Exception as e:
             print(e, 'while loading', file)
             continue
         transcripts = pd.concat([transcripts, df])
 
     # set index
-    transcripts['id'] = transcripts['id'].apply(lambda s: s.split('|')[0])
-    transcripts.set_index('id', inplace=True)
-    assert not transcripts.index.has_duplicates, "duplicate transcript ID's detected"
+    transcripts[ID_COL] = transcripts[ID_COL].apply(lambda s: s.split('|')[0])
+    transcripts.set_index(ID_COL, inplace=True)
+    assert not transcripts.index.has_duplicates, "duplicate transcript ID's detected in fasta file"
 
     return transcripts
 
@@ -101,6 +109,9 @@ def process_data(transcript_seq: str):
 
 
 def prediction_transform(predictions: np.array, **params):
+    if len(params) == 0:
+        with open('transform_params.pkl', 'rb') as f:
+            params = pickle.load(f)
 
     if UNIT_INTERVAL_MAP == 'sigmoid':
         return 1 - 1 / (1 + np.exp(params['a'] * predictions + params['b']))
@@ -135,23 +146,47 @@ def prediction_transform(predictions: np.array, **params):
         raise NotImplementedError
 
 
-def titration_ratio(guide: np.array, parent: np.array):
-    return 1 - np.clip(parent - guide, a_min=0.0, a_max=1.0)
+def get_on_target_predictions(transcripts: pd.DataFrame, model: tf.keras.Model, status_bar=None, status_text=None):
 
+    # loop over transcripts
+    predictions = pd.DataFrame()
+    for i, (index, row) in enumerate(transcripts.iterrows()):
 
-def predict_on_target(transcript_seq: str, model: tf.keras.Model):
+        # parse transcript sequence
+        target_seq, guide_seq, model_inputs = process_data(row[SEQ_COL])
 
-    # parse transcript sequence
-    target_seq, guide_seq, model_inputs = process_data(transcript_seq)
+        # get predictions
+        lfc_estimate = model.predict(model_inputs, batch_size=BATCH_SIZE_COMPUTE, verbose=False)
+        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],
+            GUIDE_COL: guide_seq,
+            SCORE_COL: scores})])
 
-    # get predictions
-    normalized_lfc = model.predict(model_inputs, batch_size=BATCH_SIZE_COMPUTE, verbose=False)
-    predictions = pd.DataFrame({'Guide': guide_seq, 'Normalized LFC': tf.squeeze(normalized_lfc).numpy()})
-    predictions = predictions.sort_values('Normalized LFC')
+        # progress update
+        percent_complete = 100 * min((i + 1) / len(transcripts), 1)
+        update_text = 'Evaluating on-target guides for each transcript: {:.2f}%'.format(percent_complete)
+        if status_bar:
+            status_text.text()
+            status_bar.progress(int(100 * min((i + 1) / len(transcripts), 1)))
+        print('\r' + update_text, end='')
+    print('')
 
     return predictions
 
 
+def top_guides_per_transcript(predictions: pd.DataFrame):
+
+    top_guides = pd.DataFrame()
+    for transcript in predictions[ID_COL].unique():
+        df = predictions.loc[predictions[ID_COL] == transcript]
+        df = df.sort_values(SCORE_COL, ascending=False).reset_index(drop=True).iloc[:NUM_TOP_GUIDES]
+        top_guides = pd.concat([top_guides, df])
+
+    return top_guides.reset_index(drop=True)
+
+
 def find_off_targets(top_guides: pd.DataFrame, status_bar, status_text):
 
     # load reference transcripts
@@ -171,7 +206,7 @@ def find_off_targets(top_guides: pd.DataFrame, status_bar, status_text):
         i += BATCH_SIZE_SCAN
 
         # find locations of off-targets
-        transcripts = one_hot_encode_sequence(df_batch['seq'].values.tolist(), add_context_padding=False)
+        transcripts = one_hot_encode_sequence(df_batch[SEQ_COL].values.tolist(), add_context_padding=False)
         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()
 
@@ -183,7 +218,7 @@ def find_off_targets(top_guides: pd.DataFrame, status_bar, status_text):
                 '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]],
+                'Target': df_batch[SEQ_COL].values[loc_off_targets[:, 0]],
                 'Mismatches': tf.gather_nd(num_mismatches, loc_off_targets).numpy().astype(int),
                 'Midpoint': loc_off_targets[:, 1],
             }).to_dict('records')
@@ -224,12 +259,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(model_inputs, batch_size=BATCH_SIZE_COMPUTE, verbose=False)
+    off_targets[SCORE_COL] = model.predict(model_inputs, batch_size=BATCH_SIZE_COMPUTE, verbose=False)
 
-    return off_targets.sort_values('Normalized LFC')
+    return off_targets.sort_values(SCORE_COL)
 
 
-def tiger_exhibit(transcripts: pd.DataFrame, status_bar=None, status_text=None, check_off_targets=False):
+def tiger_exhibit(transcripts: pd.DataFrame, run_mode: str, check_off_targets: bool, status_bar=None, status_text=None):
 
     # load model
     if os.path.exists('model'):
@@ -238,31 +273,30 @@ def tiger_exhibit(transcripts: pd.DataFrame, status_bar=None, status_text=None,
         print('no saved model!')
         exit()
 
-    # find top guides for each transcript
-    print('Finding top guides for each transcript')
-    on_target_predictions = pd.DataFrame(columns=['On-target ID', 'Guide', 'Normalized LFC'])
-    for i, (index, row) in enumerate(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]])
+    # evaluate all on-target guides per transcript
+    on_target_predictions = get_on_target_predictions(transcripts, tiger, status_bar, status_text)
 
-        # progress update
-        if status_bar:
-            status_text.text("Scanning for on-targets Percent complete: {:.2f}%".format(100 * min((i + 1) / len(transcripts), 1)))
-            status_bar.progress(int(100 * min((i + 1) / len(transcripts), 1)))
-        print('\rPercent complete: {:.2f}%'.format(100 * min((i + 1) / len(transcripts), 1)), end='')
-    print('')
-
-    # predict off-target effects for top guides
+    # initialize other outputs
     off_target_predictions = pd.DataFrame()
+
+    if run_mode == RUN_MODE_ALL_PM:
+        return on_target_predictions, off_target_predictions
+
+    elif run_mode == RUN_MODE_TITRATION:  # TODO: and titration candidates
+        on_target_predictions = top_guides_per_transcript(on_target_predictions)
+
+    else:
+        raise NotImplementedError
+
+    # check off-target effects for top guides
     if check_off_targets:
         off_targets = find_off_targets(on_target_predictions,  status_bar, status_text)
         off_target_predictions = predict_off_target(off_targets, model=tiger)
 
     # reverse guide sequences
-    on_target_predictions['Guide'] = on_target_predictions['Guide'].apply(lambda s: s[::-1])
+    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'] = off_target_predictions['Guide'].apply(lambda s: s[::-1])
+        off_target_predictions[GUIDE_COL] = off_target_predictions[GUIDE_COL].apply(lambda s: s[::-1])
 
     return on_target_predictions.reset_index(drop=True), off_target_predictions.reset_index(drop=True)
 
@@ -279,9 +313,11 @@ if __name__ == '__main__':
     # simple test case
     if args.simple_test:
         # first 50 from EIF3B-003's CDS
-        simple_test = pd.DataFrame(dict(id=['ManualEntry'], seq=['ATGCAGGACGCGGAGAACGTGGCGGTGCCCGAGGCGGCCGAGGAGCGCGC']))
-        simple_test.set_index('id', inplace=True)
-        df_on_target, df_off_target = tiger_exhibit(simple_test, check_off_targets=args.check_off_targets)
+        simple_test = pd.DataFrame({
+            ID_COL: ['ManualEntry'],
+            SEQ_COL: ['ATGCAGGACGCGGAGAACGTGGCGGTGCCCGAGGCGGCCGAGGAGCGCGC']})
+        simple_test.set_index(ID_COL, inplace=True)
+        df_on_target, df_off_target = tiger_exhibit(simple_test, check_off_targets=args.off_target)
         df_on_target.to_csv('on_target.csv')
         if args.check_off_targets:
             df_off_target.to_csv('off_target.csv')
@@ -306,7 +342,9 @@ 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], check_off_targets=args.check_off_targets)
+            df_on_target, df_off_target = tiger_exhibit(df_transcripts[idx:idx_stop],
+                                                        run_mode=RUN_MODE_TITRATION,
+                                                        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')