File size: 2,679 Bytes
fae1173 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | import ablang
import numpy as np
import pandas as pd
import scipy.special
heavy_ablang = ablang.pretrained("heavy")
heavy_ablang.freeze()
light_ablang = ablang.pretrained("light")
light_ablang.freeze()
ab_dict = dict(
cr6261= ('EVQLVESGAEVKKPGSSVKVSCKASGGPFRSYAISWVRQAPGQGPEWMGGIIPIFGTTKYAPKFQGRVTITADDFAGTVYMELSSLRSEDTAMYYCAKHMGYQVRETMDVWGKGTTVTVSS',
'QSVLTQPPSVSAAPGQKVTISCSGSSSNIGNDYVSWYQQLPGTAPKLLIYDNNKRPSGIPDRFSGSKSGTSATLGITGLQTGDEANYYCATWDRRPTAYVVFGGGTKLTVL'),
cr9114= ('QVQLVQSGAEVKKPGSSVKVSCKSSGGTSNNYAISWVRQAPGQGLDWMGGISPIFGSTAYAQKFQGRVTISADIFSNTAYMELNSLTSEDTAVYFCARHGNYYYYSGMDVWGQGTTVTVSS',
'QSALTQPPAVSGTPGQRVTISCSGSDSNIGRRSVNWYQQFPGTAPKLLIYSNDQRPSVVPDRFSGSKSGTSASLAISGLQSEDEAEYYCAAWDDSLKGAVFGGGTQLTVL'),
g6 = ('EVQLVESGGGLVQPGGSLRLSCAASGFTISDYWIHWVRQAPGKGLEWVAGITPAGGYTYYADSVKGRFTISADTSKNTAYLQMNSLRAEDTAVYYCARFVFFLPYAMDYWGQGTLVTV',
'DIQMTQSPSSLSASVGDRVTITCRASQDVSTAVAWYQQKPGKAPKLLIYSASFLYSGVPSRFSGSGSGTDFTLTISSLQPEDFATYYCQQSYTTPPTFGQGTKVEIK')
)
def eval_ablang(s, ab, chain):
if chain == 'hc':
fname = 'output/ab_mutagenesis_expts/'+ab+'/' + ab +'_hc_ablangScores.csv'
log_likelihoods = heavy_ablang(s, mode = 'likelihood')[0][1:-1]
alphabet = heavy_ablang.tokenizer.vocab_to_aa
elif chain == 'lc':
fname = 'output/ab_mutagenesis_expts/'+ab+'/' + ab +'_lc_ablangScores.csv'
log_likelihoods = light_ablang(s, mode = 'likelihood')[0][1:-1]
alphabet = light_ablang.tokenizer.vocab_to_aa
assert (log_likelihoods.shape)[0] == len(s)
filt_alphabet = {key: value for key, value in alphabet.items() if value.isalpha()}
log_likelihood_ratio = []
for i,res_log_likelihoods in enumerate(log_likelihoods):
wt_res = s[i]
wt_index = list(filt_alphabet.values()).index(wt_res)
wt_log_likelihood = res_log_likelihoods[wt_index]
log_likelihood_ratio.extend(res_log_likelihoods-wt_log_likelihood)
res_order = [alphabet[key] for key in range(1, 21)] #extract order of residues in likelihood
mt = res_order * len(s)
wt = [char for char in s for i in range(len(res_order))]
pos = [i+1 for i in range(len(s)) for j in range(len(res_order))]
data = {
'pos': pos,
'wt': wt,
'mt': mt,
'log_likelihood' : log_likelihoods.flatten(),
'log_likelihood_ratio' : log_likelihood_ratio,
}
df = pd.DataFrame(data)
df.to_csv(fname, index = False)
def main():
for ab in ab_dict:
for s,chain in zip(ab_dict[ab], ('hc','lc')):
eval_ablang(s, ab, chain)
if __name__ == '__main__':
main() |