Spaces:
Sleeping
Sleeping
| """ | |
| https://github.com/ProteinDesignLab/protpardelle | |
| License: MIT | |
| Author: Alex Chu | |
| Configs and convenience functions for wrapping the model sample() function. | |
| """ | |
| import argparse | |
| import time | |
| from typing import Optional, Tuple | |
| import torch | |
| from torchtyping import TensorType | |
| from core import residue_constants | |
| from core import utils | |
| import diffusion | |
| def default_backbone_sampling_config(): | |
| config = argparse.Namespace( | |
| n_steps=500, | |
| s_churn=200, | |
| step_scale=1.2, | |
| sidechain_mode=False, | |
| noise_schedule=lambda t: diffusion.noise_schedule(t, s_max=80, s_min=0.001), | |
| ) | |
| return config | |
| def default_allatom_sampling_config(): | |
| noise_schedule = lambda t: diffusion.noise_schedule(t, s_max=80, s_min=0.001) | |
| stage2 = argparse.Namespace( | |
| apply_cond_proportion=1.0, | |
| n_steps=200, | |
| s_churn=100, | |
| step_scale=1.2, | |
| sidechain_mode=True, | |
| skip_mpnn_proportion=1.0, | |
| noise_schedule=noise_schedule, | |
| ) | |
| config = argparse.Namespace( | |
| n_steps=500, | |
| s_churn=200, | |
| step_scale=1.2, | |
| sidechain_mode=True, | |
| skip_mpnn_proportion=0.6, | |
| use_fullmpnn=False, | |
| use_fullmpnn_for_final=True, | |
| anneal_seq_resampling_rate="linear", | |
| noise_schedule=noise_schedule, | |
| stage_2=stage2, | |
| ) | |
| return config | |
| def draw_backbone_samples( | |
| model: torch.nn.Module, | |
| seq_mask: TensorType["b n", float] = None, | |
| n_samples: int = None, | |
| sample_length_range: Tuple[int] = (50, 512), | |
| pdb_save_path: Optional[str] = None, | |
| return_aux: bool = False, | |
| return_sampling_runtime: bool = False, | |
| **sampling_kwargs, | |
| ): | |
| device = model.device | |
| if seq_mask is None: | |
| assert n_samples is not None | |
| seq_mask = model.make_seq_mask_for_sampling( | |
| n_samples=n_samples, | |
| min_len=sample_length_range[0], | |
| max_len=sample_length_range[1], | |
| ) | |
| start = time.time() | |
| aux = model.sample( | |
| seq_mask=seq_mask, return_last=False, return_aux=True, **sampling_kwargs | |
| ) | |
| aux["runtime"] = time.time() - start | |
| seq_lens = seq_mask.sum(-1).long() | |
| cropped_samp_coords = [ | |
| s[: seq_lens[i], model.bb_idxs] for i, s in enumerate(aux["xt_traj"][-1]) | |
| ] | |
| if pdb_save_path is not None: | |
| gly_aatype = (seq_mask * residue_constants.restype_order["G"]).long() | |
| trimmed_aatype = [a[: seq_lens[i]] for i, a in enumerate(gly_aatype)] | |
| atom_mask = utils.atom37_mask_from_aatype(gly_aatype, seq_mask).cpu() | |
| for i in range(len(cropped_samp_coords)): | |
| utils.write_coords_to_pdb( | |
| cropped_samp_coords[i], | |
| f"{pdb_save_path}{i}.pdb", | |
| batched=False, | |
| aatype=trimmed_aatype[i], | |
| atom_mask=atom_mask[i], | |
| ) | |
| if return_aux: | |
| return aux | |
| else: | |
| if return_sampling_runtime: | |
| return cropped_samp_coords, seq_mask, aux["runtime"] | |
| else: | |
| return cropped_samp_coords, seq_mask | |
| def draw_allatom_samples( | |
| model: torch.nn.Module, | |
| seq_mask: TensorType["b n", float] = None, | |
| n_samples: int = None, | |
| sample_length_range: Tuple[int] = (50, 512), | |
| two_stage_sampling: bool = True, | |
| pdb_save_path: Optional[str] = None, | |
| return_aux: bool = False, | |
| return_sampling_runtime: bool = False, | |
| **sampling_kwargs, | |
| ): | |
| """Implement the default 2-stage all-atom sampling routine.""" | |
| def save_allatom_samples(aux, path): | |
| seq_lens = aux["seq_mask"].sum(-1).long() | |
| cropped_samp_coords = [ | |
| c[: seq_lens[i]] for i, c in enumerate(aux["xt_traj"][-1]) | |
| ] | |
| cropped_samp_aatypes = [ | |
| s[: seq_lens[i]] for i, s in enumerate(aux["st_traj"][-1]) | |
| ] | |
| samp_atom_mask = utils.atom37_mask_from_aatype( | |
| aux["st_traj"][-1].to(device), seq_mask | |
| ) | |
| samp_atom_mask = [m[: seq_lens[i]] for i, m in enumerate(samp_atom_mask)] | |
| for i, c in enumerate(cropped_samp_coords): | |
| utils.write_coords_to_pdb( | |
| c, | |
| f"{path}{i}.pdb", | |
| batched=False, | |
| aatype=cropped_samp_aatypes[i], | |
| atom_mask=samp_atom_mask[i], | |
| conect=True, | |
| ) | |
| device = model.device | |
| if seq_mask is None: | |
| assert n_samples is not None | |
| seq_mask = model.make_seq_mask_for_sampling( | |
| n_samples=n_samples, | |
| min_len=sample_length_range[0], | |
| max_len=sample_length_range[1], | |
| ) | |
| sampling_runtime = 0.0 | |
| # Stage 1 sampling | |
| start = time.time() | |
| if "stage_2" in sampling_kwargs: | |
| stage_2_kwargs = vars(sampling_kwargs.pop("stage_2")) | |
| aux = model.sample( | |
| seq_mask=seq_mask, | |
| return_last=False, | |
| return_aux=True, | |
| **sampling_kwargs, | |
| ) | |
| sampling_runtime = time.time() - start | |
| if pdb_save_path is not None and two_stage_sampling: | |
| save_allatom_samples(aux, pdb_save_path + "_init") | |
| # Stage 2 sampling (sidechain refinement only) | |
| if two_stage_sampling: | |
| samp_seq = aux["st_traj"][-1] | |
| samp_coords = aux["xt_traj"][-1] | |
| cond_atom_mask = utils.atom37_mask_from_aatype((seq_mask * 7).long(), seq_mask) | |
| aux = {f"stage1_{k}": v for k, v in aux.items()} | |
| start = time.time() | |
| stage2_aux = model.sample( | |
| gt_cond_atom_mask=cond_atom_mask.to(device), # condition on backbone | |
| gt_cond_seq_mask=seq_mask.to(device), | |
| gt_coords=samp_coords.to(device), | |
| gt_aatype=samp_seq.to(device), | |
| seq_mask=seq_mask, | |
| return_last=False, | |
| return_aux=True, | |
| **stage_2_kwargs, | |
| ) | |
| sampling_runtime += time.time() - start | |
| aux = {**aux, **stage2_aux} | |
| if pdb_save_path is not None: | |
| save_allatom_samples(aux, pdb_save_path + "_samp") | |
| aux["runtime"] = sampling_runtime | |
| # Process outputs, crop to correct length | |
| if return_aux: | |
| return aux | |
| else: | |
| xt_traj = aux["xt_traj"] | |
| st_traj = aux["st_traj"] | |
| seq_mask = aux["seq_mask"] | |
| seq_lens = seq_mask.sum(-1).long() | |
| cropped_samp_coords = [c[: seq_lens[i]] for i, c in enumerate(xt_traj[-1])] | |
| cropped_samp_aatypes = [s[: seq_lens[i]] for i, s in enumerate(st_traj[-1])] | |
| samp_atom_mask = utils.atom37_mask_from_aatype(st_traj[-1].to(device), seq_mask) | |
| samp_atom_mask = [m[: seq_lens[i]] for i, m in enumerate(samp_atom_mask)] | |
| orig_xt_traj = aux["stage1_xt_traj"] | |
| stage1_coords = [c[: seq_lens[i]] for i, c in enumerate(orig_xt_traj[-1])] | |
| ret = ( | |
| cropped_samp_coords, | |
| cropped_samp_aatypes, | |
| samp_atom_mask, | |
| stage1_coords, | |
| seq_mask, | |
| ) | |
| if return_sampling_runtime: | |
| ret = ret + (sampling_runtime,) | |
| return ret | |