Spaces:
Sleeping
Sleeping
| """ | |
| https://github.com/ProteinDesignLab/protpardelle | |
| License: MIT | |
| Author: Alex Chu | |
| Top-level model definitions. | |
| Typically these are initialized with config rather than arguments. | |
| """ | |
| import argparse | |
| from functools import partial | |
| import os | |
| from typing import Callable, List, Optional | |
| from einops import rearrange, repeat | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from torchtyping import TensorType | |
| from core import protein_mpnn | |
| from core import residue_constants | |
| from core import utils | |
| import diffusion | |
| import evaluation | |
| import modules | |
| class MiniMPNN(nn.Module): | |
| """Wrapper for ProteinMPNN network to predict sequence from structure.""" | |
| def __init__(self, config: argparse.Namespace): | |
| super().__init__() | |
| self.config = config | |
| self.model_config = cfg = config.model.mpnn_model | |
| self.n_tokens = config.data.n_aatype_tokens | |
| self.seq_emb_dim = cfg.n_channel | |
| time_cond_dim = cfg.n_channel * cfg.noise_cond_mult | |
| self.noise_block = modules.NoiseConditioningBlock(cfg.n_channel, time_cond_dim) | |
| self.token_embedding = nn.Linear(self.n_tokens, self.seq_emb_dim) | |
| self.mpnn_net = modules.NoiseConditionalProteinMPNN( | |
| n_channel=cfg.n_channel, | |
| n_layers=cfg.n_layers, | |
| n_neighbors=cfg.n_neighbors, | |
| time_cond_dim=time_cond_dim, | |
| vocab_size=config.data.n_aatype_tokens, | |
| input_S_is_embeddings=True, | |
| ) | |
| self.proj_out = nn.Linear(cfg.n_channel, self.n_tokens) | |
| def forward( | |
| self, | |
| denoised_coords: TensorType["b n a x", float], | |
| coords_noise_level: TensorType["b", float], | |
| seq_mask: TensorType["b n", float], | |
| residue_index: TensorType["b n", int], | |
| seq_self_cond: Optional[TensorType["b n t", float]] = None, # logprobs | |
| return_embeddings: bool = False, | |
| ): | |
| coords_noise_level_scaled = 0.25 * torch.log(coords_noise_level) | |
| noise_cond = self.noise_block(coords_noise_level_scaled) | |
| b, n, _, _ = denoised_coords.shape | |
| if seq_self_cond is None or not self.model_config.use_self_conditioning: | |
| seq_emb_in = torch.zeros(b, n, self.seq_emb_dim).to(denoised_coords) | |
| else: | |
| seq_emb_in = self.token_embedding(seq_self_cond.exp()) | |
| node_embs, encoder_embs = self.mpnn_net( | |
| denoised_coords, seq_emb_in, seq_mask, residue_index, noise_cond | |
| ) | |
| logits = self.proj_out(node_embs) | |
| pred_logprobs = F.log_softmax(logits, -1) | |
| if return_embeddings: | |
| return pred_logprobs, node_embs, encoder_embs | |
| return pred_logprobs | |
| class CoordinateDenoiser(nn.Module): | |
| """Wrapper for U-ViT module to denoise structure coordinates.""" | |
| def __init__(self, config: argparse.Namespace): | |
| super().__init__() | |
| self.config = config | |
| # Configuration | |
| self.sigma_data = config.data.sigma_data | |
| m_cfg = config.model.struct_model | |
| nc = m_cfg.n_channel | |
| bb_atoms = ["N", "CA", "C", "O"] | |
| n_atoms = config.model.struct_model.n_atoms | |
| self.use_conv = len(m_cfg.uvit.n_filt_per_layer) > 0 | |
| if self.use_conv and n_atoms == 37: | |
| n_atoms += 1 # make it an even number | |
| self.n_atoms = n_atoms | |
| self.bb_idxs = [residue_constants.atom_order[a] for a in bb_atoms] | |
| n_xyz = 9 if config.model.crop_conditional else 6 | |
| nc_in = n_xyz * n_atoms # xyz + selfcond xyz + maybe cropcond xyz | |
| # Neural networks | |
| n_noise_channel = nc * m_cfg.noise_cond_mult | |
| self.net = modules.TimeCondUViT( | |
| seq_len=config.data.fixed_size, | |
| patch_size=m_cfg.uvit.patch_size, | |
| dim=nc, | |
| depth=m_cfg.uvit.n_layers, | |
| n_filt_per_layer=m_cfg.uvit.n_filt_per_layer, | |
| heads=m_cfg.uvit.n_heads, | |
| dim_head=m_cfg.uvit.dim_head, | |
| conv_skip_connection=m_cfg.uvit.conv_skip_connection, | |
| n_atoms=n_atoms, | |
| channels_per_atom=n_xyz, | |
| time_cond_dim=n_noise_channel, | |
| position_embedding_type=m_cfg.uvit.position_embedding_type, | |
| ) | |
| self.noise_block = modules.NoiseConditioningBlock(nc, n_noise_channel) | |
| def forward( | |
| self, | |
| noisy_coords: TensorType["b n a x", float], | |
| noise_level: TensorType["b", float], | |
| seq_mask: TensorType["b n", float], | |
| residue_index: Optional[TensorType["b n", int]] = None, | |
| struct_self_cond: Optional[TensorType["b n a x", float]] = None, | |
| struct_crop_cond: Optional[TensorType["b n a x", float]] = None, | |
| ): | |
| # Prep inputs and time conditioning | |
| actual_var_data = self.sigma_data**2 | |
| var_noisy_coords = noise_level**2 + actual_var_data | |
| emb = noisy_coords / utils.expand(var_noisy_coords.sqrt(), noisy_coords) | |
| struct_noise_scaled = 0.25 * torch.log(noise_level) | |
| noise_cond = self.noise_block(struct_noise_scaled) | |
| # Prepare self- and crop-conditioning and concatenate along channels | |
| if struct_self_cond is None: | |
| struct_self_cond = torch.zeros_like(noisy_coords) | |
| if self.config.model.crop_conditional: | |
| if struct_crop_cond is None: | |
| struct_crop_cond = torch.zeros_like(noisy_coords) | |
| else: | |
| struct_crop_cond = struct_crop_cond / self.sigma_data | |
| emb = torch.cat([emb, struct_self_cond, struct_crop_cond], -1) | |
| else: | |
| emb = torch.cat([emb, struct_self_cond], -1) | |
| # Run neural network | |
| emb = self.net(emb, noise_cond, seq_mask=seq_mask, residue_index=residue_index) | |
| # Preconditioning from Karras et al. | |
| out_scale = noise_level * actual_var_data**0.5 / torch.sqrt(var_noisy_coords) | |
| skip_scale = actual_var_data / var_noisy_coords | |
| emb = emb * utils.expand(out_scale, emb) | |
| skip_info = noisy_coords * utils.expand(skip_scale, noisy_coords) | |
| denoised_coords_x0 = emb + skip_info | |
| # Don't use atom mask; denoise all atoms | |
| denoised_coords_x0 *= utils.expand(seq_mask, denoised_coords_x0) | |
| return denoised_coords_x0 | |
| class Protpardelle(nn.Module): | |
| """All-atom protein diffusion-based generative model. | |
| This class wraps a structure denoising network and a sequence prediction network | |
| to do structure/sequence co-design (for all-atom generation), or backbone generation. | |
| It can be trained for one of four main tasks. To produce the all-atom (co-design) | |
| Protpardelle model, we will typically pretrain an 'allatom' model, then use this | |
| to train a 'seqdes' model. A 'seqdes' model can be trained with either a backbone | |
| or allatom denoiser. The two can be combined to yield all-atom (co-design) Protpardelle | |
| without further training. | |
| 'backbone': train only a backbone coords denoiser. | |
| 'seqdes': train only a mini-MPNN, using a pretrained coords denoiser. | |
| 'allatom': train only an allatom coords denoiser (cannot do all-atom generation | |
| by itself). | |
| 'codesign': train both an allatom denoiser and mini-MPNN at once. | |
| """ | |
| def __init__(self, config: argparse.Namespace, device: str = "cpu"): | |
| super().__init__() | |
| self.config = config | |
| self.device = device | |
| self.task = config.model.task | |
| self.n_tokens = config.data.n_aatype_tokens | |
| self.use_mpnn_model = self.task in ["seqdes", "codesign"] | |
| # Modules | |
| self.all_modules = {} | |
| self.bb_idxs = [0, 1, 2, 4] | |
| self.n_atoms = 37 | |
| self.struct_model = CoordinateDenoiser(config) | |
| self.all_modules["struct_model"] = self.struct_model | |
| self.bb_idxs = self.struct_model.bb_idxs | |
| self.n_atoms = self.struct_model.n_atoms | |
| if self.use_mpnn_model: | |
| self.mpnn_model = MiniMPNN(config) | |
| self.all_modules["mpnn_model"] = self.mpnn_model | |
| # Load any pretrained modules | |
| for module_name in self.config.model.pretrained_modules: | |
| self.load_pretrained_module(module_name) | |
| # Diffusion-related | |
| self.sigma_data = self.struct_model.sigma_data | |
| self.training_noise_schedule = partial( | |
| diffusion.noise_schedule, | |
| sigma_data=self.sigma_data, | |
| **vars(config.diffusion.training), | |
| ) | |
| self.sampling_noise_schedule_default = self.make_sampling_noise_schedule() | |
| def load_pretrained_module(self, module_name: str, ckpt_path: Optional[str] = None): | |
| """Load pretrained weights for a given module name.""" | |
| assert module_name in ["struct_model", "mpnn_model"], module_name | |
| # Load pretrained checkpoint | |
| if ckpt_path is None: | |
| ckpt_path = getattr(self.config.model, f"{module_name}_checkpoint") | |
| ckpt_path = os.path.join(self.config.train.home_dir, ckpt_path) | |
| ckpt_dict = torch.load(ckpt_path, map_location=self.device) | |
| model_state_dict = ckpt_dict["model_state_dict"] | |
| # Get only submodule state_dict | |
| submodule_state_dict = { | |
| sk[len(module_name) + 1 :]: sv | |
| for sk, sv in model_state_dict.items() | |
| if sk.startswith(module_name) | |
| } | |
| # Load into module | |
| module = dict(self.named_modules())[module_name] | |
| module.load_state_dict(submodule_state_dict) | |
| # Freeze unneeded modules | |
| if module_name == "struct_model": | |
| self.struct_model = module | |
| if self.task == "seqdes": | |
| for p in module.parameters(): | |
| p.requires_grad = False | |
| if module_name == "mpnn_model": | |
| self.mpnn_model = module | |
| if self.task not in ["codesign", "seqdes"]: | |
| for p in module.parameters(): | |
| p.requires_grad = False | |
| return module | |
| def load_minimpnn(self, mpnn_ckpt_path: Optional[str] = None): | |
| """Convert an allatom model to a codesign model.""" | |
| if mpnn_ckpt_path is None: | |
| mpnn_ckpt_path = "checkpoints/minimpnn_state_dict.pth" | |
| self.mpnn_model = MiniMPNN(self.config).to(self.device) | |
| self.load_pretrained_module("mpnn_model", ckpt_path=mpnn_ckpt_path) | |
| self.use_mpnn_model = True | |
| return | |
| def remove_minimpnn(self): | |
| """Revert a codesign model to an allatom model to a codesign model.""" | |
| self.use_mpnn_model = False | |
| self.mpnn_model = None | |
| self.all_modules["mpnn_model"] = None | |
| def make_sampling_noise_schedule(self, **noise_kwargs): | |
| """Make the default sampling noise schedule function.""" | |
| noise_schedule_kwargs = vars(self.config.diffusion.sampling) | |
| if len(noise_kwargs) > 0: | |
| noise_schedule_kwargs.update(noise_kwargs) | |
| return partial(diffusion.noise_schedule, **noise_schedule_kwargs) | |
| def forward( | |
| self, | |
| *, | |
| noisy_coords: TensorType["b n a x", float], | |
| noise_level: TensorType["b", float], | |
| seq_mask: TensorType["b n", float], | |
| residue_index: TensorType["b n", int], | |
| struct_self_cond: Optional[TensorType["b n a x", float]] = None, | |
| struct_crop_cond: Optional[TensorType["b n a x", float]] = None, | |
| seq_self_cond: Optional[TensorType["b n t", float]] = None, # logprobs | |
| run_struct_model: bool = True, | |
| run_mpnn_model: bool = True, | |
| ): | |
| """Main forward function for denoising/co-design. | |
| Arguments: | |
| noisy_coords: noisy array of xyz coordinates. | |
| noise_level: std of noise for each example in the batch. | |
| seq_mask: mask indicating which indexes contain data. | |
| residue_index: residue ordering. This is used by proteinMPNN, but currently | |
| only used by the diffusion model when the 'absolute_residx' or | |
| 'relative' position_embedding_type is specified. | |
| struct_self_cond: denoised coordinates from the previous step, scaled | |
| down by sigma data. | |
| struct_crop_cond: unnoised coordinates. unscaled (scaled down by sigma | |
| data inside the denoiser) | |
| seq_self_cond: mpnn-predicted sequence logprobs from the previous step. | |
| run_struct_model: flag to optionally not run structure denoiser. | |
| run_mpnn_model: flag to optionally not run mini-mpnn. | |
| """ | |
| # Coordinate denoiser | |
| denoised_x0 = noisy_coords | |
| if run_struct_model: | |
| denoised_x0 = self.struct_model( | |
| noisy_coords, | |
| noise_level, | |
| seq_mask, | |
| residue_index=residue_index, | |
| struct_self_cond=struct_self_cond, | |
| struct_crop_cond=struct_crop_cond, | |
| ) | |
| # Mini-MPNN | |
| aatype_logprobs = None | |
| if self.use_mpnn_model and run_mpnn_model: | |
| aatype_logprobs = self.mpnn_model( | |
| denoised_x0.detach(), | |
| noise_level, | |
| seq_mask, | |
| residue_index, | |
| seq_self_cond=seq_self_cond, | |
| return_embeddings=False, | |
| ) | |
| aatype_logprobs = aatype_logprobs * seq_mask[..., None] | |
| # Process outputs | |
| if aatype_logprobs is None: | |
| aatype_logprobs = repeat(seq_mask, "b n -> b n t", t=self.n_tokens) | |
| aatype_logprobs = torch.ones_like(aatype_logprobs) | |
| aatype_logprobs = F.log_softmax(aatype_logprobs, -1) | |
| struct_self_cond_out = denoised_x0.detach() / self.sigma_data | |
| seq_self_cond_out = aatype_logprobs.detach() | |
| return denoised_x0, aatype_logprobs, struct_self_cond_out, seq_self_cond_out | |
| def make_seq_mask_for_sampling( | |
| self, | |
| prot_lens: Optional[TensorType["b", int]] = None, | |
| n_samples: int = 1, | |
| min_len: int = 50, | |
| max_len: Optional[int] = None, | |
| ): | |
| """Makes a sequence mask of varying protein lengths (only input required | |
| to begin sampling). | |
| """ | |
| if max_len is None: | |
| max_len = self.config.data.fixed_size | |
| if prot_lens is None: | |
| possible_lens = np.arange(min_len, max_len) | |
| prot_lens = torch.Tensor(np.random.choice(possible_lens, n_samples)) | |
| else: | |
| n_samples = len(prot_lens) | |
| max_len = max(prot_lens) | |
| mask = repeat(torch.arange(max_len), "n -> b n", b=n_samples) | |
| mask = (mask < prot_lens[:, None]).float().to(self.device) | |
| return mask | |
| def sample( | |
| self, | |
| *, | |
| seq_mask: TensorType["b n", float] = None, | |
| n_samples: int = 1, | |
| min_len: int = 50, | |
| max_len: int = 512, | |
| residue_index: TensorType["b n", int] = None, | |
| gt_coords: TensorType["b n a x", float] = None, | |
| gt_coords_traj: List[TensorType["b n a x", float]] = None, | |
| gt_cond_atom_mask: TensorType["b n a", float] = None, | |
| gt_aatype: TensorType["b n", int] = None, | |
| gt_cond_seq_mask: TensorType["b n", float] = None, | |
| apply_cond_proportion: float = 1.0, | |
| n_steps: int = 200, | |
| step_scale: float = 1.2, | |
| s_churn: float = 50.0, | |
| noise_scale: float = 1.0, | |
| s_t_min: float = 0.01, | |
| s_t_max: float = 50.0, | |
| temperature: float = 1.0, | |
| top_p: float = 1.0, | |
| disallow_aas: List[int] = [4, 20], # cys, unk | |
| sidechain_mode: bool = False, | |
| skip_mpnn_proportion: float = 0.7, | |
| anneal_seq_resampling_rate: Optional[str] = None, # linear, cosine | |
| use_fullmpnn: bool = False, | |
| use_fullmpnn_for_final: bool = True, | |
| use_reconstruction_guidance: bool = False, | |
| use_classifier_free_guidance: bool = False, # defaults to replacement guidance if these are all false | |
| guidance_scale: float = 1.0, | |
| noise_schedule: Optional[Callable] = None, | |
| tqdm_pbar: Optional[Callable] = None, | |
| return_last: bool = True, | |
| return_aux: bool = False, | |
| ): | |
| """Sampling function for backbone or all-atom diffusion. All arguments are optional. | |
| Arguments: | |
| seq_mask: mask defining the number and lengths of proteins to be sampled. | |
| n_samples: number of samples to draw (if seq_mask not provided). | |
| min_len: minimum length of proteins to be sampled (if seq_mask not provided). | |
| max_len: maximum length of proteins to be sampled (if seq_mask not provided). | |
| residue_index: residue index of proteins to be sampled. | |
| gt_coords: conditioning information for coords. | |
| gt_coords_traj: conditioning information for coords specified for each timestep | |
| (if gt_coords is not provided). | |
| gt_cond_atom_mask: mask identifying atoms to apply gt_coords. | |
| gt_aatype: conditioning information for sequence. | |
| gt_cond_seq_mask: sequence positions to apply gt_aatype. | |
| apply_cond_proportion: the proportion of timesteps to apply the conditioning. | |
| e.g. if 0.5, then the first 50% of steps use conditioning, and the last 50% | |
| are unconditional. | |
| n_steps: number of denoising steps (ODE discretizations). | |
| step_scale: scale to apply to the score. | |
| s_churn: gamma = s_churn / n_steps describes the additional noise to add | |
| relatively at each denoising step. Use 0.0 for deterministic sampling or | |
| 0.2 * n_steps as a rough default for stochastic sampling. | |
| noise_scale: scale to apply to gamma. | |
| s_t_min: don't apply s_churn below this noise level. | |
| s_t_max: don't apply s_churn above this noise level. | |
| temperature: scale to apply to aatype logits. | |
| top_p: don't tokens which fall outside this proportion of the total probability. | |
| disallow_aas: don't sample these token indices. | |
| sidechain_mode: whether to do all-atom sampling (False for backbone-only). | |
| skip_mpnn_proportion: proportion of timesteps from the start to skip running | |
| mini-MPNN. | |
| anneal_seq_resampling_rate: whether and how to decay the probability of | |
| running mini-MPNN. None, 'linear', or 'cosine' | |
| use_fullmpnn: use "full" ProteinMPNN at each step. | |
| use_fullmpnn_for_final: use "full" ProteinMPNN at the final step. | |
| use_reconstruction_guidance: use reconstruction guidance on the conditioning. | |
| use_classifier_free_guidance: use classifier-free guidance on the conditioning. | |
| guidance_scale: weight for reconstruction/classifier-free guidance. | |
| noise_schedule: specify the noise level timesteps for sampling. | |
| tqdm_pbar: progress bar in interactive contexts. | |
| return_last: return only the sampled structure and sequence. | |
| return_aux: return a dict of everything associated with the sampling run. | |
| """ | |
| def ode_step(sigma_in, sigma_next, xt_in, x0_pred, gamma, guidance_in=None): | |
| if gamma > 0: | |
| t_hat = sigma_in + gamma * sigma_in | |
| sigma_delta = torch.sqrt(t_hat**2 - sigma_in**2) | |
| noisier_x = xt_in + utils.expand( | |
| sigma_delta, xt_in | |
| ) * noise_scale * torch.randn_like(xt_in).to(xt_in) | |
| xt_in = noisier_x * utils.expand(seq_mask, noisier_x) | |
| sigma_in = t_hat | |
| mask = (sigma_in > 0).float() | |
| score = (xt_in - x0_pred) / utils.expand(sigma_in.clamp(min=1e-6), xt_in) | |
| score = score * utils.expand(mask, score) | |
| if use_reconstruction_guidance: | |
| guidance, guidance_mask = guidance_in | |
| guidance = guidance * guidance_mask[..., None] | |
| guidance_std = guidance[guidance_mask.bool()].var().sqrt() | |
| score_std = score[guidance_mask.bool()].var().sqrt() | |
| score = score + guidance * guidance_scale | |
| if use_classifier_free_guidance: | |
| # guidance_in is the unconditional x0 (x0_pred is the conditional x0) | |
| # guidance_scale = 1 + w from Ho paper | |
| # ==0: use only unconditional score; <1: interpolate the scores; | |
| # ==1: use only conditional score; >1: skew towards conditional score | |
| uncond_x0 = guidance_in | |
| uncond_score = (xt_in - uncond_x0) / utils.expand( | |
| sigma_in.clamp(min=1e-6), xt_in | |
| ) | |
| uncond_score = uncond_score * utils.expand(mask, uncond_score) | |
| score = guidance_scale * score + (1 - guidance_scale) * uncond_score | |
| step = score * step_scale * utils.expand(sigma_next - sigma_in, score) | |
| new_xt = xt_in + step | |
| return new_xt | |
| def sample_aatype(logprobs): | |
| # Top-p truncation | |
| probs = F.softmax(logprobs.clone(), dim=-1) | |
| sorted_prob, sorted_idxs = torch.sort(probs, descending=True) | |
| cumsum_prob = torch.cumsum(sorted_prob, dim=-1) | |
| sorted_indices_to_remove = cumsum_prob > top_p | |
| sorted_indices_to_remove[..., 0] = 0 | |
| sorted_prob[sorted_indices_to_remove] = 0 | |
| orig_probs = torch.scatter( | |
| torch.zeros_like(sorted_prob), | |
| dim=-1, | |
| index=sorted_idxs, | |
| src=sorted_prob, | |
| ) | |
| # Apply temperature and disallowed AAs and sample | |
| assert temperature >= 0.0 | |
| scaled_logits = orig_probs.clamp(min=1e-9).log() / (temperature + 1e-4) | |
| if disallow_aas: | |
| unwanted_mask = torch.zeros(scaled_logits.shape[-1]).to(scaled_logits) | |
| unwanted_mask[disallow_aas] = 1 | |
| scaled_logits -= unwanted_mask * 1e10 | |
| orig_probs = F.softmax(scaled_logits, dim=-1) | |
| categorical = torch.distributions.Categorical(probs=orig_probs) | |
| samp_aatype = categorical.sample() | |
| return samp_aatype | |
| def design_with_fullmpnn(batched_coords, seq_mask): | |
| seq_lens = seq_mask.sum(-1).long() | |
| designed_seqs = [ | |
| evaluation.design_sequence(c[: seq_lens[i]], model=fullmpnn_model)[0] | |
| for i, c in enumerate(batched_coords) | |
| ] | |
| designed_aatypes, _ = utils.batched_seq_to_aatype_and_mask( | |
| designed_seqs, max_len=seq_mask.shape[-1] | |
| ) | |
| return designed_aatypes | |
| # Initialize masks/features | |
| if seq_mask is None: # Sample random lengths | |
| assert gt_aatype is None # Don't condition on aatype without seq_mask | |
| seq_mask = self.make_seq_mask_for_sampling( | |
| n_samples=n_samples, | |
| min_len=min_len, | |
| max_len=max_len, | |
| ) | |
| if residue_index is None: | |
| residue_index = torch.arange(seq_mask.shape[-1]) | |
| residue_index = repeat(residue_index, "n -> b n", b=seq_mask.shape[0]) | |
| residue_index = residue_index.to(seq_mask) * seq_mask | |
| if use_fullmpnn or use_fullmpnn_for_final: | |
| fullmpnn_model = protein_mpnn.get_mpnn_model( | |
| path_to_model_weights=self.config.train.home_dir | |
| + "/ProteinMPNN/vanilla_model_weights", | |
| device=self.device, | |
| ) | |
| # Initialize noise schedule/parameters | |
| to_batch_size = lambda x: x * torch.ones(seq_mask.shape[0]).to(self.device) | |
| s_t_min = s_t_min * self.sigma_data | |
| s_t_max = s_t_max * self.sigma_data | |
| if noise_schedule is None: | |
| noise_schedule = self.sampling_noise_schedule_default | |
| sigma = noise_schedule(1) | |
| timesteps = torch.linspace(1, 0, n_steps + 1) | |
| # Set up conditioning/guidance information | |
| crop_cond_coords = None | |
| if gt_coords is None: | |
| coords_shape = seq_mask.shape + (self.n_atoms, 3) | |
| xt = torch.randn(*coords_shape).to(self.device) * sigma | |
| xt *= utils.expand(seq_mask, xt) | |
| else: | |
| assert gt_coords_traj is None | |
| noise_levels = [to_batch_size(noise_schedule(t)) for t in timesteps] | |
| gt_coords_traj = [ | |
| diffusion.noise_coords(gt_coords, nl) for nl in noise_levels | |
| ] | |
| xt = gt_coords_traj[0] | |
| if gt_cond_atom_mask is not None: | |
| crop_cond_coords = gt_coords * gt_cond_atom_mask[..., None] | |
| gt_atom_mask = None | |
| if gt_aatype is not None: | |
| gt_atom_mask = utils.atom37_mask_from_aatype(gt_aatype, seq_mask) | |
| fake_logits = repeat(seq_mask, "b n -> b n t", t=self.n_tokens) | |
| s_hat = (sample_aatype(fake_logits) * seq_mask).long() | |
| # Initialize superposition for all-atom sampling | |
| if sidechain_mode: | |
| b, n = seq_mask.shape[:2] | |
| # Latest predicted x0 for sidechain superpositions | |
| atom73_state_0 = torch.zeros(b, n, 73, 3).to(xt) | |
| # Current state xt for sidechain superpositions (denoised to different levels) | |
| atom73_state_t = torch.randn(b, n, 73, 3).to(xt) * sigma | |
| # Noise level of xt | |
| sigma73_last = torch.ones(b, n, 73).to(xt) * sigma | |
| # Seqhat and mask used to choose sidechains for euler step (b, n) | |
| s_hat = (seq_mask * 7).long() | |
| mask37 = utils.atom37_mask_from_aatype(s_hat, seq_mask).bool() | |
| mask73 = utils.atom73_mask_from_aatype(s_hat, seq_mask).bool() | |
| begin_mpnn_step = int(n_steps * skip_mpnn_proportion) | |
| # Prepare to run sampling trajectory | |
| sigma = to_batch_size(sigma) | |
| x0 = None | |
| x0_prev = None | |
| x_self_cond = None | |
| s_logprobs = None | |
| s_self_cond = None | |
| if tqdm_pbar is None: | |
| tqdm_pbar = lambda x: x | |
| torch.set_grad_enabled(False) | |
| # *t_traj is the denoising trajectory; *0_traj is the evolution of predicted clean data | |
| # s0 are aatype probs of shape (b n t); s_hat are discrete aatype of shape (b n) | |
| xt_traj, x0_traj, st_traj, s0_traj = [], [], [], [] | |
| # Sampling trajectory | |
| for i, t in tqdm_pbar(enumerate(iter(timesteps[1:]))): | |
| # Set up noise levels | |
| sigma_next = noise_schedule(t) | |
| if i == n_steps - 1: | |
| sigma_next *= 0 | |
| gamma = ( | |
| s_churn / n_steps | |
| if (sigma_next >= s_t_min and sigma_next <= s_t_max) | |
| else 0.0 | |
| ) | |
| sigma_next = to_batch_size(sigma_next) | |
| if sidechain_mode: | |
| # Fill in noise for masked positions since xt is initialized to zeros at each step | |
| dummy_fill_noise = torch.randn_like(xt) * utils.expand(sigma, xt) | |
| zero_atom_mask = utils.atom37_mask_from_aatype(s_hat, seq_mask) | |
| dummy_fill_mask = 1 - zero_atom_mask[..., None] | |
| xt = xt * zero_atom_mask[..., None] + dummy_fill_noise * dummy_fill_mask | |
| else: # backbone only | |
| bb_seq = (seq_mask * residue_constants.restype_order["G"]).long() | |
| bb_atom_mask = utils.atom37_mask_from_aatype(bb_seq, seq_mask) | |
| xt *= bb_atom_mask[..., None] | |
| # Enable grad for reconstruction guidance | |
| if use_reconstruction_guidance: | |
| torch.set_grad_enabled(True) | |
| xt.requires_grad = True | |
| # Run denoising network | |
| run_mpnn = not sidechain_mode or i > begin_mpnn_step | |
| x0, s_logprobs, x_self_cond, s_self_cond = self.forward( | |
| noisy_coords=xt, | |
| noise_level=sigma, | |
| seq_mask=seq_mask, | |
| residue_index=residue_index, | |
| struct_self_cond=x_self_cond, | |
| struct_crop_cond=crop_cond_coords, | |
| seq_self_cond=s_self_cond, | |
| run_mpnn_model=run_mpnn, | |
| ) | |
| # Compute additional stuff for guidance | |
| if use_reconstruction_guidance: | |
| loss = (x0 - gt_coords).pow(2).sum(-1) | |
| loss = loss * gt_cond_atom_mask | |
| loss = loss.sum() / gt_cond_atom_mask.sum().clamp(min=1) | |
| xt.retain_grad() | |
| loss.backward() | |
| guidance = xt.grad.clone() | |
| xt.grad *= 0 | |
| torch.set_grad_enabled(False) | |
| if use_classifier_free_guidance: | |
| assert not use_reconstruction_guidance | |
| uncond_x0, _, _, _ = self.forward( | |
| noisy_coords=xt, | |
| noise_level=sigma, | |
| seq_mask=seq_mask, | |
| residue_index=residue_index, | |
| struct_self_cond=x_self_cond, | |
| seq_self_cond=s_self_cond, | |
| run_mpnn_model=run_mpnn, | |
| ) | |
| # Structure denoising step | |
| if not sidechain_mode: # backbone | |
| if sigma[0] > 0: | |
| xt = ode_step(sigma, sigma_next, xt, x0, gamma) | |
| else: | |
| xt = x0 | |
| else: # allatom | |
| # Write x0 into atom73_state_0 for atoms corresponding to old seqhat | |
| atom73_state_0[mask73] = x0[mask37] | |
| # Determine sequence resampling probability | |
| if anneal_seq_resampling_rate is not None: | |
| step_time = 1 - (i - begin_mpnn_step) / max( | |
| 1, n_steps - begin_mpnn_step | |
| ) | |
| if anneal_seq_resampling_rate == "linear": | |
| resampling_rate = step_time | |
| elif anneal_seq_resampling_rate == "cosine": | |
| k = 2 | |
| resampling_rate = ( | |
| 1 + np.cos(2 * np.pi * (step_time - 0.5)) | |
| ) / k | |
| resample_this_step = np.random.uniform() < resampling_rate | |
| # Resample sequence or design with full ProteinMPNN | |
| if i == n_steps - 1 and use_fullmpnn_for_final: | |
| s_hat = design_with_fullmpnn(x0, seq_mask).to(x0.device) | |
| elif anneal_seq_resampling_rate is None or resample_this_step: | |
| if run_mpnn and use_fullmpnn: | |
| s_hat = design_with_fullmpnn(x0, seq_mask).to(x0.device) | |
| else: | |
| s_hat = sample_aatype(s_logprobs) | |
| # Overwrite s_hat with any conditioning information | |
| if (i + 1) / n_steps <= apply_cond_proportion: | |
| if gt_cond_seq_mask is not None and gt_aatype is not None: | |
| s_hat = ( | |
| 1 - gt_cond_seq_mask | |
| ) * s_hat + gt_cond_seq_mask * gt_aatype | |
| s_hat = s_hat.long() | |
| # Set masks for collapsing superposition using new sequence | |
| mask37 = utils.atom37_mask_from_aatype(s_hat, seq_mask).bool() | |
| mask73 = utils.atom73_mask_from_aatype(s_hat, seq_mask).bool() | |
| # Determine prev noise levels for atoms corresponding to new sequence | |
| step_sigma_prev = ( | |
| torch.ones(*xt.shape[:-1]).to(xt) * sigma[..., None, None] | |
| ) | |
| step_sigma_prev[mask37] = sigma73_last[mask73] # b, n, 37 | |
| step_sigma_next = sigma_next[..., None, None] # b, 1, 1 | |
| # Denoising step on atoms corresponding to new sequence | |
| b, n = mask37.shape[:2] | |
| step_xt = torch.zeros(b, n, 37, 3).to(xt) | |
| step_x0 = torch.zeros(b, n, 37, 3).to(xt) | |
| step_xt[mask37] = atom73_state_t[mask73] | |
| step_x0[mask37] = atom73_state_0[mask73] | |
| guidance_in = None | |
| if (i + 1) / n_steps <= apply_cond_proportion: | |
| if use_reconstruction_guidance: | |
| guidance_in = (guidance, mask37.float()) | |
| elif use_classifier_free_guidance: | |
| guidance_in = uncond_x0 | |
| step_xt = ode_step( | |
| step_sigma_prev, | |
| step_sigma_next, | |
| step_xt, | |
| step_x0, | |
| gamma, | |
| guidance_in=guidance_in, | |
| ) | |
| xt = step_xt | |
| # Write new xt into atom73_state_t for atoms corresponding to new seqhat and update sigma_last | |
| atom73_state_t[mask73] = step_xt[mask37] | |
| sigma73_last[mask73] = step_sigma_next[0].item() | |
| # Replacement guidance if conditioning information provided | |
| if (i + 1) / n_steps <= apply_cond_proportion: | |
| if gt_coords_traj is not None: | |
| if gt_cond_atom_mask is None: | |
| xt = gt_coords_traj[i + 1] | |
| else: | |
| xt = (1 - gt_cond_atom_mask)[ | |
| ..., None | |
| ] * xt + gt_cond_atom_mask[..., None] * gt_coords_traj[i + 1] | |
| sigma = sigma_next | |
| # Logging | |
| xt_scale = self.sigma_data / utils.expand( | |
| torch.sqrt(sigma_next**2 + self.sigma_data**2), xt | |
| ) | |
| scaled_xt = xt * xt_scale | |
| xt_traj.append(scaled_xt.cpu()) | |
| x0_traj.append(x0.cpu()) | |
| st_traj.append(s_hat.cpu()) | |
| s0_traj.append(s_logprobs.cpu()) | |
| if return_last: | |
| return xt, s_hat, seq_mask | |
| elif return_aux: | |
| return { | |
| "x": xt, | |
| "s": s_hat, | |
| "seq_mask": seq_mask, | |
| "xt_traj": xt_traj, | |
| "x0_traj": x0_traj, | |
| "st_traj": st_traj, | |
| "s0_traj": s0_traj, | |
| } | |
| else: | |
| return xt_traj, x0_traj, st_traj, s0_traj, seq_mask | |