load_transcripts function

tiger.py

@@ -1,3 +1,4 @@
+import argparse
 import os
 import gzip
 import numpy as np
@@ -13,6 +14,7 @@ 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
+REFERENCE_TRANSCRIPTS = ('gencode.v19.pc_transcripts.fa.gz', 'gencode.v19.lncRNA_transcripts.fa.gz')
 
 # configure GPUs
 for gpu in tf.config.list_physical_devices('GPU'):
@@ -21,6 +23,30 @@ if len(tf.config.list_physical_devices('GPU')) > 0:
     tf.config.experimental.set_visible_devices(tf.config.list_physical_devices('GPU')[0], 'GPU')
 
 
+def load_transcripts(fasta_files):
+
+    # load all transcripts from fasta files into a DataFrame
+    transcripts = pd.DataFrame()
+    for file in 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'])
+            else:
+                df = pd.DataFrame([(t.id, str(t.seq)) for t in SeqIO.parse(f, 'fasta')], columns=['id', 'seq'])
+        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
+
+    return transcripts
+
+
 def sequence_complement(sequence: list):
     return [''.join([NUCLEOTIDE_COMPLEMENT[nt] for nt in list(seq)]) for seq in sequence]
 
@@ -87,14 +113,7 @@ def predict_on_target(transcript_seq: str, model: tf.keras.Model):
 def find_off_targets(guides, batch_size=500):
 
     # load reference transcripts
-    df_transcripts = pd.DataFrame()
-    for transcripts in ['gencode.v19.pc_transcripts.fa.gz', 'gencode.v19.lncRNA_transcripts.fa.gz']:
-        with gzip.open(os.path.join('transcripts', transcripts), 'rt') as file:
-            df = pd.DataFrame([(t.id, str(t.seq)) for t in SeqIO.parse(file, 'fasta')], columns=['id', 'seq'])
-        df_transcripts = pd.concat([df_transcripts, df])
-    df_transcripts['id'] = df_transcripts['id'].apply(lambda s: s.split('|')[0])
-    df_transcripts.set_index('id', inplace=True)
-    assert not df_transcripts.index.has_duplicates
+    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)
@@ -105,9 +124,9 @@ def find_off_targets(guides, batch_size=500):
     i = 0
     print('Scanning for off-targets')
     df_off_targets = pd.DataFrame()
-    while i < len(df_transcripts):
+    while i < len(reference_transcripts):
         # select batch
-        df_batch = df_transcripts.iloc[i:min(i + batch_size, len(df_transcripts))]
+        df_batch = reference_transcripts.iloc[i:min(i + batch_size, len(reference_transcripts))]
         i += batch_size
 
         # find and log off-targets
@@ -124,7 +143,7 @@ def find_off_targets(guides, batch_size=500):
         })])
 
         # progress update
-        print('\rPercent complete: {:.2f}%'.format(100 * min(i / len(df_transcripts), 1)), end='')
+        print('\rPercent complete: {:.2f}%'.format(100 * min(i / len(reference_transcripts), 1)), end='')
     print('')
 
     # trim transcripts to targets