openfold.model.structure_module

Classes

AngleResnet(c_in, c_hidden, no_blocks, ...)	Implements Algorithm 20, lines 11-14
AngleResnetBlock(c_hidden)
BackboneUpdate(c_s)	Implements part of Algorithm 23.
InvariantPointAttention(c_s, c_z, c_hidden, ...)	Implements Algorithm 22.
InvariantPointAttentionMultimer(c_s, c_z, ...)	Implements Algorithm 22.
PointProjection(c_hidden, num_points, ...[, ...])
StructureModule(c_s, c_z, c_ipa, c_resnet, ...)
StructureModuleTransition(c, num_layers, ...)
StructureModuleTransitionLayer(c)

class AngleResnet(c_in, c_hidden, no_blocks, no_angles, epsilon)

Bases: Module

Implements Algorithm 20, lines 11-14

forward(s, s_initial)

Parameters:

• s (Tensor) – [*, C_hidden] single embedding

• s_initial (Tensor) – [*, C_hidden] single embedding as of the start of the StructureModule

Returns:

[*, no_angles, 2] predicted angles

Return type:

Tuple[Tensor, Tensor]

class AngleResnetBlock(c_hidden)

Bases: Module

forward(a)

Parameters:

a (Tensor)

Return type:

Tensor

class BackboneUpdate(c_s)

Bases: Module

Implements part of Algorithm 23.

forward(s)

Parameters:

• [*

• N_res

• representation (C_s] single)

• s (Tensor)

Returns:

[*, N_res, 6] update vector

Return type:

Tuple[Tensor, Tensor]

class InvariantPointAttention(c_s, c_z, c_hidden, no_heads, no_qk_points, no_v_points, inf=100000.0, eps=1e-08, is_multimer=False)

Bases: Module

Implements Algorithm 22.

Parameters:

• c_s (int)

• c_z (int)

• c_hidden (int)

• no_heads (int)

• no_qk_points (int)

• no_v_points (int)

• inf (float)

• eps (float)

• is_multimer (bool)

forward(s, z, r, mask, inplace_safe=False, _offload_inference=False, _z_reference_list=None)

Parameters:

• s (Tensor) – [*, N_res, C_s] single representation

• z (Tensor) – [*, N_res, N_res, C_z] pair representation

• r (Rigid | Rigid3Array) – [*, N_res] transformation object

• mask (Tensor) – [*, N_res] mask

• inplace_safe (bool)

• _offload_inference (bool)

• _z_reference_list (Sequence[Tensor] | None)

Returns:

[*, N_res, C_s] single representation update

Return type:

Tensor

class InvariantPointAttentionMultimer(c_s, c_z, c_hidden, no_heads, no_qk_points, no_v_points, inf=100000.0, eps=1e-08, is_multimer=True)

Bases: Module

Implements Algorithm 22.

Parameters:

• c_s (int)

• c_z (int)

• c_hidden (int)

• no_heads (int)

• no_qk_points (int)

• no_v_points (int)

• inf (float)

• eps (float)

• is_multimer (bool)

forward(s, z, r, mask, inplace_safe=False, _offload_inference=False, _z_reference_list=None)

Parameters:

• s (Tensor) – [*, N_res, C_s] single representation

• z (Tensor | None) – [*, N_res, N_res, C_z] pair representation

• r (Rigid | Rigid3Array) – [*, N_res] transformation object

• mask (Tensor) – [*, N_res] mask

• inplace_safe (bool)

• _offload_inference (bool)

• _z_reference_list (Sequence[Tensor] | None)

Returns:

[*, N_res, C_s] single representation update

Return type:

Tensor

class PointProjection(c_hidden, num_points, no_heads, is_multimer, return_local_points=False)

Bases: Module

Parameters:

• c_hidden (int)

• num_points (int)

• no_heads (int)

• is_multimer (bool)

• return_local_points (bool)

forward(activations, rigids)

Parameters:

• activations (Tensor)

• rigids (Rigid | Rigid3Array)

Return type:

Tensor | Tuple[Tensor, Tensor]

class StructureModule(c_s, c_z, c_ipa, c_resnet, no_heads_ipa, no_qk_points, no_v_points, dropout_rate, no_blocks, no_transition_layers, no_resnet_blocks, no_angles, trans_scale_factor, epsilon, inf, is_multimer=False, **kwargs)

Bases: Module

forward(evoformer_output_dict, aatype, mask=None, inplace_safe=False, _offload_inference=False)

Parameters:

• s – [*, N_res, C_s] single representation

• z – [*, N_res, N_res, C_z] pair representation

• aatype – [*, N_res] amino acid indices

• mask – Optional [*, N_res] sequence mask

Returns:

A dictionary of outputs

frames_and_literature_positions_to_atom14_pos(r, f)

torsion_angles_to_frames(r, alpha, f)

class StructureModuleTransition(c, num_layers, dropout_rate)

Bases: Module

forward(s)

class StructureModuleTransitionLayer(c)

Bases: Module

forward(s)