| """Helpers for working with DNA/RNA data""" |
|
|
| import os |
| import numpy as np |
|
|
| |
| dna_vocab = {"A":0, |
| "C":1, |
| "G":2, |
| "T":3, |
| "*":4} |
| rna_vocab = {"A":0, |
| "C":1, |
| "G":2, |
| "U":3, |
| "*":4} |
| dna_nt_only_vocab = {k:v for k,v in dna_vocab.items() if k in "ACGT"} |
| rna_nt_only_vocab = {k:v for k,v in dna_vocab.items() if k in "ACGU"} |
|
|
| rev_dna_vocab = {v:k for k,v in dna_nt_only_vocab.items()} |
| rev_rna_vocab = {v:k for k,v in rna_vocab.items()} |
| rev_dna_nt_only_vocab = {v:k for k,v in dna_vocab.items()} |
| rev_rna_nt_only_vocab = {v:k for k,v in rna_nt_only_vocab.items()} |
|
|
| def get_vocab(vocab_name, vocab_order=None): |
| if vocab_name=="dna": |
| charmap = dna_vocab |
| elif vocab_name=="rna": |
| charmap = rna_vocab |
| elif vocab_name=="dna_nt_only": |
| charmap = dna_nt_only_vocab |
| elif vocab_name=="rna_nt_only": |
| charmap = rna_nt_only_vocab |
| else: |
| raise Exception("Unknown vocabulary name.") |
|
|
| if vocab_order: |
| if set(vocab_order) != set(charmap): |
| raise ValueError("Provided `vocab` and `vocab_order` arguments are not compatible") |
| else: |
| charmap = {c: idx for idx, c in enumerate(vocab_order)} |
|
|
| rev_charmap = {v: k for k, v in charmap.items()} |
| return charmap, rev_charmap |
|
|
| def _process_line(line, max_len, charmap): |
| chars = line.strip() |
| I = np.eye(len(charmap)) |
| try: |
| base = [I[charmap[c]] for c in chars] |
| if len(chars) < max_len: |
| extra = [] |
| if "*" in charmap: extra = [I[charmap["*"]]] * (max_len - len(chars)) |
| else: |
| extra = [] |
| arr = np.array(base + extra) |
| except: |
| raise Exception("Unable to process line: {}".format(chars)) |
| return np.expand_dims(arr, 0) |
|
|
| def load(data_loc, max_seq_len=None, vocab="dna", vocab_order=None, data_start_line=0, scores=False, valid=False, test=False, filenames=None, annotate=False): |
| charmap, _ = get_vocab(vocab, vocab_order) |
|
|
| if filenames: |
| if type(filenames)==list: |
| seq_filenames = filenames |
| else: |
| seq_filenames = [filenames] |
| else: |
| seq_filenames = ["train_data.txt"] |
| if valid: |
| seq_filenames.append("valid_data.txt") |
| if test: |
| seq_filenames.append("test_data.txt") |
|
|
| data = [] |
| for name in seq_filenames: |
| with open(os.path.join(data_loc, name)) as f: |
| lines = f.readlines() |
| lines = lines[data_start_line:] |
| if not max_seq_len: |
| print("Warning: max_seq_len not provided. Inferring size from data.") |
| max_seq_len = len(max(lines, key=len)) - 1 |
| lines = [_process_line(l, max_seq_len, charmap) for l in lines] |
| data.append(np.vstack(lines)) |
| |
| if scores: |
| score_filename = name.split("_")[0] + "_vals.txt" |
| scores_array = np.loadtxt(os.path.join(data_loc, score_filename), skiprows=data_start_line) |
| data.append(scores_array) |
|
|
| if annotate: |
| ann_filename = name.split("_")[0] + "_ann.txt" |
| ann_array = np.loadtxt(os.path.join(data_loc, ann_filename), skiprows=data_start_line) |
| ann_array = np.expand_dims(ann_array, 2) |
| data.append(ann_array) |
| if type(data) == list and len(data) == 1: |
| data = data[0] |
| return data |
| |
