DOCSEARCHbyAlgolia
Submit Docs

egnn-pytorch

repository·main·Indexed 19 days ago

https://github.com/lucidrains/egnn-pytorch

An implementation of Equivariant Graph Neural Networks (EGNN) for processing geometric data, specifically optimized for molecular modeling and protein structures. It provides the EGNN layer for equivariant message passing and higher-level wrappers like EGNN_Network and EGNN_Sparse_Network. Key features include support for masking, sparse neighbor attention via adjacency matrices, Fourier features for relative distance encoding, and utilities for calculating protein covalent bonds and N-th degree adjacency matrices.

Tokens
2.5K
Snippets
8
Records
10
Agent score
18%

What's inside egnn-pytorch

  1. Improve EGNN stability for high neighbor counts

    main

    When using a high number of neighbors (num_nearest_neighbors), the architecture can become unstable. To mitigate this, use the following settings:

    1. Set norm_coors = True to normalize relative coordinates.
    2. Set coor_weights_clamp_value to an absolute clamped value for coordinate weights.
    import torch
from egnn_pytorch import EGNN_Network

net = EGNN_Network(
    num_tokens = 21,
    dim = 32,
    depth = 3,
    num_nearest_neighbors = 32,
    norm_coors = True,
    coor_weights_clamp_value = 2.
)

  3. Install EGNN and dependencies

    main

    To use the EGNN implementation, you need to install the core package along with several geometric deep learning dependencies. Note that the example uses sidechainnet and geometric-vector-perceptron as part of its environment setup.

    # Install core dependencies
!pip install sidechainnet proDy einops

# Install PyTorch Geometric dependencies (adjust versions for your CUDA/Torch setup)
!pip install torch-scatter -f https://pytorch-geometric.com/whl/torch-1.8.0+cu101.html
!pip install torch-sparse -f https://pytorch-geometric.com/whl/torch-1.8.0+cu101.html
!pip install torch-cluster -f https://pytorch-geometric.com/whl/torch-1.8.0+cu101.html
!pip install torch-spline-conv -f https://pytorch-geometric.com/whl/torch-1.8.0+cu101.html
!pip install torch-geometric

  4. Use the EGNN_Network for full models

    main

    EGNN_Network is a higher-level wrapper for building complete networks. It supports masking, sparse neighbor attention via adjacency matrices, and automatic Nth-order neighbor determination.

    Key features:

    • Masking: Pass a boolean mask to handle variable sequence lengths.
    • Sparse Neighbors: Set only_sparse_neighbors = True and provide an adj_mat to restrict message passing to specific connections.
    • Adjacency Embeddings: Use num_adj_degrees and adj_dim to automatically fetch Nth-order neighbors and pass adjacency degree embeddings to the edge MLP.
    • Continuous Edges: Pass edges for continuous edge features.
    import torch
from egnn_pytorch import EGNN_Network

# Full network with masking
net = EGNN_Network(
    num_tokens = 21,
    num_positions = 1024,
    dim = 32,
    depth = 3,
    num_nearest_neighbors = 8,
    coor_weights_clamp_value = 2.
)

feats = torch.randint(0, 21, (1, 1024))
coors = torch.randn(1, 1024, 3)
mask = torch.ones_like(feats).bool()

feats_out, coors_out = net(feats, coors, mask = mask)

  5. Use the EGNN layer

    main

    The EGNN class implements a single layer of the E(n)-Equivariant Graph Neural Network. It can be used with or without explicit edge features.

    Basic usage (no edges): Pass feats (features) and coors (coordinates).

    Usage with edges: If you provide edge_dim during initialization, you must pass edges during the forward pass.

    import torch
from egnn_pytorch import EGNN

# Basic usage
layer1 = EGNN(dim = 512)
feats = torch.randn(1, 16, 512)
coors = torch.randn(1, 16, 3)
feats, coors = layer1(feats, coors)

# Usage with edges
layer2 = EGNN(dim = 512, edge_dim = 4)
edges = torch.randn(1, 16, 16, 4)
feats, coors = layer2(feats, coors, edges)

  6. Configure EGNN layer parameters

    main

    The EGNN class accepts several parameters to control the behavior of the equivariant message passing:

    ParameterDescription
    dimInput dimension
    edge_dimDimension of the edges, if exists, should be > 0
    m_dimHidden model dimension
    fourier_featuresNumber of fourier features for encoding of relative distance
    num_nearest_neighborsCap the number of neighbors doing message passing by relative distance
    dropoutDropout rate
    norm_featsWhether to layernorm the features
    norm_coorsWhether to normalize the coordinates
    update_featsWhether to update features
    update_coorsWhether to update coordinates
    only_sparse_neighborsIf True, only allow message passing along adjacent neighbors via adj_mat
    valid_radiusThe valid radius each node considers for message passing
    m_pool_methodPooling method for output node representation ('mean' or 'sum')
    soft_edgesExtra GLU on the edges for stabilization
    coor_weights_clamp_valueClamping of the coordinate updates for stabilization
    model = EGNN(
    dim = 512,
    edge_dim = 0,
    m_dim = 16,
    fourier_features = 0,
    num_nearest_neighbors = 0,
    dropout = 0.0,
    norm_feats = False,
    norm_coors = False,
    update_feats = True,
    update_coors = True,
    only_sparse_neighbors = False,
    valid_radius = float('inf'),
    m_pool_method = 'sum',
    soft_edges = False,
    coor_weights_clamp_value = None
)

  7. Perform a forward pass with EGNN_Sparse_Network

    main

    The forward method of EGNN_Sparse_Network takes node features, edge indices, batch information, and edge attributes to predict updated states (including coordinates).

    Arguments:

    • x: Node features tensor.
    • edge_index: Tensor of shape [2, num_edges] representing connectivity.
    • batch: Batch index tensor.
    • edge_attr: Edge attribute tensor.
    • recalc_edge: If provided, overrides edge recalculation logic.
    • verbose: Boolean for logging.
    # preds will contain updated features and coordinates
preds = model.forward(
    x, 
    edge_index, 
    batch=batch, 
    edge_attr=edge_attrs, 
    recalc_edge=None, 
    verbose=False
)

  8. Initialize EGNN_Sparse_Network

    main

    The EGNN_Sparse_Network is the primary model class. It is configured with dimensions for features, positions, and edges, and supports Fourier features for encoding.

    Key parameters:

    • n_layers: Number of EGNN layers.
    • feats_dim: Dimension of node features.
    • pos_dim: Dimension of position coordinates (typically 3).
    • edge_attr_dim: Dimension of edge attributes.
    • m_dim: Dimension of the message passing space.
    • fourier_features: Number of Fourier features.
    • embedding_nums / embedding_dims: Lists defining the number and dimensions of embeddings for different components (e.g., amino acids, atoms).
    • update_coors / update_feats: Booleans to enable/disable coordinate and feature updates.
    • recalc: Boolean to determine if edges should be recalculated.
    model = EGNN_Sparse_Network(
    n_layers=4,
    feats_dim=2,
    pos_dim=3,
    edge_attr_dim=1,
    m_dim=32,
    fourier_features=4,
    embedding_nums=[36, 20],
    embedding_dims=[16, 16],
    edge_embedding_nums=[3],
    edge_embedding_dims=[2],
    update_coors=True,
    update_feats=True,
    norm_feats=False,
    norm_coors=False,
    recalc=False
)

  9. Compute N-th degree adjacency matrix

    main

    The nth_deg_adjacency function calculates the N-th degree adjacency matrix of a graph. It can operate on dense tensors or use torch-sparse for efficient sparse matrix multiplication.

    Arguments:

    • adj_mat: The base adjacency matrix.
    • n: The degree of connectivity to compute.
    • sparse: Boolean; if True, uses torch_sparse.spspmm for computation.

    Returns:

    • adj_mat: The resulting N-th degree adjacency matrix.
    • attr_mat: A matrix where values represent the degree of connectivity (e.g., 1 for neighbors, 2 for neighbors of neighbors).
    # Compute 2nd degree adjacency using sparse implementation
adj_mat, attr_mat = nth_deg_adjacency(adj_mat, n=2, sparse=True)

  10. Calculate protein covalent bonds

    main

    The prot_covalent_bond function identifies the indices and attributes of covalent bonds within a protein sequence. It can also compute N-th degree adjacency (e.g., neighbors of neighbors) to expand the connectivity graph.

    Arguments:

    • seq: Protein sequence string.
    • adj_degree: The degree of adjacency to compute (e.g., 1 for immediate neighbors, 2 for neighbors of neighbors).
    • cloud_mask: (Optional) A mask selecting present atoms.

    Returns:

    • edge_idxs: Tensor of shape [2, num_edges] containing the indices of the bonds.
    • edge_attrs: Tensor containing the connectivity degree/attribute for each edge.
    # Get immediate covalent bonds
edge_idxs, edge_attrs = prot_covalent_bond(seq, adj_degree=1)

# Get 2nd degree adjacency (neighbors of neighbors)
edge_idxs, edge_attrs = prot_covalent_bond(seq, adj_degree=2)