Constructs an identity transformation. Args: shape: The desired shape d

Applies the inverse of the transformation to a coordinate tensor. Args: pts: A [*, 3] coordinate tensor

Unique key to lookup bonds.

Unique key to lookup bonds.

Construct a feature dict of sequence features.

Extracts loop associated with a prefix from mmCIF data as a dictionary. Args: prefix: Prefix shared by each of the data items in the loop.

Simulates the reverse SDE for 1 step Args: x_t: [..., 3] current positions at time t in angstroms. score_t: [..., 3]

Performs matrix multiplication of two rotation matrix tensors. Written out by hand to avoid AMP downcasting. Args:

Applies a rotation to a vector. Written out by hand to avoid transfer to avoid AMP downcasting. Args: r: [*, 3,

Convert a batch of rotation matrices to rotation vectors (logarithmic map from SO(3) to so(3)). The standard logarithmic map can be derived f

Extract a rotation vector from the so(3) skew matrix basis. Args: skew_matrices (torch.Tensor): Skew matrices. Returns:

Softmax, but without automatic casting to fp32 when the input is of type bfloat16

Args: param_dict_list: A list of (nested) Param dicts to stack. The structure of each dict must be the identical

Pack and pad batch into torch tensors

Map a vector into the corresponding skew matrix so(3) basis. ``` [ 0 -z y] [x,y,z] -> [ z 0 -x] [ -y

Module for sampling rotations from the USO(3) distribution. Can be used to generate initial unbiased samples in the reverse process.

Adds a harmonic potential that restrains the system to a structure.

Generate the lookup table for sampling from the target SO(3) CDF. The table is 2D, with the rows corresponding to different sigma val

Returns PDB id and chain id for an HHSearch Hit.

Returns a pdb string provided OpenMM topology and positions.

Queries the database chunk using Jackhmmer.

Auxiliary routine for computing rotation angle and rotation axis from unit quaternions. To avoid complications, rotations vectors with angles

Computes the relative indices for each residue with respect to the original sequence.

Align structures in a ChemGraph batch to a reference, e.g. for RMSD computation. This uses the sparse formulation of pytorch geometric. If th

Move the molecule center to zero for sparse position tensors. Args: pos: [N, 3] batch positions of atoms in the molecule in sparse b

Chunk a list of blocks and run each chunk with activation checkpointing. We define a "block" as a callable whose only inputs are the outp

Define chi-angle rigid groups via one-hot representations.

Define chi-angle rigid groups via one-hot representations.

Adds missing atoms to Protein instance. Args: prot: A `protein.Protein` instance. checks: A `bool` specifying whether to add addition

Composes the transformation with a quaternion update vector of shape [*, 6], where the final 6 columns represent the x, y, an

Conditional variance of p(xt|x0). Var[x_t|x_0] = conditional_var(t)*I

Moves the transformation object to GPU memory Returns: A version of the transformation on GPU

Time-dependent diffusion coefficient.

Time-dependent drift coefficient.

Computes several checks for structural violations.

Args: q_x: [*, Q, C_q] query data kv_x: [*, K, C_k] key data biases:

Convert frames to their idealized all atom representation. Args: r: All rigid groups. [..., N, 8, 3] aatype: Residue types. [...,

Takes a PDB string and constructs a Protein object. WARNING: All non-standard residue types will be converted into UNK. All non-standard atoms

Auxiliary function for generating a lookup table from IGSO(3) expansions and their derivatives. Takes a basic function and loops over differe

Generate a weight matrix with variance according to He initialization. Based on a random (semi-)orthogonal matrix neural networks are exp

Truncated sum of IGSO(3) distribution. This function approximates the power series in equation 5 of "DENOISING DIFFUSION PROBABILISTIC MODELS

Args: x: [batch_size, residue_num, residue_num, embed_size] rigids: [batch_size, residue_num, 3, 3] num_n

Negative log probabilities

Negative log probabilities

Apply a Tensor -> Tensor function to underlying translation and rotation tensors, mapping over the translation/rotation dimen