openfold.model.model

Classes

AlphaFold(config)

Alphafold 2.

class AlphaFold(config)

Bases: Module

Alphafold 2.

Implements Algorithm 2 (but with training).

embed_templates(batch, feats, z, pair_mask, templ_dim, inplace_safe)
forward(batch)
Parameters:

batch

Dictionary of arguments outlined in Algorithm 2. Keys must include the official names of the features in the supplement subsection 1.2.9.

The final dimension of each input must have length equal to the number of recycling iterations.

Features (without the recycling dimension):

”aatype” ([*, N_res]):

Contrary to the supplement, this tensor of residue indices is not one-hot.

”target_feat” ([*, N_res, C_tf])

One-hot encoding of the target sequence. C_tf is config.model.input_embedder.tf_dim.

”residue_index” ([*, N_res])

Tensor whose final dimension consists of consecutive indices from 0 to N_res.

”msa_feat” ([*, N_seq, N_res, C_msa])

MSA features, constructed as in the supplement. C_msa is config.model.input_embedder.msa_dim.

”seq_mask” ([*, N_res])

1-D sequence mask

”msa_mask” ([*, N_seq, N_res])

MSA mask

”pair_mask” ([*, N_res, N_res])

2-D pair mask

”extra_msa_mask” ([*, N_extra, N_res])

Extra MSA mask

”template_mask” ([*, N_templ])

Template mask (on the level of templates, not residues)

”template_aatype” ([*, N_templ, N_res])

Tensor of template residue indices (indices greater than 19 are clamped to 20 (Unknown))

”template_all_atom_positions”

([*, N_templ, N_res, 37, 3]) Template atom coordinates in atom37 format

”template_all_atom_mask” ([*, N_templ, N_res, 37])

Template atom coordinate mask

”template_pseudo_beta” ([*, N_templ, N_res, 3])

Positions of template carbon “pseudo-beta” atoms (i.e. C_beta for all residues but glycine, for for which C_alpha is used instead)

”template_pseudo_beta_mask” ([*, N_templ, N_res])

Pseudo-beta mask

iteration(feats, prevs, _recycle=True)
tolerance_reached(prev_pos, next_pos, mask, eps=1e-08)

Early stopping criteria based on criteria used in AF2Complex: https://www.nature.com/articles/s41467-022-29394-2 :param prev_pos: Previous atom positions in atom37/14 representation :param next_pos: Current atom positions in atom37/14 representation :param mask: 1-D sequence mask :param eps: Epsilon used in square root calculation

Returns:

Whether to stop recycling early based on the desired tolerance.

Return type:

bool