import torch
import pandas as pd
import seaborn as sns
import matplotlib.pylab as plt
from mlm import *
import matplotlib.pyplot as plt
from predict import predict
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model = torch.load('./trans_best.pt', map_location=device)
model.eval()
def read_vocab(vocab_file):
vocab_dict = {}
with open(vocab_file, 'r') as vocab_file:
for i, line in enumerate(vocab_file):
vocab = line.replace("\n", "")
vocab_dict[vocab] = i
return vocab_dict
vocab = read_vocab('vocab.txt')
def tokenizer(seq):
x = torch.zeros(len(seq)+2)
x[0] = 2
for i, aa in enumerate(seq):
w = vocab.get(aa)
if w:
x[i+1] = w
else:
x[i+1] = 1
x[-1] = 3
return x.long()
def check_heatmap():
inp = tokenizer(raw_sequence).view(1, -1)
map = []
for i, w in enumerate(raw_sequence):
wild = vocab[w]
inp[0][i] = 4
with torch.no_grad():
y = model(inp.to(device))
map.append(list(y[0].cpu()[i].numpy() - y[0].cpu()[i].numpy()[wild])[5:])
inp = tokenizer(raw_sequence).view(1, -1)
return map
map = check_heatmap()
plot_x = []
plot_y = []
for x, y in zip([i for i in range(len(map))], np.array(map).mean(axis=-1)):
if x+36 < 280: # we remove the unstructured C-terminal
plot_x.append(x+36)
plot_y.append(y)
with open("mutations-BhrPETase-new.fasta", 'a') as ofile:
x = (top10['pos']-36).to_numpy()
inds = []
scores = []
for i in x:
masked_seq = raw_sequence[:i] + "" + raw_sequence[i+1:]
predicted_sequence, predicted_scores, mask_idx = predict(masked_seq)
predicted_sequence = predicted_sequence[5:-5:]
res = predicted_sequence[i]
inds.append(mask_idx)
scores.append(predicted_scores[4:-4:])
ofile.write(f"{raw_sequence[i]}{i+1}{res}\n")
ofile.write(f"{raw_sequence[:i] + res + raw_sequence[i+1:]}\n")
ofile.write(f"Old: {raw_sequence[i]}\n")
ofile.write(f"New: {predicted_sequence[i]}\n\n")
amino_acids = []
for i in range(len(scores[0])):
amino_acids.append(alphabet.get_tok(4+i))
print(np.array(scores[0]).argmax(axis=0))
# ploting the heatmap of the scores
fig, ax = plt.subplots(figsize=(8,6))
sns.heatmap(np.array(scores), cmap='coolwarm', ax=ax, square=True, linewidth=0.5, cbar_kws={"shrink": 0.5})
# Increase cell size
ax.set_aspect(1.5)
# Rename the x_axis labels using the indices of the mutations
ax.set_yticklabels([f"{i+2} ({raw_sequence[i]})" for i in inds])
ax.set_xticklabels(amino_acids, rotation=0, ha="right", rotation_mode="anchor")
ax.set_title("Predicted Scores of Mutations on Top 10 Positions on BhrPETase")
ax.set_xlabel("Amino Acid Mutations")
ax.set_ylabel("Position on BhrPETase")
# Control the size of the color bar
plt.setp(ax.get_yticklabels(), rotation=0, ha="right",
rotation_mode="anchor")
plt.savefig("heatmap-BhrPETase-new.png", dpi=300)
