| import os, string
|
| import numpy as np
|
| import jax
|
| import jax.numpy as jnp
|
|
|
| ALPHABET = list("ARNDCQEGHILKMFPSTWYV-")
|
|
|
| def parse_fasta(filename, a3m=False, stop=100000):
|
| '''function to parse fasta file'''
|
|
|
| if a3m:
|
|
|
|
|
| rm_lc = str.maketrans(dict.fromkeys(string.ascii_lowercase))
|
|
|
| header, sequence = [],[]
|
| lines = open(filename, "r")
|
| for line in lines:
|
| line = line.rstrip()
|
| if len(line) > 0:
|
| if line[0] == ">":
|
| if len(header) == stop:
|
| break
|
| else:
|
| header.append(line[1:])
|
| sequence.append([])
|
| else:
|
| if a3m: line = line.translate(rm_lc)
|
| else: line = line.upper()
|
| sequence[-1].append(line)
|
| lines.close()
|
| sequence = [''.join(seq) for seq in sequence]
|
|
|
| return header, sequence
|
|
|
| def mk_msa(seqs):
|
| '''one hot encode msa'''
|
| states = len(ALPHABET)
|
| a2n = {a:n for n,a in enumerate(ALPHABET)}
|
| msa_ori = np.array([[a2n.get(aa, states-1) for aa in seq] for seq in seqs])
|
| return np.eye(states)[msa_ori]
|
|
|
| def get_eff(msa, eff_cutoff=0.8):
|
| '''compute weight per sequence'''
|
| if msa.shape[0] > 10000:
|
|
|
| msa = msa.argmax(-1)
|
| def get_w(seq): return 1/((seq==msa).mean(-1) > eff_cutoff).sum()
|
| return jax.lax.scan(lambda _,x:(_,get_w(x)),None,msa,unroll=2)[1]
|
| else:
|
|
|
| msa_ident = jnp.tensordot(msa,msa,[[1,2],[1,2]])/msa.shape[1]
|
| return 1/(msa_ident >= eff_cutoff).sum(-1)
|
|
|
| def ar_mask(order, diag=True):
|
| '''compute autoregressive mask, given order of positions'''
|
| L = order.shape[0]
|
| r = order[::-1].argsort()
|
| tri = jnp.triu(jnp.ones((L,L)),k=not diag)
|
| return tri[r[None,:],r[:,None]]
|
|
|
