terminator.models.layers.graph_features.ProteinFeatures

class terminator.models.layers.graph_features.ProteinFeatures(edge_features, node_features, num_positional_embeddings=16, num_rbf=16, top_k=30, features_type='full', augment_eps=0.0, dropout=0.1)

Bases: Module

Protein backbone featurization based on Ingraham et al NeurIPS

Variables:

• embeddings (PositionalEncodings) – Module to generate differential positional embeddings for edges

• dropout (nn.Dropout) – Dropout module

• edge_embeddings (node_embeddings,) – Embedding layers for nodes and edges

• norm_edges (norm_nodes,) – Normalization layers for node and edge features

__init__(edge_features, node_features, num_positional_embeddings=16, num_rbf=16, top_k=30, features_type='full', augment_eps=0.0, dropout=0.1)

Extract protein features

Methods

__init__(edge_features, node_features[, ...])	Extract protein features
forward(X, mask)	Featurize coordinates as an attributed graph

Attributes

T_destination	alias of TypeVar('T_destination', bound=Mapping[str, Tensor])
dump_patches	This allows better BC support for load_state_dict().

_contacts(D_neighbors, mask_neighbors, cutoff=8)

Contacts

_dist(X, mask, eps=1e-06)

Pairwise euclidean distances

_hbonds(X, E_idx, mask_neighbors, eps=0.001)

Hydrogen bonds and contact map

_quaternions(R, eps=1e-10)

Convert a batch of 3D rotations [R] to quaternions [Q] R […,3,3] Q […,4]

forward(X, mask)

Featurize coordinates as an attributed graph

Parameters:

• X (torch.Tensor) – Backbone coordinates Shape: n_batch x seq_len x 4 x 3

• mask (torch.ByteTensor) – Mask for residues Shape: n_batch x seq_len

Returns:

• V (torch.Tensor) – Node embeddings Shape: n_batches x seq_len x n_hidden

• E (torch.Tensor) – Edge embeddings in kNN dense form Shape: n_batches x seq_len x k x n_hidden

• E_idx (torch.LongTensor) – Edge indices Shape: n_batches x seq_len x k x n_hidden