"""Utility functions for experiments.""" import logging import torch import os import re import random import esm import numpy as np import pandas as pd import random from analysis import utils as au from pytorch_lightning.utilities.rank_zero import rank_zero_only from data.residue_constants import restype_order from data.repr import get_pre_repr from data import utils as du from data.residue_constants import restype_atom37_mask from openfold.data import data_transforms from openfold.utils import rigid_utils from data.cal_trans_rotmats import cal_trans_rotmats class LengthDataset(torch.utils.data.Dataset): def __init__(self, samples_cfg): self._samples_cfg = samples_cfg # self._all_sample_seqs = ['LSEEEKKELEKEKRKKEIVEEVYKELKEEGKIKNLPKEEFMKKGLEILEKNEGKLKTDEEAKEELL', # 'PSPEELEKKAKEERKLRIVKEVGEELRKEGLIKDLPEEEFLKKGLEILKENEGKLKTEEEAKEALLEKFK', # 'PMLIEVLCRTDKVEELIKEIEELTKDKILEVKVEKIDENTVKIEIVLEKKEAAEKAAKWLSEVEGCEVIEMREV', # 'PTKITVLCPKSKVEELIEEIKEKTNDKILSVEVEEISPDSVKINIILETEEAALKAAEWLSEVEGCEVLEISEVELE', # 'SSSVKKRYKLTVKIEGASEEEFKELAELAKKLGEELGGLIEFEGDKENGKLTLLMESKEKAEKVGEALKEAGVKGGYTIEEFD', # 'VTSITKRYKLTVKITGASAAEFAALGAAAEAQGKALGGLLSFTADAANGTITVLMDTKEKAEKIGDALKALGVKGGYTISEFLEAD', # 'SKIEETKKKIAEGNYEEIKKLKEEIEKEKKKFEEEEKKEKEKAEELLKKDPEKGKKEKAKKEAEFEKKKKEYEEILKIIEKALKGKE', # 'SRIEEVKKQIEESDKEGVKELKKEILKEYEKFKKEAEKEKAEAEKLKKEDPEKGAKEEAELKKKHEEEKKEYEKILEIIEKRLKGAEEGK', # 'GEEALKLMEEELAAAKTEEAKKFMEGLKKMIEEIAKAMATGDPEVIEEGKKRLLEWGKEVGEKGKKEGNPFLIELEKIIEYMAEGEIEEGLKKLMEFLKKKR', # 'GAEALALMDEMLAAAKREEDKAFYARLRELVRRLAAALATGDPAVLAAGRAEAAAEGDALGAEGRATGDPFLVELAAIVAALAAGTPEEGLAALAAFLRAKAAAR'] # self._all_filename = ['P450'] * 250 # self._all_sample_seqs = [('GKLPPGPSPLPVLGNLLQMDRKGLLRSFLRLREKYGDVFTVYLGSRPVVVLCGTDAIREALVDQAEAFSGRGKIAVVDPIFQGYGVIFANGERWRALRRFSLATMRDFGMGKRSVEERIQEEARCLVEELRKSKGALLDNTLLFHSITSNIICSIVFGKRFDYKDPVFLRLLDLFFQSFSLISSFSSQVFELFSGFLKYFPGTHRQIYRNLQEINTFIGQSVEKHRATLDPSNPRDFIDVYLLRMEKDKSDPSSEFHHQNLILTVLSLFFAGTETTSTTLRYGFLLMLKYPHVTERVQKEIEQVIGSHRPPALDDRAKMPYTDAVIHEIQRLGDLIPFGVPHTVTKDTQFRGYVIPKNTEVFPVLSSALHDPRYFETPNTFNPGHFLDANGALKRNEGFMPFSLGKRICLGEGIARTELFLFFTTILQNFSIASPVPPEDIDLTPRESGVGNVPPSYQIRFLARH',0)] * 250 validcsv = pd.read_csv(self._samples_cfg.validset_path) self._all_sample_seqs = [] self._all_filename = [] prob_num = 500 exp_prob = np.exp([-prob/prob_num*2 for prob in range(prob_num)]).cumsum() exp_prob = exp_prob/np.max(exp_prob) for idx in range(len(validcsv['seq'])): # if idx >= 0 and idx < 15: # if idx >= 15 and idx < 30: # if idx >= 30 and idx < 45: # if idx >= 45 and idx < 60: # if idx >= 60 and idx < 75: # if idx >= 75 and idx < 90: # if idx >= 90 and idx < 105: # pass # else: # continue # if not re.search('2wsi_A',validcsv['file'][idx]): # continue self._all_filename += [validcsv['file'][idx]] * self._samples_cfg.sample_num for batch_idx in range(self._samples_cfg.sample_num): rand = random.random() for prob in range(prob_num): if rand < exp_prob[prob]: energy = torch.tensor(prob/prob_num) break self._all_sample_seqs += [(validcsv['seq'][idx], energy)] self._all_sample_ids = self._all_sample_seqs # Load ESM-2 model self.device_esm=f'cuda:{torch.cuda.current_device()}' self.model_esm2, self.alphabet = esm.pretrained.esm2_t33_650M_UR50D() self.batch_converter = self.alphabet.get_batch_converter() self.model_esm2.eval().cuda(self.device_esm) # disables dropout for deterministic results self.model_esm2.requires_grad_(False) self._folding_model = esm.pretrained.esmfold_v1().eval() self._folding_model = self._folding_model.to(self.device_esm) self.esm_savepath = self._samples_cfg.esm_savepath self.device_esm=f'cuda:{torch.cuda.current_device()}' self._folding_model = esm.pretrained.esmfold_v1().eval() self._folding_model.requires_grad_(False) self._folding_model.to(self.device_esm) def run_folding(self, sequence, save_path): """Run ESMFold on sequence.""" with torch.no_grad(): output = self._folding_model.infer_pdb(sequence) self._folding_model.to("cpu") with open(save_path, "w") as f: f.write(output) return output def __len__(self): return len(self._all_sample_ids) def __getitem__(self, idx): seq, energy = self._all_sample_ids[idx] aatype = torch.tensor([restype_order[s] for s in seq]) num_res = len(aatype) node_repr_pre, pair_repr_pre = get_pre_repr(aatype, self.model_esm2, self.alphabet, self.batch_converter, device = self.device_esm) # (B,L,d_node_pre=1280), (B,L,L,d_edge_pre=20) node_repr_pre = node_repr_pre[0].cpu() pair_repr_pre = pair_repr_pre[0].cpu() motif_mask = torch.ones(aatype.shape) save_path = os.path.join(self.esm_savepath, "esm_" + self._all_filename[idx] + ".pdb") if not os.path.exists(save_path): seq_string = seq with torch.no_grad(): output = self._folding_model.infer_pdb(seq_string) with open(save_path, "w") as f: f.write(output) trans_esmfold, rotmats_esmfold = cal_trans_rotmats(save_path) batch = { 'filename':self._all_filename[idx], 'trans_esmfold': trans_esmfold, 'rotmats_esmfold': rotmats_esmfold, 'motif_mask': motif_mask, 'res_mask': torch.ones(num_res).int(), 'num_res': num_res, 'energy': energy, 'aatype': aatype, 'seq': seq, 'node_repr_pre': node_repr_pre, 'pair_repr_pre': pair_repr_pre, } return batch def save_traj( sample: np.ndarray, bb_prot_traj: np.ndarray, x0_traj: np.ndarray, diffuse_mask: np.ndarray, output_dir: str, aatype = None, index=0, ): """Writes final sample and reverse diffusion trajectory. Args: bb_prot_traj: [T, N, 37, 3] atom37 sampled diffusion states. T is number of time steps. First time step is t=eps, i.e. bb_prot_traj[0] is the final sample after reverse diffusion. N is number of residues. x0_traj: [T, N, 3] x_0 predictions of C-alpha at each time step. aatype: [T, N, 21] amino acid probability vector trajectory. res_mask: [N] residue mask. diffuse_mask: [N] which residues are diffused. output_dir: where to save samples. Returns: Dictionary with paths to saved samples. 'sample_path': PDB file of final state of reverse trajectory. 'traj_path': PDB file os all intermediate diffused states. 'x0_traj_path': PDB file of C-alpha x_0 predictions at each state. b_factors are set to 100 for diffused residues and 0 for motif residues if there are any. """ # Write sample. diffuse_mask = diffuse_mask.astype(bool) # (B,L) sample_path = os.path.join(output_dir, 'sample_'+str(index)+'.pdb') prot_traj_path = os.path.join(output_dir, 'bb_traj_'+str(index)+'.pdb') x0_traj_path = os.path.join(output_dir, 'x0_traj_'+str(index)+'.pdb') # Use b-factors to specify which residues are diffused. b_factors = np.tile((diffuse_mask * 100)[:, None], (1, 37)) sample_path = au.write_prot_to_pdb( sample, sample_path, b_factors=b_factors, no_indexing=True, aatype=aatype, ) prot_traj_path = au.write_prot_to_pdb( bb_prot_traj, prot_traj_path, b_factors=b_factors, no_indexing=True, aatype=aatype, ) x0_traj_path = au.write_prot_to_pdb( x0_traj, x0_traj_path, b_factors=b_factors, no_indexing=True, aatype=aatype ) return { 'sample_path': sample_path, 'traj_path': prot_traj_path, 'x0_traj_path': x0_traj_path, } def get_pylogger(name=__name__) -> logging.Logger: """Initializes multi-GPU-friendly python command line logger.""" logger = logging.getLogger(name) # this ensures all logging levels get marked with the rank zero decorator # otherwise logs would get multiplied for each GPU process in multi-GPU setup logging_levels = ("debug", "info", "warning", "error", "exception", "fatal", "critical") for level in logging_levels: setattr(logger, level, rank_zero_only(getattr(logger, level))) return logger def flatten_dict(raw_dict): """Flattens a nested dict.""" flattened = [] for k, v in raw_dict.items(): if isinstance(v, dict): flattened.extend([ (f'{k}:{i}', j) for i, j in flatten_dict(v) ]) else: flattened.append((k, v)) return flattened