# Copyright Generate Biomedicines, Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Models for generating protein sequence and side chain conformations given backbones. These can be used for sequence design and packing. """ from types import SimpleNamespace from typing import Callable, Literal, Optional, Tuple, Union import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.checkpoint import checkpoint from chroma import constants from chroma.data.xcs import validate_XC from chroma.layers import complexity, graph from chroma.layers.structure import diffusion, potts, protein_graph, sidechain from chroma.layers.structure.protein_graph_allatom import ( EdgeSidechainsDirect, NodeChiRBF, ) from chroma.utility.model import load_model as utility_load_model class GraphDesign(nn.Module): """Graph-based sequence design and sidechain packing. Given a fixed backbone, a GraphDesign model yields probabilities of residue type and angles by position. It encodes backbones with a `BackboneEncoderGNN` and then autoregressively factorizes the joint distribution of sequence and sidechain conformations given these graph embeddings. Optional first order marginal and Potts sequence decoders are also available. Some `GraphDesign` models are trained in a diffusion-aware mannner to model sequence likelihoods given a noised structure and a particular time point along a forwards diffusion process. Args: See documention of `structure.protein_graph.ProteinFeatureGraph`, and `graph.GraphNN` for more details. dim_nodes (int): Hidden dimension of node tensors of underlying GNNs. dim_edges (int): Hidden dimension of edge tensors of underlying GNNs. num_neighbors (int): Number of neighbors per node for underlying GNNs. node_features (tuple): List of node feature specifications for structure encoder. Features can be given as strings or as dictionaries. edge_features (tuple): List of edge feature specifications for structure encoder. Features can be given as strings or as dictionaries. sequence_embedding (str): How to represent sequence when decoding. Currently the only option is `linear`. sidechain_embedding (str): How to represent chi angles when decoding. Options include `chi_linear` for a simple linear layer, `chi_rbf` for a featurization based on smooth binning of chi angles, `X_direct` which directly encodes the all-atom coordinates using random Fourier features, and `mixed_chi_X` which uses both the featurizations of `chi_rbf` and of `X_direct`. sidechains (bool): Whether to use a joint sequence/sidechain autoregressive model to decode the backbones. num_layers (int): Number of layers of underlying GNNs. Can be overridden for the structure encoder by `num_layers_encoder`. num_layers_encoder (int, optional): Number of layers for structure encoder GNN. dropout (float): Dropout fraction used for all encoders and decoders except for the marginal sequence likelihood decoder in `decoder_S_marginals`. node_mlp_layers (int): Number of hidden layers for node update function of underlying GNNs. node_mlp_dim (int, optional): Dimension of hidden layers for node update function of underlying GNNs, defaults to match output dimension. edge_update (bool): Whether to include an edge update step in the GNNs. edge_mlp_layers (int): Number of hidden layers for edge update function of underlying GNNs. edge_mlp_dim (int, optional): Dimension of hidden layers for edge update function of underlying GNNs, defaults to match output dimension. skip_connect_input (bool): Whether to include skip connections between layers of underlying GNNs. mlp_activation (str): MLP nonlinearity function, `relu` or `softplus` accepted. num_alphabet (int): Number of possible residues for sequence decoder. num_chi_bins (int): Number of chi bins for smooth binning of chi angles used when `sidechain_embedding` is `chi_rbf` or `mixed_chi_X`. decoder_num_hidden (int): Dimension of hidden decoder layers. label_smoothing (float): Level of smoothing to apply to sequence and sidechain labels. separate_packing (bool): If True, then autoregressively factorize sequence and sidechains in two stages where the full sequence is predicted before all of the chi angles. Otherwise an interleaved factorization will be used that autoregressively predicts both the residue identity and chi angles in an alternating manner. Default is True. graph_criterion (str): Graph criterion for structure encoder, defines how neighbors are chosen. See `chroma.models.graph_design.BackboneEncoderGNN` for allowed values. graph_random_min_local (int): Minimum number of neighbors in GNN that come from local neighborhood, before random neighbors are chosen. graph_attentional (bool): Currently unused, previously used for experimental GNN attention mechanism. graph_num_attention_heads (int): Currently unused, previously used for experimental GNN attention mechanism. predict_S_marginals (bool): Whether to train marginal sequence decoder. predict_S_potts (bool): Whether to train Potts sequence decoder. potts_parameterization (str): How to parametrize Potts sequence decoder, see `chroma.layer.structure.potts` for allowed values. potts_num_factors (int, optional): Number of factors to use for Potts sequence decoder. potts_symmetric_J (bool): Whether to force J tensor of Potts model to be symmetric. noise_schedule (str, optional): Noise schedule for mapping between diffusion time and noise level, see chroma.layers.structure.diffusion.DiffusionChainCov for allowed values. If not set, model should only be provided with denoised backbones. noise_covariance_model (str): Covariance mode for mapping between diffusion time and noise level, see chroma.layers.structure.diffusion.DiffusionChainCov for allowed values. noise_complex_scaling (bool): Whether to scale noise for complexes. noise_beta_range (Tuple[float, float]): Minimum and maximum noise levels for noise schedule. noise_log_snr_range (Tuple[float, float]): Range of log signal-to-noise ratio for noising. Inputs: X (torch.Tensor): Backbone coordinates with shape `(num_batch, num_residues, num_atoms, 3)`. C (torch.LongTensor): Chain map with shape `(num_batch, num_residues)`. S (torch.LongTensor): Sequence tensor with shape `(num_batch, num_residues)`. t (torch.Tensor, optional): Diffusion timesteps corresponding to noisy input backbones, of shape `(num_batch)`. Use zeros when passing structures without noise. sample_noise (bool, optional): Whether to apply noise to input backbones. permute_idx (torch.LongTensor, optional): Permutation tensor for fixing the autoregressive decoding order `(num_batch, num_residues)`. If `None` (default), a random decoding order will be generated. priority (torch.Tensor, optional): Priority values for constraining residue orderings with shape `(num_batch, num_residues)`. If residues are assigned to integer-valued groups, the sampled permutation will be ordered such that all residues within a lower- valued priority group will occur before residues with higher-valued priority assignments. Outputs (dict): logp_S (torch.Tensor): Sequence log likelihoods per residue with shape `(num_batch, num_residues)`. logp_chi (torch.Tensor): Chi angle Log likelihoods per residue with shape `(num_batch, num_residues, 4)`. logp_S_marginals (torch.Tensor, optional): Sequence log likelihoods per residue from marginal decoder with shape `(num_batch, num_residues)`. logp_S_potts (torch.Tensor, optional): Sequence log likelihoods per residue from Potts decoder with shape `(num_batch, num_residues)`. chi (torch.Tensor): Chi angles with shape `(num_batch, num_residues, 4)`. mask_chi (torch.Tensor): Chi angle mask with shape `(num_batch, num_residues, 4)`. node_h_chi (torch.Tensor): Node features used for predicting chi angles with shape `(num_batch, num_residues, dim_nodes)`. mask_i (torch.Tensor): Node mask with shape `(num_batch, num_residues)`. mask_ij (torch.Tensor): Edge mask with shape `(num_batch, num_nodes, num_neighbors)`. mask_ij_causal (torch.Tensor): Causal edge mask for autoregressive decoding with shape `(num_batch, num_nodes, num_neighbors)`. edge_idx (torch.LongTensor): Edge indices for neighbors with shape `(num_batch, num_residues, num_neighbors)`. permute_idx (torch.LongTensor, optional): Permutation tensor that was used for the autoregressive decoding order with shape `(num_batch, num_residues)`. X_noise (torch.Tensor): Noised structure coordinates with shape `(num_batch, num_residues, num_atoms, 3)`. """ def __init__( self, dim_nodes: int = 128, dim_edges: int = 128, num_neighbors: int = 30, node_features: tuple = (("internal_coords", {"log_lengths": True}),), edge_features: tuple = ( "distances_2mer", "orientations_2mer", "distances_chain", ), sequence_embedding: str = "linear", sidechain_embedding: str = "chi_rbf", sidechains: bool = True, num_layers: int = 3, num_layers_encoder: Optional[int] = None, dropout: float = 0.1, node_mlp_layers: int = 1, node_mlp_dim: Optional[int] = None, edge_update: bool = True, edge_mlp_layers: int = 1, edge_mlp_dim: Optional[int] = None, skip_connect_input: bool = False, mlp_activation: str = "softplus", num_alphabet: int = 20, num_chi_bins: int = 20, decoder_num_hidden: int = 512, label_smoothing: float = 0.1, separate_packing: bool = True, graph_criterion: str = "knn", graph_random_min_local: int = 20, graph_attentional: bool = False, graph_num_attention_heads: int = 4, predict_S_marginals: bool = False, predict_S_potts: bool = False, potts_parameterization: str = "factor", potts_num_factors: Optional[int] = None, potts_symmetric_J: bool = True, noise_schedule: Optional[str] = None, noise_covariance_model: str = "brownian", noise_complex_scaling: bool = False, noise_beta_range: Tuple[float, float] = (0.2, 70.0), noise_log_snr_range: Tuple[float, float] = (-7.0, 13.5), checkpoint_gradients: bool = False, **kwargs ) -> None: """Initialize GraphDesign network.""" super(GraphDesign, self).__init__() # Save configuration in kwargs self.kwargs = locals() self.kwargs.pop("self") for key in list(self.kwargs.keys()): if key.startswith("__") and key.endswith("__"): self.kwargs.pop(key) args = SimpleNamespace(**self.kwargs) # Important global options self.dim_nodes = dim_nodes self.dim_edges = dim_edges self.num_alphabet = num_alphabet self.num_chi_bins = num_chi_bins self.separate_packing = separate_packing self.sidechains = sidechains self.predict_S_potts = predict_S_potts self.traversal = ProteinTraversalSpatial() # Encoder GNN process backbone self.encoder = BackboneEncoderGNN( dim_nodes=args.dim_nodes, dim_edges=args.dim_edges, num_neighbors=args.num_neighbors, node_features=args.node_features, edge_features=args.edge_features, num_layers=( args.num_layers if args.num_layers_encoder is None else args.num_layers_encoder ), node_mlp_layers=args.node_mlp_layers, node_mlp_dim=args.node_mlp_dim, edge_update=args.edge_update, edge_mlp_layers=args.edge_mlp_layers, edge_mlp_dim=args.edge_mlp_dim, mlp_activation=args.mlp_activation, dropout=args.dropout, skip_connect_input=args.skip_connect_input, graph_criterion=args.graph_criterion, graph_random_min_local=args.graph_random_min_local, checkpoint_gradients=checkpoint_gradients, ) # Time features for diffusion if args.noise_schedule is not None: self.noise_perturb = diffusion.DiffusionChainCov( noise_schedule=args.noise_schedule, beta_min=args.noise_beta_range[0], beta_max=args.noise_beta_range[1], log_snr_range=args.noise_log_snr_range, covariance_model=args.noise_covariance_model, complex_scaling=args.noise_complex_scaling, ) self.time_features = diffusion.NoiseTimeEmbedding( dim_embedding=args.dim_nodes, noise_schedule=self.noise_perturb.noise_schedule, ) # Decoder GNN process backbone if self.sidechains: self.decoder = SidechainDecoderGNN( dim_nodes=args.dim_nodes, dim_edges=args.dim_edges, num_neighbors=args.num_neighbors, predict_S=True, predict_chi=(not args.separate_packing), sequence_embedding=args.sequence_embedding, sidechain_embedding=args.sidechain_embedding, num_layers=args.num_layers, node_mlp_layers=args.node_mlp_layers, node_mlp_dim=args.node_mlp_dim, edge_update=args.edge_update, edge_mlp_layers=args.edge_mlp_layers, edge_mlp_dim=args.edge_mlp_dim, mlp_activation=args.mlp_activation, dropout=args.dropout, skip_connect_input=args.skip_connect_input, num_alphabet=args.num_alphabet, num_chi_bins=args.num_chi_bins, decoder_num_hidden=args.decoder_num_hidden, label_smoothing=args.label_smoothing, checkpoint_gradients=checkpoint_gradients, ) if args.predict_S_marginals: self.decoder_S_marginals = NodePredictorS( num_alphabet=args.num_alphabet, dim_nodes=args.dim_nodes, dim_hidden=args.decoder_num_hidden, label_smoothing=args.label_smoothing, ) if args.predict_S_potts: self.decoder_S_potts = potts.GraphPotts( dim_nodes=args.dim_nodes, dim_edges=args.dim_edges, num_states=args.num_alphabet, parameterization=args.potts_parameterization, num_factors=args.potts_num_factors, symmetric_J=args.potts_symmetric_J, dropout=args.dropout, label_smoothing=args.label_smoothing, ) if args.separate_packing: # Optionally do a two-stage autoregressive prediction self.embed_S = nn.Embedding(args.num_alphabet, args.dim_nodes) self.encoder_S_gnn = graph.GraphNN( dim_nodes=args.dim_nodes, dim_edges=args.dim_edges, num_layers=args.num_layers, node_mlp_layers=args.node_mlp_layers, node_mlp_dim=args.node_mlp_dim, edge_update=args.edge_update, edge_mlp_layers=args.edge_mlp_layers, edge_mlp_dim=args.edge_mlp_dim, mlp_activation=args.mlp_activation, dropout=args.dropout, norm="transformer", scale=args.num_neighbors, skip_connect_input=args.skip_connect_input, checkpoint_gradients=checkpoint_gradients, ) self.decoder_chi = SidechainDecoderGNN( dim_nodes=args.dim_nodes, dim_edges=args.dim_edges, num_neighbors=args.num_neighbors, predict_S=False, predict_chi=True, sequence_embedding=args.sequence_embedding, sidechain_embedding=args.sidechain_embedding, num_layers=args.num_layers, node_mlp_layers=args.node_mlp_layers, node_mlp_dim=args.node_mlp_dim, edge_update=args.edge_update, edge_mlp_layers=args.edge_mlp_layers, edge_mlp_dim=args.edge_mlp_dim, mlp_activation=args.mlp_activation, dropout=args.dropout, skip_connect_input=args.skip_connect_input, num_alphabet=args.num_alphabet, num_chi_bins=args.num_chi_bins, decoder_num_hidden=args.decoder_num_hidden, label_smoothing=args.label_smoothing, checkpoint_gradients=checkpoint_gradients, ) if sidechains: self.chi_to_X = sidechain.SideChainBuilder() self.X_to_chi = sidechain.ChiAngles() self.loss_rmsd = sidechain.LossSideChainRMSD() self.loss_clash = sidechain.LossSidechainClashes() self.loss_eps = 1e-5 @validate_XC() def forward( self, X: torch.Tensor, C: torch.LongTensor, S: torch.LongTensor, t: Optional[torch.Tensor] = None, sample_noise: bool = False, permute_idx: Optional[torch.LongTensor] = None, priority: Optional[torch.LongTensor] = None, ) -> dict: # Sample noisy backbones X_noise = X if sample_noise and hasattr(self, "noise_perturb"): X_bb = X[:, :, :4, :] _schedule = self.noise_perturb.noise_schedule t = self.noise_perturb.sample_t(C, t) X_noise_bb = self.noise_perturb(X_bb, C, t=t) if self.sidechains: # Rebuild sidechains on noised backbone from native chi angles chi, mask_chi = self.X_to_chi(X, C, S) X_noise, mask_X = self.chi_to_X(X_noise_bb, C, S, chi) else: pass # TODO IDK what to return here node_h, edge_h, edge_idx, mask_i, mask_ij = self.encode(X_noise, C, t=t) logp_S_marginals = None if self.kwargs["predict_S_marginals"]: logp_S_marginals, _ = self.decoder_S_marginals(S, node_h, mask_i) logp_S_potts = None if self.kwargs["predict_S_potts"]: logp_S_potts = self.decoder_S_potts.loss( S, node_h, edge_h, edge_idx, mask_i, mask_ij ) # Sample random permutations and build autoregressive mask if permute_idx is None: permute_idx = self.traversal(X, C, priority=priority) if self.sidechains: # In one-stage packing, predict S and chi angles in an interleaved manner ( logp_S, logp_chi, chi, mask_chi, node_h_chi, _, _, _, mask_ij_causal, ) = self.decoder( X_noise, C, S, node_h, edge_h, edge_idx, mask_i, mask_ij, permute_idx ) else: logp_S = (None,) logp_chi = None chi = None mask_chi = None node_h_chi = None mask_ij_causal = None if self.separate_packing: # In two-stage packing, re-process embeddings with sequence node_h = node_h + mask_i.unsqueeze(-1) * self.embed_S(S) node_h, edge_h = self.encoder_S_gnn( node_h, edge_h, edge_idx, mask_i, mask_ij ) _, logp_chi, chi, mask_chi, node_h_chi, _, _, _, _ = self.decoder_chi( X_noise, C, S, node_h, edge_h, edge_idx, mask_i, mask_ij, permute_idx ) if t is None: t = torch.zeros(C.size(0), device=C.device) outputs = { "logp_S": logp_S, "logp_chi": logp_chi, "logp_S_marginals": logp_S_marginals, "logp_S_potts": logp_S_potts, "chi": chi, "mask_chi": mask_chi, "node_h_chi": node_h_chi, "mask_i": mask_i, "mask_ij": mask_ij, "mask_ij_causal": mask_ij_causal, "edge_idx": edge_idx, "permute_idx": permute_idx, "X_noise": X_noise, "t": t, } return outputs def set_gradient_checkpointing(self, flag: bool): """Sets gradient checkpointing to `flag` on all relevant modules""" self.encoder.checkpoint_gradients = flag self.encoder.gnn.checkpoint_gradients = flag if self.sidechains: self.decoder.checkpoint_gradients = flag self.decoder.gnn.checkpoint_gradients = flag if self.separate_packing: self.encoder_S_gnn.checkpoint_gradients = flag self.decoder_chi.checkpoint_gradients = flag self.decoder_chi.gnn.checkpoint_gradients = flag @validate_XC() def encode( self, X: torch.Tensor, C: torch.Tensor, t: Optional[torch.Tensor] = None ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Encode the backbone and (optionally) the noise level. Args: X (torch.Tensor): Backbone coordinates with shape `(num_batch, num_residues, num_atoms, 3)`. C (torch.LongTensor): Chain map with shape `(num_batch, num_residues)`. t (torch.Tensor, optional): Diffusion timesteps corresponding to noisy input backbones, of shape `(num_batch)`. Default is no noise. Returns: node_h (torch.Tensor): Node features with shape `(num_batch, num_residues, dim_nodes)`. edge_h (torch.Tensor): Edge features with shape `(num_batch, num_residues, num_neighbors, dim_edges)`. edge_idx (torch.LongTensor): Edge indices for neighbors with shape `(num_batch, num_residues, num_neighbors)`. mask_i (torch.Tensor): Node mask with shape `(num_batch, num_residues)`. mask_ij (torch.Tensor): Edge mask with shape `(num_batch, num_nodes, num_neighbors)`. """ node_h_aux = None if hasattr(self, "time_features"): t = 0.0 if t is None else t node_h_aux = self.time_features(t) node_h, edge_h, edge_idx, mask_i, mask_ij = self.encoder( X, C, node_h_aux=node_h_aux ) return node_h, edge_h, edge_idx, mask_i, mask_ij @validate_XC() def predict_marginals( self, X: torch.Tensor, C: torch.Tensor, t: Optional[torch.Tensor] = None ) -> Tuple[torch.Tensor, torch.Tensor]: """Predict sequence marginal likelihoods. Args: X (torch.Tensor): Backbone coordinates with shape `(num_batch, num_residues, num_atoms, 3)`. C (torch.LongTensor): Chain map with shape `(num_batch, num_residues)`. t (torch.Tensor, optional): Diffusion timesteps corresponding to noisy input backbones, of shape `(num_batch)`. Default is no noise. Returns: log_probs_S (torch.Tensor): Node-wise sequence log probabilities with shape `(num_batch, num_residues, 20)`. mask_i (torch.Tensor): Node mask with shape `(num_batch, num_residues)`. """ if not self.kwargs["predict_S_marginals"]: raise Exception( "This version of GraphDesign was not trained with marginal prediction" ) node_h, edge_h, edge_idx, mask_i, mask_ij = self.encode(X, C, t) log_probs_S = self.decoder_S_marginals.log_probs_S(node_h, mask_i) return log_probs_S, mask_i @validate_XC() def predict_potts( self, X: torch.Tensor, C: torch.Tensor, t: Optional[torch.Tensor] = None ) -> Tuple[torch.Tensor, torch.Tensor, torch.LongTensor]: """Predict sequence Potts model. Args: X (torch.Tensor): Backbone coordinates with shape `(num_batch, num_residues, num_atoms, 3)`. C (torch.LongTensor): Chain map with shape `(num_batch, num_residues)`. t (torch.Tensor, optional): Diffusion timesteps corresponding to noisy input backbones, of shape `(num_batch)`. Default is no noise. Returns: h (torch.Tensor): The h tensor of a Potts model with dimensions `(seq_length, n_tokens)`. J (torch.Tensor): The J tensor of a Potts model with dimensions `(seq_length, seq_length, n_tokens, n_tokens)`. edge_idx (torch.LongTensor): Edge indices for neighbors with shape `(num_batch, num_residues, num_neighbors)` from GNN encoding. """ if not self.kwargs["predict_S_potts"]: raise Exception( "This version of GraphDesign was not trained with Potts prediction" ) node_h, edge_h, edge_idx, mask_i, mask_ij = self.encode(X, C, t) h, J = self.decoder_S_potts(node_h, edge_h, edge_idx, mask_i, mask_ij) return h, J, edge_idx @validate_XC() def loss( self, X: torch.Tensor, C: torch.LongTensor, S: torch.LongTensor, t: Optional[torch.Tensor] = None, permute_idx: Optional[torch.LongTensor] = None, sample_noise: bool = False, batched: bool = True, **kwargs ) -> dict: """Compute losses used for training. Args: X (torch.Tensor): Backbone coordinates with shape `(num_batch, num_residues, num_atoms, 3)`. C (torch.LongTensor): Chain map with shape `(num_batch, num_residues)`. S (torch.LongTensor): Sequence tensor with shape `(num_batch, num_residues)`. t (torch.Tensor, optional): Diffusion timesteps corresponding to noisy input backbones, of shape `(num_batch)`. Default is no noise. permute_idx (torch.LongTensor, optional): Permutation tensor for fixing the autoregressive decoding order `(num_batch, num_residues)`. If `None` (default), a random decoding order will be generated. sample_noise (bool): Whether to apply noise to input backbones. batched (bool): Whether to batch average losses. Returns (dict): neglogp (torch.Tensor): Sum of `neglogp_S` and `neglogp_chi` with shape `(num_batch, num_residues)`. neglogp_S (torch.Tensor): Average negative log probability per residue identity with shape `(num_batch, num_residues)`. neglogp_S_marginals (torch.Tensor): Average negative log probability per residue identity from marginal decoder with shape `(num_batch, num_residues)`. neglogp_S_potts (torch.Tensor): Average negative log probability per residue identity from Potts decoder with shape `(num_batch, num_residues)`. neglogp_chi (torch.Tensor): Average negative log probability per chi angle with shape `(num_batch, num_residues)`. mask_chi (torch.Tensor): Chi angle mask with shape `(batch_size, num_residues, 4)`. rmsd (torch.Tensor): Average RMSD per side-chain after sampling. clash (torch.Tensor): Average number of clashes per side-chain after sampling. permute_idx (LongTensor, optional): Permutation tensor that was used for the autoregressive decoding order with shape `(num_batch, num_residues)`. """ o = self.forward( X, C, S, t=t, permute_idx=permute_idx, sample_noise=sample_noise ) # Aggregate into per-residue scores for the batch if batched: _avg = lambda m, l: (m * l).sum() / (m.sum() + self.loss_eps) else: _avg = lambda m, l: (m * l).sum(dim=tuple(range(1, l.dim()))) / ( m.sum(dim=tuple(range(1, l.dim()))) + self.loss_eps ) mask_S = o["mask_i"] neglogp_S = -_avg(mask_S, o["logp_S"]) neglogp_chi = -_avg(o["mask_chi"], o["logp_chi"]) neglogp = neglogp_S + neglogp_chi if o["logp_S_marginals"] is not None: neglogp_S_marginals = -_avg(mask_S, o["logp_S_marginals"]) neglogp = neglogp + neglogp_S_marginals else: neglogp_S_marginals = None if o["logp_S_potts"] is not None: neglogp_S_potts = -_avg(mask_S, o["logp_S_potts"]) neglogp = neglogp + neglogp_S_potts else: neglogp_S_potts = None # Evaluate sampled side chains decoder = self.decoder_chi if self.separate_packing else self.decoder chi_sample = decoder.decoder_chi.sample( S, o["mask_chi"], o["node_h_chi"], o["mask_i"], temperature=0.01 ) X_sample, mask_X = self.chi_to_X(o["X_noise"][:, :, :4, :], C, S, chi_sample) # RMSD loss rmsd_i = self.loss_rmsd(o["X_noise"], X_sample, C, S) rmsd = _avg(mask_S, rmsd_i) # Clash loss measures clashes generated to the past clashes = self.loss_clash( X_sample, C, S, edge_idx=o["edge_idx"], mask_ij=o["mask_ij_causal"] ) clash = _avg(mask_S, clashes) losses = { "neglogp": neglogp, "neglogp_S": neglogp_S, "neglogp_S_marginals": neglogp_S_marginals, "neglogp_S_potts": neglogp_S_potts, "neglogp_chi": neglogp_chi, "mask_chi": o["mask_chi"], "rmsd": rmsd, "clash": clash, "permute_idx": o["permute_idx"], "t": o["t"], } return losses @torch.no_grad() @validate_XC() def sample( self, X: torch.Tensor, C: torch.LongTensor, S: Optional[torch.LongTensor] = None, t: Optional[Union[float, torch.Tensor]] = None, t_packing: Optional[Union[float, torch.Tensor]] = None, mask_sample: Optional[torch.Tensor] = None, permute_idx: Optional[torch.LongTensor] = None, temperature_S: float = 0.1, temperature_chi: float = 1e-3, clamped: bool = False, resample_chi: bool = True, return_scores: bool = False, top_p_S: Optional[float] = None, ban_S: Optional[tuple] = None, sampling_method: Literal["potts", "autoregressive"] = "autoregressive", regularization: Optional[str] = "LCP", potts_sweeps: int = 500, potts_proposal: Literal["dlmc", "chromatic"] = "dlmc", verbose: bool = False, symmetry_order: Optional[int] = None, ) -> tuple: """Sample sequence and side chain conformations given an input structure. Args: X (torch.Tensor): All atom coordinates with shape `(num_batch, num_residues, 14, 3)`. C (torch.LongTensor): Chain map with shape `(num_batch, num_residues)`. S (torch.LongTensor): Sequence tensor with shape `(num_batch, num_residues)`. t (float or torch.Tensor, optional): Diffusion time for models trained with diffusion augmentation of input structures. Setting `t=0` or `t=None` will condition the model to treat the structure as exact coordinates, while values of `t > 0` will condition the model to treat structures as though they were drawn from noise-augmented ensembles with that noise level. Default is `None`, while for robust design we recommend `t=0.5`. May be a float or a tensor of shape `(num_batch)`. t_packing (float or torch.Tensor, optional): Potentially separate diffusion time for packing. mask_sample (torch.Tensor, optional): Binary tensor mask indicating positions to be sampled with shape `(num_batch, num_residues)` or position-specific valid amino acid choices with shape `(num_batch, num_residues, num_alphabet)`. If `None` (default), all positions will be sampled. permute_idx (LongTensor, optional): Permutation tensor for fixing the autoregressive decoding order `(num_batch, num_residues)`. If `None` (default), a random decoding order will be generated. temperature_S (float): Temperature parameter for sampling sequence tokens. A value of `temperature_S=1.0` corresponds to the model's unadjusted positions, though because of training such as label smoothing values less than 1.0 are recommended. Default is `0.1`. temperature_chi (float): Temperature parameter for sampling chi angles. Even if a high temperature sequence is sampled, this is recommended to always be low. Default is `1E-3`. clamped (bool): If `True`, no sampling is done and the likelihood values will be calculated for the input sequence and structure. Used for validating the sequential versus parallel decoding modes. Default is `False`. resample_chi (bool): If `True`, all chi angles will be resampled, even for sequence positions that were not sampled (i.e. the model will perform global repacking). Default is `True`. return_scores (bool): If `True`, return dictionary containing likelihood scores similar to those produced by `forward`. top_p_S (float, optional): Option to perform top-p sampling for autoregressive sequence decoding. If not `None` it will be the top-p value [1]. [1] Holtzman et al. The Curious Case of Neural Text Degeneration. (2020) ban_S (tuple, optional): An optional set of token indices from `chroma.constants.AA20` to ban during sampling. sampling_method (str): Sampling method for decoding sequence from structure. If `autoregressive`, sequences will be designed by ancestral sampling with the autoregessive decoder head. If `potts`, sequences will be designed via MCMC with the potts decoder head. regularization (str, optional): Optional sequence regularization to use during decoding. Can be `LCP` for Local Composition Perplexity regularization which penalizes local sequence windows from having unnaturally low compositional entropies. (Implemented for both `potts` and `autoregressive`) potts_sweeps (int): Number of sweeps to perform for MCMC sampling of `potts` decoder. A sweep corresponds to a sufficient number of Monte Carlo steps such that every position could have changed. potts_proposal (str): MCMC proposal for Potts sampling. Currently implemented proposals are `dlmc` for Discrete Langevin Monte Carlo [1] or `chromatic` for Gibbs sampling with graph coloring. [1] Sun et al. Discrete Langevin Sampler via Wasserstein Gradient Flow (2023). symmetry_order (int, optional): Optional integer argument to enable symmetric sequence decoding under `symmetry_order`-order symmetry. The first `(num_nodes // symmetry_order)` states will be free to move, and all consecutively tiled sets of states will be locked to these during decoding. Internally this is accomplished by summing the parameters Potts model under a symmetry constraint into this reduced sized system and then back imputing at the end. Currently only implemented for Potts models. Returns: X_sample (torch.Tensor): Sampled all atom coordinates with shape `(num_batch, num_residues, 14, 3)`. S_sample (torch.LongTensor): Sampled sequence tensor with shape `(num_batch, num_residues)`. permute_idx (torch.LongTensor): Permutation tensor that was used for the autoregressive decoding order with shape `(num_batch, num_residues)`. scores (dict, optional): Dictionary containing likelihood scores similar to those produced by `forward`. """ if X.shape[2] == 4: X = F.pad(X, [0, 0, 0, 10]) alphabet = constants.AA20 node_h, edge_h, edge_idx, mask_i, mask_ij = self.encode(X, C, t=t) # Process sampling mask logits_init = torch.zeros( list(C.shape) + [len(alphabet)], device=C.device ).float() if ban_S is not None: ban_S = [alphabet.index(c) for c in ban_S] mask_sample, mask_sample_1D, S_init = potts.init_sampling_masks( logits_init, mask_sample, S=S, ban_S=ban_S ) if not clamped: S = S_init # Sample random permutations and build autoregressive mask if permute_idx is None: permute_idx = self.traversal(X, C, priority=mask_sample_1D) if symmetry_order is not None and not (sampling_method == "potts"): raise NotImplementedError( "Symmetric decoding is currently only supported for Potts models" ) if sampling_method == "potts": if not self.kwargs["predict_S_potts"]: raise Exception( "This GraphDesign model was not trained with Potts prediction" ) # Complexity regularization penalty_func = None mask_ij_coloring = None edge_idx_coloring = None if regularization == "LCP": C_complexity = ( C if symmetry_order is None else C[:, : C.shape[1] // symmetry_order] ) penalty_func = lambda _S: complexity.complexity_lcp(_S, C_complexity) # edge_idx_coloring, mask_ij_coloring = complexity.graph_lcp(C, edge_idx, mask_ij) S_sample, _ = self.decoder_S_potts.sample( node_h, edge_h, edge_idx, mask_i, mask_ij, S=S, mask_sample=mask_sample, temperature=temperature_S, num_sweeps=potts_sweeps, penalty_func=penalty_func, proposal=potts_proposal, rejection_step=(potts_proposal == "chromatic"), verbose=verbose, edge_idx_coloring=edge_idx_coloring, mask_ij_coloring=mask_ij_coloring, symmetry_order=symmetry_order, ) chi_sample, logp_S, logp_chi = None, None, None else: # Sample sequence (and chi angles if one-stage) # Complexity regularization bias_S_func = None if regularization == "LCP": bias_S_func = complexity.complexity_scores_lcp_t S_sample, chi_sample, logp_S, logp_chi, _ = self.decoder.decode( X, C, S, node_h, edge_h, edge_idx, mask_i, mask_ij, permute_idx, temperature_S=temperature_S, temperature_chi=temperature_chi, sample=not clamped, mask_sample=mask_sample, resample_chi=resample_chi, top_p_S=top_p_S, ban_S=ban_S, bias_S_func=bias_S_func, ) if self.separate_packing: if t != t_packing: node_h, edge_h, edge_idx, mask_i, mask_ij = self.encode( X, C, t=t_packing ) # In two-stage packing, re-process embeddings with sequence node_h = node_h + mask_i.unsqueeze(-1) * self.embed_S(S_sample) node_h, edge_h = self.encoder_S_gnn( node_h, edge_h, edge_idx, mask_i, mask_ij ) _, chi_sample, _, logp_chi, _ = self.decoder_chi.decode( X, C, S_sample, node_h, edge_h, edge_idx, mask_i, mask_ij, permute_idx, temperature_chi=temperature_chi, sample=not clamped, mask_sample=mask_sample_1D, resample_chi=resample_chi, ) # Rebuild side chains X_sample, mask_X = self.chi_to_X(X[:, :, :4, :], C, S_sample, chi_sample) if return_scores: if sampling_method == "potts": raise NotImplementedError # Summarize mask_chi = sidechain.chi_mask(C, S_sample) neglogp_S = -(mask_i * logp_S).sum([1]) / ( (mask_i).sum([1]) + self.loss_eps ) neglogp_chi = -(mask_chi * logp_chi).sum([1, 2]) / ( mask_chi.sum([1, 2]) + self.loss_eps ) scores = { "neglogp_S": neglogp_S, "neglogp_chi": neglogp_chi, "logp_S": logp_S, "logp_chi": logp_chi, "mask_i": mask_i, "mask_chi": mask_chi, } return X_sample, S_sample, permute_idx, scores else: return X_sample, S_sample, permute_idx @validate_XC() def pack( self, X: torch.Tensor, C: torch.LongTensor, S: torch.LongTensor, permute_idx: Optional[torch.LongTensor] = None, temperature_chi: float = 1e-3, clamped: bool = False, resample_chi: bool = True, return_scores: bool = False, ) -> tuple: """Sample side chain conformations given an input structure. Args: X (torch.Tensor): All atom coordinates with shape `(num_batch, num_residues, 14, 3)`. C (torch.LongTensor): Chain map with shape `(num_batch, num_residues)`. S (torch.LongTensor): Sequence tensor with shape `(num_batch, num_residues)`. permute_idx (LongTensor, optional): Permutation tensor for fixing the autoregressive decoding order `(num_batch, num_residues)`. If `None` (default), a random decoding order will be generated. temperature_chi (float): Temperature parameter for sampling chi angles. Even if a high temperature sequence is sampled, this is recommended to always be low. Default is `1E-3`. clamped (bool): If `True`, no sampling is done and the likelihood values will be calculated for the input sequence and structure. Used for validating the sequential versus parallel decoding modes. Default is `False` resample_chi (bool): If `True`, all chi angles will be resampled, even for sequence positions that were not sampled (i.e. global repacking). Default is `True`. return_scores (bool): If `True`, return dictionary containing likelihood scores similar to those produced by `forward`. Returns: X_sample (torch.Tensor): Sampled all atom coordinates with shape `(num_batch, num_residues, 14, 3)`. neglogp_chi (torch.Tensor, optional): Average negative log probability per chi angle. permute_idx (torch.LongTensor): Permutation tensor that was used for the autoregressive decoding order with shape `(num_batch, num_residues)`. scores (dict, optional): Dictionary containing likelihood scores similar to those produced by `forward`. """ assert self.separate_packing with torch.no_grad(): if X.shape[2] == 4: X = F.pad(X, [0, 0, 0, 10]) node_h, edge_h, edge_idx, mask_i, mask_ij = self.encode(X, C) # Sample random permutations and build autoregressive mask if permute_idx is None: permute_idx = self.traversal(X, C) # In two-stage packing, re-process embeddings with sequence node_h = node_h + mask_i.unsqueeze(-1) * self.embed_S(S) node_h, edge_h = self.encoder_S_gnn( node_h, edge_h, edge_idx, mask_i, mask_ij ) _, chi_sample, _, logp_chi, _ = self.decoder_chi.decode( X, C, S, node_h, edge_h, edge_idx, mask_i, mask_ij, permute_idx, temperature_chi=temperature_chi, sample=not clamped, resample_chi=resample_chi, ) X_sample, mask_X = self.chi_to_X(X[:, :, :4, :], C, S, chi_sample) # Summarize mask_chi = sidechain.chi_mask(C, S) neglogp_chi = -(mask_chi * logp_chi).sum([1, 2]) / ( mask_chi.sum([1, 2]) + self.loss_eps ) if return_scores: scores = { "neglogp_chi": neglogp_chi, "logp_chi": logp_chi, "mask_i": mask_i, "mask_chi": mask_chi, } return X_sample, permute_idx, scores else: return X_sample, permute_idx return X_sample, neglogp_chi, permute_idx class BackboneEncoderGNN(nn.Module): """Graph Neural Network for processing protein structure into graph embeddings. Args: See documention of `structure.protein_graph.ProteinFeatureGraph`, and `graph.GraphNN` for more details. dim_nodes (int): Hidden dimension of node tensors. dim_edges (int): Hidden dimension of edge tensors. num_neighbors (int): Number of neighbors per nodes. node_features (tuple): List of node feature specifications. Features can be given as strings or as dictionaries. edge_features (tuple): List of edge feature specifications. Features can be given as strings or as dictionaries. num_layers (int): Number of layers. node_mlp_layers (int): Number of hidden layers for node update function. node_mlp_dim (int, optional): Dimension of hidden layers for node update function, defaults to match output dimension. edge_update (bool): Whether to include an edge update step. edge_mlp_layers (int): Number of hidden layers for edge update function. edge_mlp_dim (int, optional): Dimension of hidden layers for edge update function, defaults to match output dimension. skip_connect_input (bool): Whether to include skip connections between layers. mlp_activation (str): MLP nonlinearity function, `relu` or `softplus` accepted. dropout (float): Dropout fraction. graph_distance_atom_type (int): Atom type for computing residue-residue distances for graph construction. Negative values will specify centroid across atom types. Default is `-1` (centroid). graph_cutoff (float, optional): Cutoff distance for graph construction: mask any edges further than this cutoff. Default is `None`. graph_mask_interfaces (bool): Restrict connections only to within chains, excluding-between chain interactions. Default is `False`. graph_criterion (str): Method used for building graph from distances. Currently supported methods are `{knn, random_log, random_linear}`. Default is `knn`. graph_random_min_local (int): Minimum number of neighbors in GNN that come from local neighborhood, before random neighbors are chosen. checkpoint_gradients (bool): Switch to implement gradient checkpointing during training. Inputs: X (torch.Tensor): Backbone coordinates with shape `(num_batch, num_residues, num_atoms, 3)`. C (torch.LongTensor): Chain map with shape `(num_batch, num_residues)`. node_h_aux (torch.LongTensor, optional): Auxiliary node features with shape `(num_batch, num_residues, dim_nodes)`. edge_h_aux (torch.LongTensor, optional): Auxiliary edge features with shape `(num_batch, num_residues, num_neighbors, dim_edges)`. edge_idx (torch.LongTensor, optional): Input edge indices for neighbors with shape `(num_batch, num_residues, num_neighbors)`. mask_ij (torch.Tensor, optional): Input edge mask with shape `(num_batch, num_nodes, num_neighbors)`. Outputs: node_h (torch.Tensor): Node features with shape `(num_batch, num_residues, dim_nodes)`. edge_h (torch.Tensor): Edge features with shape `(num_batch, num_residues, num_neighbors, dim_edges)`. edge_idx (torch.LongTensor): Edge indices for neighbors with shape `(num_batch, num_residues, num_neighbors)`. mask_i (torch.Tensor): Node mask with shape `(num_batch, num_residues)`. mask_ij (torch.Tensor): Edge mask with shape `(num_batch, num_nodes, num_neighbors)`. """ def __init__( self, dim_nodes: int = 128, dim_edges: int = 128, num_neighbors: int = 30, node_features: tuple = (("internal_coords", {"log_lengths": True}),), edge_features: tuple = ( "distances_2mer", "orientations_2mer", "distances_chain", ), num_layers: int = 3, node_mlp_layers: int = 1, node_mlp_dim: Optional[int] = None, edge_update: bool = True, edge_mlp_layers: int = 1, edge_mlp_dim: Optional[int] = None, skip_connect_input: bool = False, mlp_activation: str = "softplus", dropout: float = 0.1, graph_distance_atom_type: int = -1, graph_cutoff: Optional[float] = None, graph_mask_interfaces: bool = False, graph_criterion: str = "knn", graph_random_min_local: int = 20, checkpoint_gradients: bool = False, **kwargs ) -> None: """Initialize BackboneEncoderGNN.""" super(BackboneEncoderGNN, self).__init__() # Save configuration in kwargs self.kwargs = locals() self.kwargs.pop("self") for key in list(self.kwargs.keys()): if key.startswith("__") and key.endswith("__"): self.kwargs.pop(key) args = SimpleNamespace(**self.kwargs) # Important global options self.dim_nodes = dim_nodes self.dim_edges = dim_edges self.checkpoint_gradients = checkpoint_gradients graph_kwargs = { "distance_atom_type": args.graph_distance_atom_type, "cutoff": args.graph_cutoff, "mask_interfaces": args.graph_mask_interfaces, "criterion": args.graph_criterion, "random_min_local": args.graph_random_min_local, } self.feature_graph = protein_graph.ProteinFeatureGraph( dim_nodes=args.dim_nodes, dim_edges=args.dim_edges, num_neighbors=args.num_neighbors, graph_kwargs=graph_kwargs, node_features=args.node_features, edge_features=args.edge_features, ) self.gnn = graph.GraphNN( dim_nodes=args.dim_nodes, dim_edges=args.dim_edges, num_layers=args.num_layers, node_mlp_layers=args.node_mlp_layers, node_mlp_dim=args.node_mlp_dim, edge_update=args.edge_update, edge_mlp_layers=args.edge_mlp_layers, edge_mlp_dim=args.edge_mlp_dim, mlp_activation=args.mlp_activation, dropout=args.dropout, norm="transformer", scale=args.num_neighbors, skip_connect_input=args.skip_connect_input, checkpoint_gradients=checkpoint_gradients, ) @validate_XC(all_atom=False) def forward( self, X: torch.Tensor, C: torch.LongTensor, node_h_aux: Optional[torch.Tensor] = None, edge_h_aux: Optional[torch.Tensor] = None, edge_idx: Optional[torch.Tensor] = None, mask_ij: Optional[torch.Tensor] = None, ) -> Tuple[ torch.Tensor, torch.Tensor, torch.LongTensor, torch.Tensor, torch.Tensor ]: """Encode XC backbone structure into node and edge features.""" num_batch, num_residues = C.shape # Hack to enable checkpointing if self.checkpoint_gradients and (not X.requires_grad): X.requires_grad = True node_h, edge_h, edge_idx, mask_i, mask_ij = self._checkpoint( self.feature_graph, X, C, edge_idx, mask_ij ) if node_h_aux is not None: node_h = node_h + mask_i.unsqueeze(-1) * node_h_aux if edge_h_aux is not None: edge_h = edge_h + mask_ij.unsqueeze(-1) * edge_h_aux node_h, edge_h = self.gnn(node_h, edge_h, edge_idx, mask_i, mask_ij) return node_h, edge_h, edge_idx, mask_i, mask_ij def _checkpoint(self, module: nn.Module, *args) -> nn.Module: if self.checkpoint_gradients: return checkpoint(module, *args) else: return module(*args) class SidechainDecoderGNN(nn.Module): """Autoregressively generate sidechains given backbone graph embeddings. Args: See documention of `structure.protein_graph.ProteinFeatureGraph`, and `graph.GraphNN` for more details. dim_nodes (int): Hidden dimension of node tensors. dim_edges (int): Hidden dimension of edge tensors. num_neighbors (int): Number of neighbors per nodes. predict_S (bool): Whether to predict sequence. predict_chi (bool): Whether to predict chi angles. sequence_embedding (str): How to represent sequence when decoding. Currently the only option is `linear`. sidechain_embedding (str): How to represent chi angles when decoding. Options include `chi_linear` for a simple linear layer, `chi_rbf` for a featurization based on smooth binning of chi angles, `X_direct` which directly encodes the all-atom coordinates using random Fourier features, and `mixed_chi_X` which uses both the featurizations of `chi_rbf` and of `X_direct`. num_layers (int): Number of layers. node_mlp_layers (int): Number of hidden layers for node update function. node_mlp_dim (int, optional): Dimension of hidden layers for node update function, defaults to match output dimension. edge_update (bool): Whether to include an edge update step. edge_mlp_layers (int): Number of hidden layers for edge update function. edge_mlp_dim (int, optional): Dimension of hidden layers for edge update function, defaults to match output dimension. skip_connect_input (bool): Whether to include skip connections between layers. mlp_activation (str): MLP nonlinearity function, `relu` or `softplus` accepted. dropout (float): Dropout fraction. num_alphabet (int): Number of possible residues. num_chi_bins (int): Number of chi bins for smooth binning of chi angles used when `sidechain_embedding` is `chi_rbf` or `mixed_chi_X`. decoder_num_hidden (int): Dimension of hidden layers. label_smoothing (float): Level of smoothing to apply to sequence and sidechain labels. Inputs: X (torch.Tensor): Backbone coordinates with shape `(num_batch, num_residues, num_atoms, 3)`. C (torch.LongTensor): Chain map with shape `(num_batch, num_residues)`. S (torch.LongTensor): Sequence tensor with shape `(num_batch, num_residues)`. node_h (torch.Tensor): Node features with shape `(num_batch, num_residues, dim_nodes)`. edge_h (torch.Tensor): Edge features with shape `(num_batch, num_residues, num_neighbors, dim_edges)`. edge_idx (torch.LongTensor): Edge indices for neighbors with shape `(num_batch, num_residues, num_neighbors)`. mask_i (torch.Tensor): Node mask with shape `(num_batch, num_residues)`. mask_ij (torch.Tensor): Edge mask with shape `(num_batch, num_nodes, num_neighbors)`. permute_idx (torch.LongTensor): Permutation tensor for fixing the autoregressive decoding order `(num_batch, num_residues)`. If `None` (default), a random decoding order will be generated. Outputs: logp_S (torch.Tensor): Sequence log likelihoods per residue with shape `(num_batch, num_residues)`. logp_chi (torch.Tensor): Chi angle Log likelihoods per residue with shape `(num_batch, num_residues, 4)`. chi (torch.Tensor): Sidechain chi angles in radians with shape `(num_batch, num_residues, 4)`. mask_chi (torch.Tensor): Mask for chi angles with shape `(num_batch, num_residues, 4)`. node_h (torch.Tensor): Node features with shape `(num_batch, num_residues, dim_nodes)`. edge_h (torch.Tensor): Edge features with shape `(num_batch, num_residues, num_neighbors, dim_edges)`. edge_idx (torch.LongTensor): Edge indices for neighbors with shape `(num_batch, num_residues, num_neighbors)`. mask_i (torch.Tensor): Node mask with shape `(num_batch, num_residues)`. mask_ij (torch.Tensor): Edge mask with shape `(num_batch, num_nodes, num_neighbors)`. """ def __init__( self, dim_nodes: int = 128, dim_edges: int = 128, num_neighbors: int = 30, predict_S: bool = True, predict_chi: bool = True, sequence_embedding: str = "linear", sidechain_embedding: str = "mixed_chi_X", num_layers: int = 3, node_mlp_layers: int = 1, node_mlp_dim: Optional[int] = None, edge_update: bool = True, edge_mlp_layers: int = 1, edge_mlp_dim: Optional[int] = None, skip_connect_input: bool = False, mlp_activation: str = "softplus", dropout: float = 0.1, num_alphabet: int = 20, num_chi_bins: int = 20, decoder_num_hidden: int = 512, label_smoothing: float = 0.1, checkpoint_gradients: bool = False, **kwargs ): super(SidechainDecoderGNN, self).__init__() # Save configuration in kwargs self.kwargs = locals() self.kwargs.pop("self") for key in list(self.kwargs.keys()): if key.startswith("__") and key.endswith("__"): self.kwargs.pop(key) args = SimpleNamespace(**self.kwargs) # Important global options self.dim_nodes = dim_nodes self.dim_edges = dim_edges self.num_alphabet = num_alphabet self.num_chi_bins = num_chi_bins # Predict S, chi or both? assert predict_S or predict_chi self.predict_S = predict_S self.predict_chi = predict_chi self.sequence_embedding = sequence_embedding self.sidechain_embedding = sidechain_embedding if self.sequence_embedding == "linear": self.W_S = nn.Embedding(num_alphabet, dim_edges) # If we are predicting chi angles, then embed them if self.predict_chi: if self.sidechain_embedding == "chi_linear": self.W_chi = nn.Linear(8, dim_edges) elif self.sidechain_embedding == "chi_rbf": self.embed_chi = NodeChiRBF( dim_out=args.dim_edges, num_chi=4, num_chi_bins=args.num_chi_bins ) elif self.sidechain_embedding == "X_direct": self.embed_X = EdgeSidechainsDirect(dim_out=dim_edges) elif self.sidechain_embedding == "mixed_chi_X": self.embed_chi = NodeChiRBF( dim_out=args.dim_edges, num_chi=4, num_chi_bins=args.num_chi_bins ) self.embed_X = EdgeSidechainsDirect(dim_out=dim_edges, basis_type="rff") # Decoder GNN process backbone self.gnn = graph.GraphNN( dim_nodes=args.dim_nodes, dim_edges=args.dim_edges, num_layers=args.num_layers, node_mlp_layers=args.node_mlp_layers, node_mlp_dim=args.node_mlp_dim, edge_update=args.edge_update, edge_mlp_layers=args.edge_mlp_layers, edge_mlp_dim=args.edge_mlp_dim, mlp_activation=args.mlp_activation, dropout=args.dropout, norm="transformer", scale=args.num_neighbors, skip_connect_input=args.skip_connect_input, checkpoint_gradients=checkpoint_gradients, ) if self.predict_S: self.decoder_S = NodePredictorS( num_alphabet=args.num_alphabet, dim_nodes=args.dim_nodes, dim_hidden=args.decoder_num_hidden, label_smoothing=args.label_smoothing, ) if self.predict_chi: self.decoder_chi = NodePredictorChi( num_alphabet=args.num_alphabet, num_chi_bins=args.num_chi_bins, dim_nodes=args.dim_nodes, dim_hidden=args.decoder_num_hidden, label_smoothing=args.label_smoothing, ) self.loss_eps = 1e-5 self.chi_to_X = sidechain.SideChainBuilder() self.X_to_chi = sidechain.ChiAngles() @validate_XC() def forward( self, X: torch.Tensor, C: torch.LongTensor, S: torch.LongTensor, node_h: torch.Tensor, edge_h: torch.Tensor, edge_idx: torch.LongTensor, mask_i: torch.Tensor, mask_ij: torch.Tensor, permute_idx: torch.LongTensor, ) -> Tuple[ torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.LongTensor, torch.Tensor, torch.Tensor, ]: """Predict sequence and chi angles autoregressively given graph features.""" # Permute graph representation ( node_h_p, edge_h_p, edge_idx_p, mask_i_p, mask_ij_p, ) = graph.permute_graph_embeddings( node_h, edge_h, edge_idx, mask_i, mask_ij, permute_idx ) # Permute sequence and side chain chi angles X_p = graph.permute_tensor(X, 1, permute_idx) C_p = graph.permute_tensor(C, 1, permute_idx) S_p = graph.permute_tensor(S, 1, permute_idx) chi, mask_chi = self.X_to_chi(X, C, S) chi_p = graph.permute_tensor(chi, -2, permute_idx) # Decode system autoregressively in the permuted coordinates node_h_p, edge_h_p, edge_idx_p, mask_i_p, mask_ij_p = self._decode_inner( X_p, C_p, S_p, chi_p, node_h_p, edge_h_p, edge_idx_p, mask_i_p, mask_ij_p ) # Unpermute graph representation permute_idx_inverse = torch.argsort(permute_idx, dim=-1) node_h, edge_h, edge_idx, mask_i, mask_ij = graph.permute_graph_embeddings( node_h_p, edge_h_p, edge_idx_p, mask_i_p, mask_ij_p, permute_idx_inverse ) # Predict per-position joint probabilities of each side-chain's sequence and structure logp_S, log_probs_S, logp_chi, log_probs_chi = None, None, None, None if self.predict_S: (logp_S, log_probs_S,) = self.decoder_S(S, node_h, mask_i) if self.predict_chi: (logp_chi, log_probs_chi,) = self.decoder_chi( S, chi, mask_chi, node_h, mask_i ) return ( logp_S, logp_chi, chi, mask_chi, node_h, edge_h, edge_idx, mask_i, mask_ij, ) def _decode_inner( self, X_p, C_p, S_p, chi_p, node_h_p, edge_h_p, edge_idx_p, mask_i_p, mask_ij_p ): # Build autoregressive mask mask_ij_p = graph.edge_mask_causal(edge_idx_p, mask_ij_p) # Add sequence context h_S_p = self.W_S(S_p) h_S_p_ij = graph.collect_neighbors(h_S_p, edge_idx_p) edge_h_p = edge_h_p + mask_ij_p.unsqueeze(-1) * h_S_p_ij # Add side chain context if self.predict_chi: if self.sidechain_embedding in ["chi_rbf", "mixed_chi_X"]: h_chi_p = self.embed_chi(chi_p) h_chi_p_ij = graph.collect_neighbors(h_chi_p, edge_idx_p) edge_h_p = edge_h_p + mask_ij_p.unsqueeze(-1) * h_chi_p_ij if self.sidechain_embedding == "mixed_chi_X": edge_feature = self.embed_X(X_p, C_p, S_p, edge_idx_p) edge_h_p = edge_h_p + mask_ij_p.unsqueeze(-1) * edge_feature # Run decoder GNN in parallel (permuted) node_h_p, edge_h_p = self.gnn( node_h_p, edge_h_p, edge_idx_p, mask_i_p, mask_ij_p ) return node_h_p, edge_h_p, edge_idx_p, mask_i_p, mask_ij_p def _decode_scatter(self, tensor, src, t): """Decoding utility function: Scatter.""" idx = (t * torch.ones_like(src)).long() tensor.scatter_(1, idx, src) def _decode_pre_func(self, t, tensors_t): """Decoding pre-step function adds features based on current S and chi.""" _scatter_t = lambda tensor, src: self._decode_scatter(tensor, src, t) # Gather relevant tensors at step t edge_h_p_t = tensors_t["edge_h_cache"][0][:, t, :, :].unsqueeze(1) edge_idx_p_t = tensors_t["edge_idx"][:, t, :].unsqueeze(1) mask_ij_p_t = tensors_t["mask_ij"][:, t, :].unsqueeze(1) # Update the edge embeddings at t with the relevant context mask_ij_p_t = mask_ij_p_t.unsqueeze(-1) # Add sequence context h_S_p_ij_t = graph.collect_neighbors(tensors_t["h_S_p"], edge_idx_p_t) edge_h_p_t = edge_h_p_t + mask_ij_p_t * h_S_p_ij_t # Add chi context if self.predict_chi: if self.sidechain_embedding in ["chi_rbf", "mixed_chi_X"]: h_chi_p_ij_t = graph.collect_neighbors( tensors_t["h_chi_p"], edge_idx_p_t ) edge_h_p_t = edge_h_p_t + mask_ij_p_t * h_chi_p_ij_t if self.sidechain_embedding == "mixed_chi_X": h_chi_p_ij_t = self.embed_X.step( t, tensors_t["X_p"], tensors_t["C_p"], tensors_t["S_p"], edge_idx_p_t, ) edge_h_p_t = edge_h_p_t + mask_ij_p_t * h_chi_p_ij_t _scatter_t(tensors_t["edge_h_cache"][0], edge_h_p_t) return tensors_t def _decode_post_func( self, t, tensors_t, S_p_input, chi_p_input, temperature_S, temperature_chi, sample, resample_chi, mask_sample, mask_sample_p=None, top_p_S=None, ban_S=None, bias_S_func=None, ): """Decoding post-step function updates S and chi.""" _scatter_t = lambda tensor, src: self._decode_scatter(tensor, src, t) # Gather relevant tensors at step t C_p_t = tensors_t["C_p"][:, t].unsqueeze(1) edge_h_p_t = tensors_t["edge_h_cache"][0][:, t, :, :].unsqueeze(1) edge_idx_p_t = tensors_t["edge_idx"][:, t, :].unsqueeze(1) mask_i_p_t = tensors_t["mask_i"][:, t].unsqueeze(1) mask_ij_p_t = tensors_t["mask_ij"][:, t, :].unsqueeze(1) node_h_p_t = tensors_t["node_h_cache"][-1][:, t, :].unsqueeze(1) idx_p_t = tensors_t["idx_p"][:, t].unsqueeze(1) # Sample updated sequence S_p_t = S_p_input[:, t].unsqueeze(1).clone() if self.predict_S and sample: bias_S = None if bias_S_func is not None: bias_S = bias_S_func( t, tensors_t["S_p"], tensors_t["C_p"], tensors_t["idx_p"], edge_idx_p_t, mask_ij_p_t, ) mask_S_t = None if mask_sample_p is not None: mask_S_t = mask_sample_p[:, t] S_p_t = self.decoder_S.sample( node_h_p_t, mask_i_p_t, temperature=temperature_S, top_p=top_p_S, bias=bias_S, mask_S=mask_S_t, ) _scatter_t(tensors_t["S_p"], S_p_t) # Sample updated side chain conformations mask_chi_p_t = sidechain.chi_mask(C_p_t, S_p_t) chi_p_t = chi_p_input[:, t].unsqueeze(1).clone() if self.predict_chi and sample: # Sample chi angles chi_p_t_sample = self.decoder_chi.sample( S_p_t, mask_chi_p_t, node_h_p_t, mask_i_p_t, temperature=temperature_chi ) if mask_sample_p is not None and not resample_chi: m = mask_sample_p[:, t].unsqueeze(-1).expand([-1, 4]) chi_p_t = torch.where(m > 0, chi_p_t_sample, chi_p_t) else: chi_p_t = chi_p_t_sample # Rebuild side chain X_p_t_bb = tensors_t["X_p"][:, t, :4, :].unsqueeze(1) X_p_t, _ = self.chi_to_X(X_p_t_bb, C_p_t, S_p_t, chi_p_t) _scatter_t(tensors_t["X_p"], X_p_t) _scatter_t(tensors_t["chi_p"], chi_p_t) # Score the updated sequence and chi angles if self.predict_S: logp_S_p_t, _ = self.decoder_S(S_p_t, node_h_p_t, mask_i_p_t) _scatter_t(tensors_t["logp_S_p"], logp_S_p_t) if self.predict_chi: logp_chi_p_t, _ = self.decoder_chi( S_p_t, chi_p_t, mask_chi_p_t, node_h_p_t, mask_i_p_t ) _scatter_t(tensors_t["logp_chi_p"], logp_chi_p_t) # Update sequence and chi features (permuted) h_S_p_t = self.W_S(S_p_t) _scatter_t(tensors_t["h_S_p"], h_S_p_t) # Cache chi embeddings if self.predict_chi and self.sidechain_embedding in ["chi_rbf", "mixed_chi_X"]: h_chi_p_t = self.embed_chi(chi_p_t) _scatter_t(tensors_t["h_chi_p"], h_chi_p_t) return tensors_t @validate_XC() def decode( self, X: torch.Tensor, C: torch.LongTensor, S: torch.LongTensor, node_h: torch.Tensor, edge_h: torch.Tensor, edge_idx: torch.LongTensor, mask_i: torch.Tensor, mask_ij: torch.Tensor, permute_idx: torch.LongTensor, temperature_S: float = 0.1, temperature_chi: float = 1e-3, sample: bool = True, mask_sample: Optional[torch.Tensor] = None, resample_chi: bool = True, top_p_S: Optional[float] = None, ban_S: Optional[tuple] = None, bias_S_func: Optional[Callable] = None, ) -> Tuple[torch.LongTensor, torch.Tensor, torch.Tensor, torch.Tensor, dict]: """Autoregressively decode sequence and chi angles from graph features. Args: X (torch.Tensor): Backbone coordinates with shape `(num_batch, num_residues, num_atoms, 3)`. C (torch.LongTensor): Chain map with shape `(num_batch, num_residues)`. S (torch.LongTensor): Sequence tensor with shape `(num_batch, num_residues)`. node_h (torch.Tensor): Node features with shape `(num_batch, num_residues, dim_nodes)`. edge_h (torch.Tensor): Edge features with shape `(num_batch, num_residues, num_neighbors, dim_edges)`. edge_idx (torch.LongTensor): Edge indices for neighbors with shape `(num_batch, num_residues, num_neighbors)`. mask_i (torch.Tensor): Node mask with shape `(num_batch, num_residues)`. mask_ij (torch.Tensor): Edge mask with shape `(num_batch, num_nodes, num_neighbors)`. temperature_chi (float): Temperature parameter for sampling chi angles. Even if a high temperature sequence is sampled, this is recommended to always be low. Default is `1E-3`. sample (bool): Whether to sample sequence and chi angles. mask_sample (torch.Tensor, optional): Binary tensor mask indicating positions to be sampled with shape `(num_batch, num_residues)`. If `None` (default), all positions will be sampled. resample_chi (bool): If `True`, all chi angles will be resampled, even for sequence positions that were not sampled (i.e. global repacking). Default is `True`. top_p_S (float, optional): Top-p cutoff for Nucleus Sampling, see Holtzman et al ICLR 2020. ban_S (tuple, optional): An optional set of token indices from `chroma.constants.AA20` to ban during sampling. Returns: S (torch.LongTensor): Sequence tensor with shape `(num_batch, num_residues)`. chi (torch.Tensor): Chi angles with shape `(num_batch, num_residues, 4)`. logp_S (torch.Tensor): Sequence log likelihoods per residue with shape `(num_batch, num_residues)`. logp_chi (torch.Tensor): Chi angle Log likelihoods per residue with shape `(num_batch, num_residues, 4)`. tensors (dict): Processed tensors from GNN decoding. """ # Permute graph representation ( node_h_p, edge_h_p, edge_idx_p, mask_i_p, mask_ij_p, ) = graph.permute_graph_embeddings( node_h, edge_h, edge_idx, mask_i, mask_ij, permute_idx ) chi, mask_chi = self.X_to_chi(X, C, S) # Build autoregressive mask mask_ij_p = graph.edge_mask_causal(edge_idx_p, mask_ij_p) # Initialize tensors B, N, K = list(edge_idx.shape) device = node_h.device idx = torch.arange(end=N, device=device)[None, :].expand(C.shape) tensors_init = { "X_p": graph.permute_tensor(X, 1, permute_idx), "C_p": graph.permute_tensor(C, 1, permute_idx), "idx_p": graph.permute_tensor(idx, 1, permute_idx), "S_p": torch.zeros_like(S), "chi_p": torch.zeros([B, N, 4], device=device), "h_S_p": torch.zeros([B, N, self.dim_edges], device=device), "h_chi_p": torch.zeros([B, N, self.dim_edges], device=device), "node_h": node_h_p, "edge_h": edge_h_p, "edge_idx": edge_idx_p, "mask_i": mask_i_p, "mask_ij": mask_ij_p, "logp_S_p": torch.zeros([B, N], device=device), "logp_chi_p": torch.zeros([B, N, 4], device=device), } # As a sanity check against future state leakage, # we initialize S and chi and zero and write in the true value # during sequential decoding S_p_input = graph.permute_tensor(S, 1, permute_idx) chi_p_input = graph.permute_tensor(chi, 1, permute_idx) mask_sample_p = None if mask_sample is not None: mask_sample_p = graph.permute_tensor(mask_sample, 1, permute_idx) # Pre-step function features current sequence and chi angles pre_step_func = self._decode_pre_func # Post-step function samples sequence and/or chi angles at step t post_step_func = lambda t, tensors_t: self._decode_post_func( t, tensors_t, S_p_input, chi_p_input, temperature_S, temperature_chi, sample, resample_chi, mask_sample, mask_sample_p, top_p_S=top_p_S, ban_S=ban_S, bias_S_func=bias_S_func, ) # Sequentially step through a forwards pass of the GNN at each # position along the node dimension (1), running _pre_func # and each iteration and _post_func after each iteration tensors = self.gnn.sequential( tensors_init, pre_step_function=pre_step_func, post_step_function=post_step_func, ) # Unpermute sampled sequence and chi angles permute_idx_inverse = torch.argsort(permute_idx, dim=-1) S = graph.permute_tensor(tensors["S_p"], 1, permute_idx_inverse) chi = graph.permute_tensor(tensors["chi_p"], 1, permute_idx_inverse) logp_S = graph.permute_tensor(tensors["logp_S_p"], 1, permute_idx_inverse) logp_chi = graph.permute_tensor(tensors["logp_chi_p"], 1, permute_idx_inverse) return S, chi, logp_S, logp_chi, tensors def _filter_logits_top_p(logits, p=0.9): """Filter logits by top-p (Nucleus sampling). See Holtzman et al, ICLR 2020. Args: logits (Tensor): Logits with shape `(..., num_classes)`. p (float): Cutoff probability. Returns: logits_filters (Tensor): Filtered logits with shape `(..., num_classes)`. """ logits_sort, indices_sort = torch.sort(logits, dim=-1, descending=True) probs_sort = F.softmax(logits_sort, dim=-1) probs_cumulative = torch.cumsum(probs_sort, dim=-1) # Remove tokens outside nucleus (aside from top token) logits_sort_filtered = logits_sort.clone() logits_sort_filtered[probs_cumulative > p] = -float("Inf") logits_sort_filtered[..., 0] = logits_sort[..., 0] # Unsort logits_filtered = logits_sort_filtered.gather(-1, indices_sort.argsort(-1)) return logits_filtered class NodePredictorS(nn.Module): """Predict sequence tokens at each node given embeddings `P(S_i | h_i)`. Args: num_alphabet (int): Number of amino acids. dim_nodes (int): Node dimension of graph input. dim_hidden (int): Hidden layer dimension. loss_eps (float): Small number to avoid division by zero errors when taking averages. label_smoothing (float): Level of smoothing to apply. Inputs: S (torch.LongTensor): Sequence tensor with shape `(num_batch, num_residues)`. node_h (torch.Tensor): Node features with shape `(num_batch, num_residues, dim_nodes)`. mask_i (torch.Tensor): Node mask with shape `(num_batch, num_residues)`. Outputs: logp_S (torch.Tensor): Log likelihoods per residue with shape `(num_batch, num_residues)`. During training, this applies label smoothing. log_probs_S (torch.Tensor): Log probabilities for each token for at each residue with shape `(num_batch, num_residues, num_alphabet)`. """ def __init__( self, num_alphabet: int, dim_nodes: int, dim_hidden: int, loss_eps: float = 1e-5, label_smoothing: float = 0.1, ) -> None: super(NodePredictorS, self).__init__() self.num_alphabet = num_alphabet self.dim_nodes = dim_nodes self.dim_hidden = dim_hidden self.loss_eps = loss_eps self.label_smoothing = label_smoothing self.training_loss = torch.nn.CrossEntropyLoss( reduction="none", label_smoothing=self.label_smoothing ) # Layers for predicting sequence and chi angles self.S_mlp = graph.MLP( dim_in=dim_nodes, dim_hidden=dim_hidden, dim_out=self.num_alphabet, num_layers_hidden=2, ) def log_probs_S(self, node_h: torch.Tensor, mask_i: torch.Tensor) -> torch.Tensor: """Compute `log P(S | X, C)`.""" mask_i_expand = mask_i.unsqueeze(-1) S_logits = self.S_mlp(node_h) log_probs_S = mask_i_expand * F.log_softmax(S_logits, -1) return log_probs_S def forward( self, S: torch.LongTensor, node_h: torch.Tensor, mask_i: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: """Evaluate chi angle joint likelihood given graph embeddings.""" log_probs_S = self.log_probs_S(node_h, mask_i) if self.training: logp_S = -self.training_loss(log_probs_S.permute([0, 2, 1]), S) else: logp_S = torch.gather(log_probs_S, 2, S.unsqueeze(-1)).squeeze(-1) return logp_S, log_probs_S def sample( self, node_h: torch.Tensor, mask_i: torch.Tensor, temperature: float = 1.0, top_p: Optional[float] = None, mask_S: Optional[torch.Tensor] = None, bias: Optional[torch.Tensor] = None, ) -> torch.LongTensor: """Sample sequence and graph embeddings. Args: node_h (torch.Tensor): Node features with shape `(num_batch, num_residues, dim_nodes)`. mask_i (torch.Tensor): Node mask with shape `(num_batch, num_residues)`. temperature (float): Temperature parameter for sampling sequence tokens. The default value of 1.0 corresponds to the model's unadjusted positions, though because of training such as label smoothing values less than 1.0 are recommended. top_p (float, optional): Top-p cutoff for Nucleus Sampling, see Holtzman et al ICLR 2020. ban_S (tuple, optional): An optional set of token indices from `chroma.constants.AA20` to ban during sampling. Returns: S_sample (torch.LongTensor): Sampled sequence of shape `(num_batch, num_residues)`. """ num_batch, num_residues, _ = node_h.shape log_probs_S = self.log_probs_S(node_h, mask_i) if bias is not None: log_probs_S = log_probs_S + bias if mask_S is not None: log_probs_S = torch.where( mask_S > 0, log_probs_S, -float("Inf") * torch.ones_like(log_probs_S) ) if top_p is not None: log_probs_S = _filter_logits_top_p(log_probs_S, p=top_p) p = torch.distributions.categorical.Categorical( logits=log_probs_S / temperature ) S_sample = p.sample() return S_sample class NodePredictorChi(nn.Module): """Predict chi angles autoregressively at each node given embeddings. Decomposes as `P(chi_i_{1-4} | h_i) = P(chi_i_4 | chi_i_<4 h_i) ... P(chi_i_1 | h_i)`. Args: num_alphabet (int): Number of amino acids. num_chi_bins (int): Number of discretization bins per chi angle. dim_nodes (int): Node dimension of graph input. dim_hidden (int): Hidden layer dimension. loss_eps (float): Small number to avoid division by zero errors when taking averages. label_smoothing (float): Level of smoothing to apply. Inputs: S (torch.LongTensor): Sequence tensor with shape `(num_batch, num_residues)`. chi (torch.Tensor): Chi angles with shape `(num_batch, num_residues, 4)`. mask_chi (torch.Tensor): Chi angle mask with shape `(num_batch, num_residues, 4)`. node_h (torch.Tensor): Node features with shape `(num_batch, num_residues, dim_nodes)`. mask_i (torch.Tensor): Node mask with shape `(num_batch, num_residues)`. Outputs: logp_chi (torch.Tensor): Log likelihoods per residue with shape `(num_batch, num_residues, 4)`. During training, this applies label smoothing. log_probs_chi (torch.Tensor): Log probabilities for each chi angle token at each residue with shape `(num_batch, num_residues, 4, num_chi_bins)`. """ def __init__( self, num_alphabet: int, num_chi_bins: int, dim_nodes: int, dim_hidden: int, loss_eps: float = 1e-5, label_smoothing: float = 0.1, ) -> None: super(NodePredictorChi, self).__init__() self.num_alphabet = num_alphabet self.num_chi_bins = num_chi_bins self.dim_nodes = dim_nodes self.dim_hidden = dim_hidden self.loss_eps = loss_eps self.label_smoothing = label_smoothing self.training_loss = torch.nn.CrossEntropyLoss( reduction="none", label_smoothing=self.label_smoothing ) self._init_chi_bins(num_chi_bins) # Layers for embedding sequence and chi angles self.W_S = nn.Embedding(num_alphabet, dim_nodes) self.chi_embedding = nn.ModuleList( [ NodeChiRBF(dim_out=dim_nodes, num_chi=i, num_chi_bins=num_chi_bins) for i in [1, 2, 3] ] ) # Layers for chi angles self.chi_mlp = nn.ModuleList( [ graph.MLP( dim_in=dim_nodes, dim_hidden=dim_hidden, dim_out=num_chi_bins, num_layers_hidden=2, ) for t in range(4) ] ) def _init_chi_bins(self, num_chi_bins): # Setup bins bins = torch.tensor( np.linspace(-np.pi, np.pi, num_chi_bins + 1), dtype=torch.float32 ).reshape([1, 1, 1, -1]) self.register_buffer("bins_left", bins[:, :, :, 0:-1]) self.register_buffer("bins_right", bins[:, :, :, 1:]) return def _log_probs_t(self, t, S, chi, node_h, mask_i): """Compute `log P(chi_t | chi_ 0: chi_t = chi[:, :, :t] if len(chi_t.shape) == 2: chi_t = chi_t.unsqueeze(-1) node_h = node_h + self.chi_embedding[t - 1](chi_t) chi_logits = mask_i_expand * self.chi_mlp[t](node_h) log_probs_chi_t = mask_i_expand * F.log_softmax(chi_logits, -1) return log_probs_chi_t def _sample_continuous(self, logits, left, right): """Reparamaterization gradients via CDF inversion""" base_shape = list(logits.shape)[:-1] CMF = torch.cumsum(F.softmax(logits, dim=-1), dim=-1) u = torch.rand(base_shape, device=logits.device) _, max_idx = torch.max((u.unsqueeze(-1) < CMF).float(), dim=-1) max_idx = max_idx.unsqueeze(-1) left = left.expand(base_shape + [-1]) right = right.expand(base_shape + [-1]) # Gather panel bounds CMF_pad = F.pad(CMF, ((1, 0))) Y_left = torch.gather(left, -1, max_idx) Y_right = torch.gather(right, -1, max_idx) CMF_left = torch.gather(CMF_pad, -1, max_idx) CMF_right = torch.gather(CMF_pad, -1, max_idx + 1) # Local CDF inversion z = Y_left + (Y_right - Y_left) * (u.unsqueeze(-1) - CMF_left) / ( CMF_right - CMF_left + 1e-5 ) z = z.squeeze(-1) return z def forward( self, S: torch.LongTensor, chi: torch.Tensor, mask_chi: torch.Tensor, node_h: torch.Tensor, mask_i: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: """Evaluate chi angle joint likelihood given graph embeddings.""" # Build the likelihood sequentially log_probs_chi_list = [] for t in range(4): log_probs_chi_t = self._log_probs_t(t, S, chi, node_h, mask_i) log_probs_chi_list.append(log_probs_chi_t) log_probs_chi = torch.stack(log_probs_chi_list, -2) # Loss function chi = chi.unsqueeze(-1) chi_onehot = ((chi >= self.bins_left) * (chi < self.bins_right)).float() if self.training: scale = self.label_smoothing / (self.num_chi_bins - 1) chi_onehot = ( chi_onehot * (1 - self.label_smoothing) + (1 - chi_onehot) * scale ) logp_chi = mask_chi * (chi_onehot * log_probs_chi).sum(-1) return logp_chi, log_probs_chi def sample( self, S: torch.LongTensor, mask_chi: torch.Tensor, node_h: torch.Tensor, mask_i: torch.Tensor, temperature: float = 1.0, ) -> torch.Tensor: """Sample chi angles given sequence and graph embeddings. Args: S (torch.LongTensor): Sequence tensor with shape `(num_batch, num_residues)`. mask_chi (torch.Tensor): Chi angle mask with shape `(num_batch, num_residues, 4)`. node_h (torch.Tensor): Node features with shape `(num_batch, num_residues, dim_nodes)`. mask_i (torch.Tensor): Node mask with shape `(num_batch, num_residues)`. temperature (float): Temperature parameter for sampling sequence tokens. The default value of 1.0 corresponds to the model's unadjusted positions, though because of training such as label smoothing values less than 1.0 are recommended. Returns: chi_sample (torch.Tensor): Chi angles with shape `(num_batch, num_residues, 4)`. """ # Sample chi angles sequentially num_batch, num_residues, _ = node_h.shape chi = torch.zeros( [num_batch, num_residues, 4], dtype=torch.float32, device=node_h.device ) left = self.bins_left.reshape([1, 1, self.num_chi_bins]) right = self.bins_right.reshape([1, 1, self.num_chi_bins]) for t in range(4): log_probs_chi_t = self._log_probs_t(t, S, chi, node_h, mask_i) chi_t = self._sample_continuous(log_probs_chi_t / temperature, left, right) chi = chi + F.pad(chi_t.unsqueeze(-1), (t, 3 - t)) return mask_chi * chi class ProteinTraversalSpatial(nn.Module): """Samples spatial correlated residue permutations in a protein. Args: smooth_alpha (float): Smoothing parameter for graph smoothing where 0 corresponds to no smoothing and 1 corresponds to maximal smoothing. Default is 1. smooth_steps (int): Number of graph smoothing steps, which must be nonnegative. More steps will increase the amount of smoothing. Default is 5. smooth_randomize (bool): Enables uniform randomization of `smooth_alpha` on the interval `(0, smooth_alpha)`. Default is True. graph_num_neighbors (int): Number of neighbors for graph construction. Default is 30. deterministic (bool): Whether to force determinism. Default is False. Inputs: X (torch.Tensor): All atom coordinates with shape `(num_batch, num_residues, 14, 3)`. C (torch.LongTensor): Chain map with shape `(num_batch, num_residues)`. priority (torch.Tensor, optional): Priority values for constraining residue orderings with shape `(num_batch, num_residues)`. If residues are assigned to integer-valued groups, the sampled permutation will be ordered such that all residues within a lower-valued priority group will occur before residues with higher-valued priority assignments. Outputs: permute_idx (LongTensor): Permutation tensor containing reordered residue indices with shape `(num_batch, num_residues)`. """ def __init__( self, smooth_alpha: float = 1.0, smooth_steps: int = 5, smooth_randomize: bool = True, graph_num_neighbors: int = 30, deterministic: bool = False, ) -> None: super(ProteinTraversalSpatial, self).__init__() self.smooth_alpha = smooth_alpha self.smooth_steps = smooth_steps self.smooth_randomize = smooth_randomize self.deterministic = deterministic self._determistic_seed = 10 self.norm_eps = 1e-5 self.protein_graph = protein_graph.ProteinGraph( num_neighbors=graph_num_neighbors ) @validate_XC() def forward( self, X: torch.Tensor, C: torch.LongTensor, priority: Optional[torch.Tensor] = None, ): # Sample random node keys if not self.deterministic: z = torch.rand_like(C.float()) else: with torch.random.fork_rng(): torch.random.manual_seed(self._determistic_seed) z = torch.rand((1, C.shape[1]), device=C.device).expand(C.shape) # Graph-based smoothing alpha = self.smooth_alpha if self.smooth_randomize and not self.deterministic: alpha = torch.rand((), device=X.device) if alpha > 0: edge_idx, mask_ij = self.protein_graph(X, C) for i in range(self.smooth_steps): z_neighbors = graph.collect_neighbors( z.unsqueeze(-1), edge_idx ).squeeze(-1) z_average = (mask_ij * z_neighbors).sum(2) / ( mask_ij.sum(2) + self.norm_eps ) z = alpha * z_average + (1.0 - alpha) * z if priority is not None: z = z + priority # Create permutation permute_idx = torch.argsort(z, dim=-1) return permute_idx def load_model( weight_file: str, device: str = "cpu", strict: bool = False, strict_unexpected: bool = True, verbose: bool = True, ) -> GraphDesign: """Load model `GraphDesign` Args: weight_file (str): The destination path of the model weights to load. Compatible with files saved by `save_model`. device (str, optional): Pytorch device specification, e.g. `'cuda'` for GPU. Default is `'cpu'`. strict (bool): Whether to require that the keys match between the input file weights and the model created from the parameters stored in the model kwargs. strict_unexpected (bool): Whether to require that there are no unexpected keys when loading model weights, as distinct from the strict option which doesn't allow for missing keys either. By default, we use this option rather than strict for ease of development when adding model features. Returns: model (GraphDesign): Instance of `GraphDesign` with loaded weights. """ return utility_load_model( weight_file, GraphDesign, device=device, strict=strict, strict_unexpected=strict_unexpected, verbose=verbose, )