Remove insertion characters from A3M MSA file for AF3 compatibility. AlphaFold3 requires MSA sequences to have uniform length, so we remove

helper function for "partial" protocol ----------------------------------------- -order=[0,1,2] - change order of specified segments

Create rotation matrices from unnormalized vectors for the x and y-axes. This creates a rotation matrix from two vectors using Gram-Schmidt ortho

Composition of rotations 'a' and 'b'.

Save trajectory data and structure with thread-safe operations. Saves trajectory information to the appropriate directory based on it

set [arg]uments

Set the path to the final structure file for this trajectory. Args: final_struct (str): Path to the final PDB structure f

Sharded apply. Applies `fun` over shards to axes, in a way similar to vmap, but does so in shards of `shard_size`. Shards are stacked after. Th

Computes Squared difference between two arrays.

Pack and pad batch into torch tensors

given coordinates a-b-c-d, return dihedral

Update quality filtering and assessment metrics. Updates metrics related to trajectory quality assessment, filtering criteria

update the sequence features

Add two vectors 'v1' and 'v2'.

Normalizes vectors 'v'. Args: v: vectors to be normalized. epsilon: small regularizer added to squared norm before taking square root. Re

Computes v1 - v2.

Connects Module. Args: inputs: Tensor with at least num_input_dims dimensions. Returns: output of shape [...] + num_output.

Constructs Linear Module. Args: num_output: Number of output channels. Can be tuple when outputting multiple dimensions. ini

Get MSA features for unclustered uniprot, for pairing.

Calculate contact map from PDB coordinates. Args: pdb_path: Path to PDB file distance_threshold: Distance threshold for cont

Calculate count-based metrics (LIA, LIR, cLIA, cLIR). Args: transformed_pae: The PAE matrix after LIS transformation. combin

Calculates expected aligned distance errors for every pair of residues. Args: alignment_confidence_breaks: [num_bins - 1] the error bin edges.

TriangleMultiplication with fused projection weights.

Returns list of atom sites; contains data not present in the structure.

Get chain lengths from PDB file. Args: pdb_path: Path to PDB file Returns: Dictionary mapping chain IDs to their length

dgram to contacts

get edge index input: mask.shape = (...,L), X.shape = (...,L,3) return: (...,L,k)

get edge index input: mask.shape = (...,L), X.shape = (...,L,3) return: (...,L,k)

Returns the first model in a Biopython structure.

Returns a basic header containing method, release date and resolution.

helix bias loss

Get existing MSAs or generate new ones for Protenix. Checks for MSAs previously generated by AF3 pipeline. If not found, generates new M

Extracts polymer information for protein chains only. Args: parsed_info: _mmcif_dict produced by the Biopython parser. Returns: A dict m

get rmsd + alignment function align based on the first L positions, computed weighted rmsd using all positions (if include_L=True) or remainin

get sidechain rmsd + alignment function

one_hot to amino acid sequence

Insert BOS/EOS around each chain and chain separator between VH and VL: <VH>|<VL>.

Returns False if data is a special mmCIF character indicating 'unset'.

kmeans implemented in jax

kmeans++ implemented in jax, for initialization

Makes a mapping from PDB-format chain ID to sequence and description.

Returns atom indices needed to compute chi angles for all residue types. Returns: A tensor of shape [residue_types=21, chis=4, atoms=4]. The re

Matrices to map atoms to symmetry partners in ambiguous case.

Matrices to map atoms to symmetry partners in ambiguous case.

Mask which atoms are ambiguous in atom14.

Mask of which atoms are present for which residue type in atom14.

Map from atom14 to atom37 per residue type.

Mask of which atoms are present for which residue type in atom37.

Map from atom37 to atom14 per residue type.

Create Map from rigidgroups to atom37 indices.

Fill the arrays above.

Returns [num_res_types, num_atom_types] mask array.

Map probabilities from ColabDesign residue order to AbLang vocabulary order.

given coordinates a-b-c, c-d (L)ength, b-c-d (A)ngle, and a-b-c-d (D)ihedral return 4th coordinate d

compute RMSD of coordinates after alignment

Return differentiable STE probabilities in AbLang vocab space, sequence string, and hard token ids.

Parses the detailed HMM HMM comparison section for a single Hit. This works on .hhr files generated from both HHBlits and HHSearch. Args: de

Runs the monomer pipeline on a single chain.

radial basis function (RBF) input: (...,L,k) output: (...,L,k,?)

radial basis function (RBF) input: (...,L,k) output: (...,L,k,?)

multiple passes through the model (aka recycle)

Add relative position encodings. For position (i, j), the value is (i-j) clipped to [-k, k] and one-hotted. When not using 'use_chain_relati

Execute Protenix structure prediction via subprocess in a conda environment. Args: input_json (list): Protenix-compatible input JSON