miniworld.models_MiniWorld_v1_5_use_interaction package

Submodules

miniworld.models_MiniWorld_v1_5_use_interaction.Attention_module module

class miniworld.models_MiniWorld_v1_5_use_interaction.Attention_module.Attention(d_query, d_key, n_head, d_hidden, d_out, p_drop=0.1)

Bases: Module

forward(query, key, value)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameter()

class miniworld.models_MiniWorld_v1_5_use_interaction.Attention_module.AttentionWithBias(d_in=256, d_bias=128, n_head=8, d_hidden=32)

Bases: Module

forward(x, bias)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameter()

class miniworld.models_MiniWorld_v1_5_use_interaction.Attention_module.BiasedAxialAttention(d_pair, d_bias, n_head, d_hidden, p_drop=0.1, is_row=True)

Bases: Module

forward(pair, sym=None, stride=-1)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameter()

class miniworld.models_MiniWorld_v1_5_use_interaction.Attention_module.FeedForwardLayer(d_model, r_ff, p_drop=0.1)

Bases: Module

forward(src)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameter()

class miniworld.models_MiniWorld_v1_5_use_interaction.Attention_module.GraphTriangleAttention(d_pair, d_state, top_k=64, d_rbf=64, n_head=8, n_query_pt=4, d_hidden=32)

Bases: Module

apply_RT(Rs, Ts, vec)

forward(pair_in, state, Rs, Ts, pair_mask=None, use_species=True)

Input:

• pair: pairwise features. shape [B, L+1, L+1, d_pair] if use_species is True, [B, L, L, d_pair] if use_species is False

• state: node features. shape [B, L+1, d_state] if use_species is True, [B, L, d_state] if use_species is False

• Rs, Ts : SE(3) Frame. shape [B, L, 3, 3], [B, L, 3]

• pair_mask: mask for valid residue pairs. shape [B, L+1, L+1] if use_species is True, [B, L, L] if use_species is False

gather_edges(edges, neighbor_idx)

gather_edges_3d(edges, neighbor_idx)

reset_parameter()

class miniworld.models_MiniWorld_v1_5_use_interaction.Attention_module.MSAColAttention(d_msa=256, n_head=8, d_hidden=32)

Bases: Module

forward(msa)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameter()

class miniworld.models_MiniWorld_v1_5_use_interaction.Attention_module.MSAColGlobalAttention(d_msa=64, n_head=8, d_hidden=8)

Bases: Module

forward(msa)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameter()

class miniworld.models_MiniWorld_v1_5_use_interaction.Attention_module.MSARowAttentionWithBias(d_msa=256, d_pair=128, n_head=8, d_hidden=32)

Bases: Module

forward(msa, pair, sym=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameter()

class miniworld.models_MiniWorld_v1_5_use_interaction.Attention_module.SequenceWeight(d_msa, n_head, d_hidden, p_drop=0.1)

Bases: Module

forward(msa)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameter()

class miniworld.models_MiniWorld_v1_5_use_interaction.Attention_module.StateAttentionGate(d_state, d_pair, n_head=8, d_hidden=32, p_drop=0.1)

Bases: Module

forward(state, pair, chain_mask)

state : (B, L, d_state) pair : (B, L, L, d_pair) chain_mask : (B, L, L)

reset_parameter()

class miniworld.models_MiniWorld_v1_5_use_interaction.Attention_module.TriangleMultiplication(d_pair, d_hidden=128, outgoing=True)

Bases: Module

forward(pair, sym=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameter()

miniworld.models_MiniWorld_v1_5_use_interaction.AuxiliaryPredictor module

class miniworld.models_MiniWorld_v1_5_use_interaction.AuxiliaryPredictor.BinderNetwork(n_hidden=64, n_bin_pae=64)

Bases: Module

forward(pae, same_chain)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameter()

class miniworld.models_MiniWorld_v1_5_use_interaction.AuxiliaryPredictor.DistanceNetwork(n_feat, p_drop=0.1)

Bases: Module

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameter()

class miniworld.models_MiniWorld_v1_5_use_interaction.AuxiliaryPredictor.ExpResolvedNetwork(d_msa, d_state, p_drop=0.1)

Bases: Module

forward(seq, state)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameter()

class miniworld.models_MiniWorld_v1_5_use_interaction.AuxiliaryPredictor.LDDTNetwork(n_feat, d_hidden=128, n_bin_lddt=50)

Bases: Module

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameter()

class miniworld.models_MiniWorld_v1_5_use_interaction.AuxiliaryPredictor.MaskedTokenNetwork(n_feat, p_drop=0.1)

Bases: Module

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameter()

class miniworld.models_MiniWorld_v1_5_use_interaction.AuxiliaryPredictor.PAENetwork(n_feat, n_bin_pae=64)

Bases: Module

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameter()

miniworld.models_MiniWorld_v1_5_use_interaction.Embeddings module

class miniworld.models_MiniWorld_v1_5_use_interaction.Embeddings.Extra_emb(d_msa=256, d_init=24, p_drop=0.1)

Bases: Module

forward(msa, seq, idx, oligo=1)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameter()

class miniworld.models_MiniWorld_v1_5_use_interaction.Embeddings.InteractionEmb(d_pair=128)

Bases: Module

forward(interaction_points)

Inputs:

• interaction_points: Interaction points (B, L, L) 0 (no information), 1 (positive), 2 (negative)

Outputs:

• emb: Interaction embedding (B, L, L, d_pair)

class miniworld.models_MiniWorld_v1_5_use_interaction.Embeddings.MSA_emb(d_msa=256, d_pair=128, d_state=32, d_init=24, minpos=-32, maxpos=32, p_drop=0.1)

Bases: Module

forward(msa, msa_species, msa_zero_pos, seq, idx, chain_mask, xyz, symmids=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameter()

class miniworld.models_MiniWorld_v1_5_use_interaction.Embeddings.PositionalEncoding2D(d_model, minpos=-31, maxpos=31, topk=2)

Bases: Module

forward(idx, chain_mask, xyz)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

get_chain_break_from_mask(chain_mask)

get_chain_embedding(xyz, chain_mask, topk=2)

class miniworld.models_MiniWorld_v1_5_use_interaction.Embeddings.Recycling(value_net, d_msa=256, d_pair=128, d_state=32, d_rbf=64)

Bases: Module

forward(target_seq, idx, chain_mask, msa, pair, state, xyz, mask_recycle=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class miniworld.models_MiniWorld_v1_5_use_interaction.Embeddings.Species_emb(d_msa=256, p_drop=0.1)

Bases: Module

forward(msa_speices)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameter()

class miniworld.models_MiniWorld_v1_5_use_interaction.Embeddings.Templ_emb(d_t1d=24, d_t2d=44, d_tor=30, d_pair=128, d_state=32, n_block=2, d_templ=64, n_head=4, d_hidden=16, p_drop=0.25)

Bases: Module

forward(t1d, t2d, alpha_t, xyz_t, mask_t, pair, state, use_checkpoint=False, p2p_crop=-1, symmids=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameter()

class miniworld.models_MiniWorld_v1_5_use_interaction.Embeddings.TemplatePairStack(n_block=2, d_templ=64, n_head=4, d_hidden=16, d_t1d=24, d_state=32, p_drop=0.25)

Bases: Module

forward(templ, rbf_feat, t1d, use_checkpoint=False, p2p_crop=-1, symmids=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class miniworld.models_MiniWorld_v1_5_use_interaction.Embeddings.TemplateTorsionStack(n_block=2, d_templ=64, d_rbf=64, n_head=4, d_hidden=16, p_drop=0.15)

Bases: Module

forward(tors, pair, rbf_feat, use_checkpoint=False)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameter()

miniworld.models_MiniWorld_v1_5_use_interaction.MiniWorld module

class miniworld.models_MiniWorld_v1_5_use_interaction.MiniWorld.MiniWorldModule(n_extra_block=4, n_main_block=8, n_ref_block=4, d_msa=256, d_pair=128, d_templ=64, d_rbf=64, n_head_msa=8, n_head_pair=4, n_head_templ=4, d_hidden=32, d_hidden_templ=64, p_drop=0.15, SE3_param_full={'l0_in_features': 32, 'l0_out_features': 16, 'num_edge_features': 32}, SE3_param_topk={'l0_in_features': 32, 'l0_out_features': 16, 'num_edge_features': 32}, value_net_param={'d_hidden': 128, 'd_node': 128, 'dropout': 0.1, 'num_layers': 3, 'topk': 16})

Bases: Module

forward(msa_latent=None, msa_species=None, msa_zero_pos=None, query_sequence=None, xyz=None, idx=None, interaction_points=None, template_1D=None, template_2D=None, template_xyz=None, template_alpha=None, mask_t=None, chain_mask=None, msa_prev=None, pair_prev=None, state_prev=None, mask_recycle=None, return_raw=False, return_full=False, use_checkpoint=False, p2p_crop=-1, topk_crop=-1, symmids=None, symmsub=None, symmRs=None, symmmeta=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

update_value_net(value_net)

miniworld.models_MiniWorld_v1_5_use_interaction.Track_module module

class miniworld.models_MiniWorld_v1_5_use_interaction.Track_module.IPA(d_seq=256, d_pair=128, d_state=16, n_head=12, d_hidden=16, n_query_points=4, n_point_values=8, top_k=64, p_drop=0.1)

Bases: Module

apply_RT(Rs, Ts, vec)

forward(seq, pair, state, Rs, Ts, chain_mask)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

gather_edges(edges, neighbor_idx)

gather_edges_3d(edges, neighbor_idx)

inverse_RT(Rs, Ts, vec)

reset_parameter()

class miniworld.models_MiniWorld_v1_5_use_interaction.Track_module.IterBlock(value_net, d_msa=256, d_pair=128, d_rbf=64, n_head_msa=8, n_head_pair=4, use_global_attn=False, d_hidden=32, d_hidden_msa=None, p_drop=0.15, SE3_param={'l0_in_features': 32, 'l0_out_features': 16, 'num_edge_features': 32})

Bases: Module

forward(seq, msa, pair, R_in, T_in, xyz_in, state, idx, chain_mask, symmids, symmsub_in, symmsub, symmRs, symmmeta, topk=0, crop=-1)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class miniworld.models_MiniWorld_v1_5_use_interaction.Track_module.IterativeSimulator(value_net, n_extra_block=4, n_main_block=12, n_ref_block=4, d_query=256, d_msa=256, d_pair=128, d_hidden=32, d_rbf=64, n_head_msa=8, n_head_pair=4, SE3_param_full={'l0_in_features': 32, 'l0_out_features': 16, 'l1_out_features': 2, 'num_edge_features': 32}, SE3_param_topk={'l0_in_features': 32, 'l0_out_features': 16, 'l1_out_features': 2, 'num_edge_features': 32}, p_drop=0.15)

Bases: Module

forward(seq, msa, pair, xyz_in, state, idx, chain_mask, symmids, symmsub, symmRs, symmmeta, use_checkpoint=False, p2p_crop=-1, topk_crop=-1)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameter()

class miniworld.models_MiniWorld_v1_5_use_interaction.Track_module.MSA2Pair(d_msa=256, d_pair=128, d_hidden=32, p_drop=0.15)

Bases: Module

forward(msa, pair, symmids, symmsub)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameter()

class miniworld.models_MiniWorld_v1_5_use_interaction.Track_module.MSAPairStr2MSA(d_msa=256, d_pair=128, n_head=8, d_state=16, d_rbf=64, d_hidden=32, p_drop=0.15, use_global_attn=False)

Bases: Module

forward(msa, pair, state, symmids, symmsub)

Update MSA with biased self-attention, using bias from Pair & Str.

Parameters:

• msa (torch.Tensor of shape (B, N, L, d_msa)) – MSA feature tensor.

• pair (torch.Tensor of shape (B, L, L, d_pair)) – Pair feature tensor.

• state (torch.Tensor of shape (B, L, d_state)) – Updated node features from the SE(3) layer.

• symmids (torch.Tensor) – Symmetry IDs.

• symmsub (torch.Tensor) – Symmetry sub-units.

Returns:

Updated MSA feature tensor.

Return type:

torch.Tensor of shape (B, N, L, d_msa)

reset_parameter()

class miniworld.models_MiniWorld_v1_5_use_interaction.Track_module.PairStr2Pair(d_pair=128, n_head=4, d_hidden=32, d_hidden_state=16, d_rbf=64, d_state=32, p_drop=0.15)

Bases: Module

forward(pair, crop=-1, use_species=True, symmids=None, symmsub=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

subblock(OP, pair, rbf_feat, crop)

class miniworld.models_MiniWorld_v1_5_use_interaction.Track_module.SCPred(d_msa=256, d_state=32, d_hidden=128, p_drop=0.15)

Bases: Module

forward(seq, state)

Predict side-chain torsion angles along with backbone torsions.

Parameters:

• seq (torch.Tensor of shape (B, L, d_msa)) – Hidden embeddings corresponding to the query sequence.

• state (torch.Tensor of shape (B, L, d_state)) – State features (output l0 feature) from the previous SE(3) layer.

Returns:

Predicted torsion angles (phi, psi, omega, chi1-4 with cos/sin, Cb bend, Cb twist, CG).

Return type:

torch.Tensor of shape (B, L, 10, 2)

reset_parameter()

class miniworld.models_MiniWorld_v1_5_use_interaction.Track_module.SeqSep(d_model, minpos=-32, maxpos=32)

Bases: Module

forward(idx)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class miniworld.models_MiniWorld_v1_5_use_interaction.Track_module.Str2Chain_IPA(d_msa=256, d_pair=128, d_state=16, n_head=8, d_hidden=32, n_query_points=4, n_point_values=8, p_drop=0.1)

Bases: Module

forward(msa, pair_in, R_in, T_in, state, chain_mask)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameter()

class miniworld.models_MiniWorld_v1_5_use_interaction.Track_module.Str2Residue_IPA(d_msa=256, d_pair=128, d_state=16, d_rbf=64, n_head=8, d_hidden=32, n_query_points=4, n_point_values=8, p_drop=0.1)

Bases: Module

forward(msa, pair_in, R_in, T_in, state, chain_mask)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameter()

class miniworld.models_MiniWorld_v1_5_use_interaction.Track_module.UpdatePairByValue(d_pair=128, d_rbf=64)

Bases: Module

forward(xyz, pair, h_E, pae_neighbor, E_idx)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameter()

miniworld.models_MiniWorld_v1_5_use_interaction.Track_module.update_symm_Rs(Rs, Ts, Lasu, symmsub_in, symmsub, symmRs)

miniworld.models_MiniWorld_v1_5_use_interaction.Track_module.update_symm_subs(Rs, Ts, pair, symmids, symmsub_in, symmsub, symmRs, metasymm)

miniworld.models_MiniWorld_v1_5_use_interaction.value_network module

class miniworld.models_MiniWorld_v1_5_use_interaction.value_network.MPNN_block(d_node, d_edge, d_hidden, dropout=0.1)

Bases: Module

forward(h_V, h_E, E_idx, mask_V=None, mask_attend=None)

Inputs:

• h_V: input node features (B, L, d_node)

• h_E: input edge features (B, L, topK, d_edge)

• E_idx: indices of topK neighbors (B, L, topK)

• mask_V: mask for missing residue

• mask_attend: mask for message passing to ignore messages from missing residues

Outputs:

• h_V: updated node features (B, L, d_node)

• h_E: updated edge features (B, L, topK, d_edge)

reset_parameter()

class miniworld.models_MiniWorld_v1_5_use_interaction.value_network.PositionWiseFeedForward(num_hidden, num_ff)

Bases: Module

forward(h_V)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameter()

class miniworld.models_MiniWorld_v1_5_use_interaction.value_network.StrSeqValueNet(d_node=128, d_rbf=64, d_hidden=128, num_layers=3, topk=16, dropout=0.1, n_bin_pae=32)

Bases: Module

forward(seq, idx, chain_mask, xyz, mask_V=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameter()

turnoff()

turnon()

miniworld.models_MiniWorld_v1_5_use_interaction.value_network.cat_neighbors_nodes(h_nodes, h_neighbors, E_idx)

miniworld.models_MiniWorld_v1_5_use_interaction.value_network.gather_edges(edges, neighbor_idx)

miniworld.models_MiniWorld_v1_5_use_interaction.value_network.gather_nodes(nodes, neighbor_idx)

Module contents