fasta directory optimizations

tiger.py

@@ -14,7 +14,9 @@ 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
+BATCH_SIZE_COMPUTE = 500
+BATCH_SIZE_SCAN = 20
+BATCH_SIZE_TRANSCRIPTS = 50
 
 # configure GPUs
 for gpu in tf.config.list_physical_devices('GPU'):
@@ -42,7 +44,7 @@ def load_transcripts(fasta_files):
     # set index
     transcripts['id'] = transcripts['id'].apply(lambda s: s.split('|')[0])
     transcripts.set_index('id', inplace=True)
-    assert not transcripts.index.has_duplicates
+    assert not transcripts.index.has_duplicates, "duplicate transcript ID's detected"
 
     return transcripts
 
@@ -72,7 +74,7 @@ def one_hot_encode_sequence(sequence: list, add_context_padding: bool = False):
         sequence = tf.concat([pad_5p, sequence, pad_3p], axis=1)
 
     # one-hot encode
-    sequence = tf.one_hot(sequence, depth=4)
+    sequence = tf.one_hot(sequence, depth=4, dtype=tf.float16)
 
     return sequence
 
@@ -103,7 +105,7 @@ def predict_on_target(transcript_seq: str, model: tf.keras.Model):
     target_seq, guide_seq, model_inputs = process_data(transcript_seq)
 
     # get predictions
-    normalized_lfc = model.predict_step(model_inputs)
+    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')
 
@@ -118,7 +120,6 @@ def find_off_targets(top_guides: pd.DataFrame):
     # one-hot encode guides to form a filter
     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
@@ -126,43 +127,48 @@ def find_off_targets(top_guides: 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_SCAN, len(reference_transcripts))]
+        i += BATCH_SIZE_SCAN
 
-        # find and log off-targets
+        # find locations of 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()
-        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],
-        })])
+
+        # off-targets discovered
+        if len(loc_off_targets) > 0:
+
+            # log off-targets
+            dict_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],
+            }).to_dict('records')
+
+            # trim transcripts to targets
+            for row in dict_off_targets:
+                start_location = row['Midpoint'] - (GUIDE_LEN // 2)
+                if start_location < CONTEXT_5P:
+                    row['Target'] = row['Target'][0:GUIDE_LEN + CONTEXT_3P]
+                    row['Target'] = 'N' * (TARGET_LEN - len(row['Target'])) + row['Target']
+                elif start_location + GUIDE_LEN + CONTEXT_3P > len(row['Target']):
+                    row['Target'] = row['Target'][start_location - CONTEXT_5P:]
+                    row['Target'] = row['Target'] + 'N' * (TARGET_LEN - len(row['Target']))
+                else:
+                    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]
+
+            # append new off-targets
+            off_targets = pd.concat([off_targets, pd.DataFrame(dict_off_targets)])
 
         # progress update
         print('\rPercent complete: {:.2f}%'.format(100 * min(i / len(reference_transcripts), 1)), end='')
     print('')
 
-    # trim transcripts to targets
-    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:
-            row['Target'] = row['Target'][0:GUIDE_LEN + CONTEXT_3P]
-            row['Target'] = 'N' * (TARGET_LEN - len(row['Target'])) + row['Target']
-        elif start_location + GUIDE_LEN + CONTEXT_3P > len(row['Target']):
-            row['Target'] = row['Target'][start_location - CONTEXT_5P:]
-            row['Target'] = row['Target'] + 'N' * (TARGET_LEN - len(row['Target']))
-        else:
-            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]
-    off_targets = pd.DataFrame(dict_off_targets)
-
     return off_targets
 
 
@@ -175,7 +181,7 @@ 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, verbose=False)
+    off_targets['Normalized LFC'] = model.predict(model_inputs, batch_size=BATCH_SIZE_COMPUTE, verbose=False)
 
     return off_targets.sort_values('Normalized LFC')
 
@@ -190,12 +196,17 @@ def tiger_exhibit(transcripts: pd.DataFrame):
         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 index, row in transcripts.iterrows():
+    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]])
 
+        # progress update
+        print('\rPercent complete: {:.2f}%'.format(100 * min((i + 1) / len(transcripts), 1)), end='')
+    print('')
+
     # predict off-target effects for top guides
     off_targets = find_off_targets(on_target_predictions)
     off_target_predictions = predict_off_target(off_targets, model=tiger)
@@ -220,16 +231,26 @@ if __name__ == '__main__':
         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)
+    # directory of fasta files
+    elif args.fasta_path is not None and os.path.exists(args.fasta_path):
+
+        # load transcripts
+        df_transcripts = load_transcripts([os.path.join(args.fasta_path, f) for f in os.listdir(args.fasta_path)])
+
+        # process in batches
+        df_on_target = pd.DataFrame()
+        df_off_target = pd.DataFrame()
+        batch = 1
+        num_batches = len(df_transcripts) // BATCH_SIZE_TRANSCRIPTS
+        num_batches += (len(df_transcripts) % BATCH_SIZE_TRANSCRIPTS > 0)
+        for t in range(0, len(df_transcripts), BATCH_SIZE_TRANSCRIPTS):
+            print('Batch {:d} of {:d}'.format(batch, num_batches))
+            t_stop = min(t + BATCH_SIZE_TRANSCRIPTS, len(df_transcripts))
+            df_on_target_new, df_off_target_new = tiger_exhibit(df_transcripts[t:t_stop])
+            df_on_target = pd.concat([df_on_target, df_on_target_new])
+            df_off_target = pd.concat([df_off_target, df_off_target_new])
+            batch += 1
+
+        # save results
+        df_on_target.to_csv('on_target.csv')
+        df_off_target.to_csv('off_target.csv')