import numpy as np import timeit import torch from utils import rna2vec from transformers import AutoTokenizer #Node class Node: #init def __init__(self, letter="", parent=None, root=False, last=False, depth=0, states=8): self.exploitation_score = 0 # Exploitaion score self.visits = 1 #How many visits self.letter = letter #This node's letter self.parent = parent #This node's parent node self.states = states #How many states in node self.children = np.array([None for _ in range(self.states)]) #This node's children self.children_stat = np.zeros(self.states, dtype=bool) #Which stat are expanded self.root = root # Is root? boolean self.last = last # Is last node? self.depth = depth # My depth self.letters =["A_", "C_", "G_", "T_", "_A", "_C", "_G", "_T"] #next_node def next_node(self, child=0): #Return next node assert self.children_stat[child] == True, "No child in here." return self.children[child] #back_parent def back_parent(self): #Go back to parent return self.parent, letters_map[self.letter] #generate_child def generate_child(self, child=0, last=False): #Generate child assert self.children_stat[child] == False, "Already tree generated child at here" self.children[child] = Node(letter=self.letters[child], parent=self, last=last, depth=self.depth+1, states=self.states) #New node self.children_stat[child] = True #Stat = True return self.children[child] #backpropagation def backpropagation(self, score=0): self.visits += 1 # +1 to visit if self.root == True: # if root, then stop return self.exploitation_score else: self.exploitation_score += score #Add score to exploitation score return self.parent.backpropagation(score=score) #Backpropagation to parent node #UCT def UCT(self): return (self.exploitation_score / self.visits) + np.sqrt(np.log(self.parent.visits) / (2 * self.visits)) #UCT score #MCTS class MCTS: def __init__(self, target_encoded, depth=20, iteration=1000, states=8, target_protein="", device='cpu', esm_alphabet=None): self.states = states #How many states self.root = Node(letter="", parent=None, root=True, last=False, states=self.states) #root node self.depth = depth #Maximum depth self.iteration = iteration #iteration for expand self.target_protein = target_protein #target protein's amino acid sequence self.device = device self.encoded_targetprotein = target_encoded self.base = "" self.candidate = "" self.letters =["A_", "C_", "G_", "T_", "_A", "_C", "_G", "_T"] self.esm_alphabet = esm_alphabet self.nt_tokenizer = AutoTokenizer.from_pretrained("InstaDeepAI/nucleotide-transformer-v2-50m-multi-species", trust_remote_code=True) def make_candidate(self, classifier): now = self.root n = 0 # rounds start_time = timeit.default_timer() #timer start while len(self.base) < self.depth * 2: #If now is last node, then stop n += 1 print(n, "round start!!!") for _ in range(self.iteration): now = self.select(classifier, now=now) #Select & Expand terminate_time = timeit.default_timer() time = terminate_time-start_time base = self.find_best_subsequence() #Find best subsequence self.base = base # print("best subsequence:", base) # print("Depth:", int(len(base)/2)) # print("%02d:%02d:%2f" % ((time//3600), (time//60)%60, time%60)) # print("=" * 80) self.root = Node(letter="", parent=None, root=True, last=False, states=self.states, depth=len(self.base)/2) now = self.root self.candidate = self.base return self.candidate #selection def select(self, classifier, now=None): if now.depth == self.depth: #If last node, then stop return self.root next_node = 0 if np.sum(now.children_stat) == self.states: #If every child is expanded, then go to best child best = 0 for i in range(self.states): if best < now.children[i].UCT(): next_node = i best = now.children[i].UCT() else: #If not, then random next_node = np.random.randint(0, self.states) if now.children_stat[next_node] == False: #If selected child is not expanded, then expand and simulate next_node = self.expand(classifier, child=next_node, now=now) return self.root #start iteration at this node return now.next_node(child=next_node) #expand def expand(self, classifier, child=None, now=None): last = False if now.depth == (self.depth-1): #If depth of this node is maximum depth -1, then next node is last last = True expanded_node = now.generate_child(child=child, last=last) #Expand score = self.simulate(classifier, target=expanded_node) #Simulate expanded_node.backpropagation(score=score) #Backporpagation return child #simulate def simulate(self, classifier, target=None): now = target #Target node sim_seq = "" while now.root != True: #Parent's letters sim_seq = now.letter + sim_seq now = now.parent sim_seq = self.base + sim_seq for i in range((self.depth * 2) - len(sim_seq)): #Random child letters r = np.random.randint(0,self.states) sim_seq += self.letters[r] sim_seq = self.reconstruct(sim_seq) scores = [] classifier.eval().to('cuda') with torch.no_grad(): sim_seq = np.array([sim_seq]) apta_toks = self.nt_tokenizer.batch_encode_plus(sim_seq, return_tensors='pt', padding='max_length', max_length=275)['input_ids'] apta_attention_mask = apta_toks != self.nt_tokenizer.pad_token_id prot_attention_mask = self.encoded_targetprotein != self.esm_alphabet.padding_idx score, _, _, _ = classifier(apta_toks.to('cuda'), self.encoded_targetprotein.to('cuda'), apta_attention_mask.to('cuda'), prot_attention_mask.to('cuda')) return score #recommend def get_candidate(self): return self.reconstruct(self.candidate) def find_best_subsequence(self): now = self.root stop = False base = self.base for _ in range((self.depth*2) - len(base)): best = 0 next_node = 0 for j in range(self.states): if now.children_stat[j] == True: if best < now.children[j].UCT(): next_node = j best = now.children[j].UCT() now = now.next_node(child=next_node) base += now.letter # if current node has no expanded children, stop reconstructing. if np.sum(now.children_stat) == 0: break return base #reconstruct def reconstruct(self, seq=""): r_seq = "" for i in range(0, len(seq), 2): if seq[i] == '_': r_seq = r_seq + seq[i+1] else: r_seq = seq[i] + r_seq return r_seq def reset(self): self.base = "" self.candidate = "" self.root = Node(letter="", parent=None, root=True, last=False, states=self.states)