protpardelle / models.py
Simon Duerr
webapp
8c639ec
raw
history blame
34.1 kB
"""
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