### Install dependencies and run denoising example

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/README.md

Commands to install the required sidechainnet package and execute the sparse denoising script.

```bash
$ pip install sidechainnet
```

```bash
$ python denoise_sparse.py
```

--------------------------------

### Install Python Packages

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/examples/egnn_test.ipynb

Installs essential Python libraries for the project. Ensure you have pip installed.

```python
!pip install sidechainnet proDy einops
```

--------------------------------

### Install egnn-pytorch

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/README.md

Install the egnn-pytorch library using pip.

```bash
pip install egnn-pytorch
```

--------------------------------

### Prepare for Training

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/examples/egnn_test.ipynb

Sets the model to training mode and initializes lists to store baseline and epoch losses. Also records the start time for epoch duration calculation.

```python
model.train()

# info records
baseline_losses = []
epoch_losses    = []
tac    = time.time()
```

--------------------------------

### Install PyTorch Geometric Dependencies

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/examples/egnn_test.ipynb

Install the required PyTorch Geometric packages using pip with specific wheel links for the environment.

```bash
!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
```

--------------------------------

### Initialize SidechainNet Dataloader

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/examples/egnn_test.ipynb

Load protein data using SidechainNet and prepare training examples.

```python
dataloaders_ = sidechainnet.load(casp_version=7, with_pytorch="dataloaders")
dataloaders_.keys() # ['train', 'train_eval', 'valid-10', ..., 'valid-90', 'test']
train_examples_storer = [get_prot(dataloader_=dataloaders_, vocab_=VOCAB,
                                  min_len=MIN_LEN, max_len=MAX_LEN, verbose=0)\
                         for i in tqdm(range(3))]
```

--------------------------------

### Check PyTorch Version

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/examples/egnn_test.ipynb

Verify the currently installed version of PyTorch.

```python
import torch
torch.__version__
```

--------------------------------

### Basic EGNN Layer Usage

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/README.md

Instantiate and use EGNN layers for feature and coordinate transformation. No specific setup or imports are required beyond PyTorch and the EGNN module.

```python
import torch
from egnn_pytorch import EGNN

layer1 = EGNN(dim = 512)
layer2 = EGNN(dim = 512)

feats = torch.randn(1, 16, 512)
coors = torch.randn(1, 16, 3)

feats, coors = layer1(feats, coors)
feats, coors = layer2(feats, coors) # (1, 16, 512), (1, 16, 3)
```

--------------------------------

### Train EGNN for Protein Backbone Denoising

Source: https://context7.com/lucidrains/egnn-pytorch/llms.txt

Example of training an EGNN network for protein backbone coordinate denoising using SidechainNet. Requires PyTorch, SidechainNet, and EGNN-PyTorch. Ensure correct data loading and preprocessing, including coordinate extraction and noise addition.

```python
import torch
import torch.nn.functional as F
from torch.optim import Adam
from einops import rearrange, repeat
import sidechainnet as scn
from egnn_pytorch import EGNN_Network

torch.set_default_dtype(torch.float64)
BATCH_SIZE = 1
GRADIENT_ACCUMULATE_EVERY = 16

def cycle(loader, len_thres=200):
    """Cycle through dataloader, filtering by sequence length."""
    while True:
        for data in loader:
            if data.seqs.shape[1] > len_thres:
                continue
            yield data

# Create EGNN network for denoising
net = EGNN_Network(
    num_tokens=21,
    num_positions=200 * 3,  # max positions for backbone atoms
    depth=5,
    dim=8,
    num_nearest_neighbors=16,
    fourier_features=2,
    norm_coors=True,
    coor_weights_clamp_value=2.
).cuda()

# Load protein data
data = scn.load(
    casp_version=12,
    thinning=30,
    with_pytorch='dataloaders',
    batch_size=BATCH_SIZE,
    dynamic_batching=False
)
dl = cycle(data['train'])
optim = Adam(net.parameters(), lr=1e-3)

# Training loop
for step in range(10000):
    for _ in range(GRADIENT_ACCUMULATE_EVERY):
        batch = next(dl)
        seqs, coords, masks = batch.seqs, batch.crds, batch.msks

        seqs = seqs.cuda().argmax(dim=-1)
        coords = coords.cuda().type(torch.float64)
        masks = masks.cuda().bool()

        # Extract backbone coordinates (N, CA, C atoms)
        coords = rearrange(coords, 'b (l s) c -> b l s c', s=14)
        coords = coords[:, :, 0:3, :]
        coords = rearrange(coords, 'b l s c -> b (l s) c')

        # Expand sequence and mask for backbone atoms
        seq = repeat(seqs, 'b n -> b (n c)', c=3)
        masks = repeat(masks, 'b n -> b (n c)', c=3)

        # Chain adjacency matrix
        i = torch.arange(seq.shape[-1], device=seq.device)
        adj_mat = (i[:, None] >= (i[None, :] - 1)) & (i[:, None] <= (i[None, :] + 1))

        # Add noise and denoise
        noised_coords = coords + torch.randn_like(coords)
        feats, denoised_coords = net(seq, noised_coords, adj_mat=adj_mat, mask=masks)

        # MSE loss on masked coordinates
        loss = F.mse_loss(denoised_coords[masks], coords[masks])
        (loss / GRADIENT_ACCUMULATE_EVERY).backward()

    print(f'Step {step}, Loss: {loss.item():.6f}')
    optim.step()
    optim.zero_grad()
```

--------------------------------

### Prepare Local Data Directory

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/examples/egnn_test.ipynb

Creates a local directory and copies the dataset from the mounted drive.

```bash
!mkdir sidechainnet_data
!cp drive/MyDrive/sidechainnet_casp7_30.pkl sidechainnet_data/sidechainnet_casp7_30.pkl
```

--------------------------------

### Initialize Optimizer and Noise Level

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/examples/egnn_test.ipynb

Sets up the Adam optimizer for model parameters and defines an initial noise level. The learning rate is set to 1e-3.

```python
noise = 1
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
```

--------------------------------

### Initialize EGNN_Network with stability settings

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/README.md

Configures an EGNN_Network with coordinate normalization and weight clamping to handle high neighbor counts.

```python
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,              # normalize the relative coordinates
    coor_weights_clamp_value = 2.   # absolute clamped value for the coordinate weights, needed if you increase the num neareest neighbors
)

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

feats_out, coors_out = net(feats, coors, mask = mask) # (1, 1024, 32), (1, 1024, 3)
```

--------------------------------

### Configure System Paths

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/examples/egnn_test.ipynb

Appends the project directories to the system path to allow importing custom modules.

```python
import os
import sys
import time

sys.path.append("geometric-vector-perceptron/geometric_vector_perceptron")
sys.path.append("geometric-vector-perceptron/examples")
sys.path.append("egnn-pytorch/egnn_pytorch")
```

--------------------------------

### Import Dependencies and Custom Modules

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/examples/egnn_test.ipynb

Imports standard libraries, PyTorch components, and custom project modules.

```python
import torch
import torch_sparse
import numpy as np
from einops import rearrange, repeat
import matplotlib.pyplot as plt
# custom model
from egnn_pytorch import *
from geometric_vector_perceptron import *
# custom utils
from data_handler import *
from data_utils import *
# process and dataset
import gc
import joblib
from tqdm import tqdm
from functools import partial
```

--------------------------------

### Visualize Input Structure with sidechainnet

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/examples/egnn_test.ipynb

Builds and visualizes a 3D structure using sidechainnet's StructureBuilder with input coordinates. Requires `sidechainnet` library and `rearrange` function.

```python
# label
sb = sidechainnet.StructureBuilder( torch.tensor([AAS.index(x) for x in seq[:-padding_seq or None]]),
                                    crd=rearrange(input_rebuilt, 'l c d -> (l c) d') ) # pred_rebuilt
sb.to_3Dmol()
```

--------------------------------

### Initialize Protein Vocabulary

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/examples/egnn_test.ipynb

Loads the protein vocabulary required for sequence processing.

```python
# data
import sidechainnet
from sidechainnet.utils.sequence import ProteinVocabulary as VOCAB
VOCAB = VOCAB()
```

--------------------------------

### Run library tests

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/README.md

Executes the test suite for the library, requiring PyTorch Geometric.

```bash
$ python setup.py test
```

--------------------------------

### Move Model to Device

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/examples/egnn_test.ipynb

Transfers the initialized model to the appropriate device (e.g., GPU). This is a standard step before training or inference to leverage hardware acceleration.

```python
model = model.to(device) # .double()
```

--------------------------------

### Full EGNN Network Initialization

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/README.md

Initialize a complete EGNN network with specified token, position, dimension, depth, and neighbor parameters. The `coor_weights_clamp_value` is important when using a high number of nearest neighbors.

```python
import torch
from egnn_pytorch import EGNN_Network

net = EGNN_Network(
    num_tokens = 21,
    num_positions = 1024,           # unless what you are passing in is an unordered set, set this to the maximum sequence length
    dim = 32,
    depth = 3,
    num_nearest_neighbors = 8,
    coor_weights_clamp_value = 2.   # absolute clamped value for the coordinate weights, needed if you increase the num neareest neighbors
)

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

feats_out, coors_out = net(feats, coors, mask = mask) # (1, 1024, 32), (1, 1024, 3)
```

--------------------------------

### Set Device Configuration

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/examples/egnn_test.ipynb

Detects and sets the computation device to CUDA if available, otherwise defaults to CPU.

```python
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # "cpu")
```

--------------------------------

### Initialize and Use EGNN_Network

Source: https://context7.com/lucidrains/egnn-pytorch/llms.txt

Instantiate a complete multi-layer EGNN network, recommended for most use cases. This class includes token and positional embeddings, and handles neighbor selection automatically. Ensure the input features match the vocabulary size and sequence length.

```python
import torch
from egnn_pytorch import EGNN_Network

# Create a complete EGNN network
net = EGNN_Network(
    num_tokens=21,          # vocabulary size for token embedding
    num_positions=1024,     # max sequence length for positional embedding
    dim=32,                 # feature dimension
    depth=3,                # number of EGNN layers
    num_nearest_neighbors=8,  # k-nearest neighbors for message passing
    coor_weights_clamp_value=2.  # clamp coordinate updates for stability
)

# Input: sequence of token IDs and 3D coordinates
feats = torch.randint(0, 21, (1, 1024))  # token IDs
coors = torch.randn(1, 1024, 3)          # 3D coordinates
mask = torch.ones_like(feats).bool()     # attention mask

# Forward pass
feats_out, coors_out = net(feats, coors, mask=mask)
# feats_out: (1, 1024, 32), coors_out: (1, 1024, 3)
```

--------------------------------

### Plot Loss Evolution and Window Means

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/examples/egnn_test.ipynb

Visualizes the training loss over epochs, including rolling window means and a baseline comparison. Requires matplotlib and numpy. The plot shows reconstruction loss, windowed averages, and a dashed line for the baseline loss.

```python
plt.figure(figsize=(15,6))
plt.title("Loss Evolution - Denoising of Gaussian-masked Coordinates (mu=0, sigma=1)")
plt.plot(epoch_losses, label="Reconstruction loss")

for window in [8,16,32]:
    plt.plot([np.mean(epoch_losses[:window][0:i+1]) for i in range(min(window, len(epoch_losses))) ] + \
             [np.mean(epoch_losses[i:i+window+1]) for i in range(len(epoch_losses)-window)], \
             label="Window mean n={0}".format(window))

plt.plot(np.ones(len(epoch_losses)) * np.mean(baseline_losses), "k--", label="Baseline")
    
# plt.yscale("log")
plt.xlim(-0.01*len(epoch_losses),1.01*len(epoch_losses))
# plt.yscale("log")
plt.ylabel("Reconstruction RMSD")
plt.xlabel("Batch number")
plt.legend()
plt.show()
```

--------------------------------

### Visualize Predicted Structure with sidechainnet

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/examples/egnn_test.ipynb

Builds and visualizes a 3D structure using sidechainnet's StructureBuilder with predicted coordinates. Requires `sidechainnet` library and `rearrange` function.

```python
# label
sb = sidechainnet.StructureBuilder( torch.tensor([AAS.index(x) for x in seq[:-padding_seq or None]]),
                                    crd=rearrange(pre_target_rebuilt, 'l c d -> (l c) d') ) # pred_rebuilt
sb.to_3Dmol()
```

--------------------------------

### Define EGNN model parameters

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/README.md

Lists all available configuration parameters for the EGNN model class.

```python
import torch
from egnn_pytorch import EGNN

model = EGNN(
    dim = dim,                         # input dimension
    edge_dim = 0,                      # dimension of the edges, if exists, should be > 0
    m_dim = 16,                        # hidden model dimension
    fourier_features = 0,              # number of fourier features for encoding of relative distance - defaults to none as in paper
    num_nearest_neighbors = 0,         # cap the number of neighbors doing message passing by relative distance
    dropout = 0.0,                     # dropout
    norm_feats = False,                # whether to layernorm the features
    norm_coors = False,                # whether to normalize the coordinates, using a strategy from the SE(3) Transformers paper    
    update_feats = True,               # whether to update features - you can build a layer that only updates one or the other
    update_coors = True,               # whether ot update coordinates
    only_sparse_neighbors = False,     # using this would only allow message passing along adjacent neighbors, using the adjacency matrix passed in 
    valid_radius = float('inf'),       # the valid radius each node considers for message passing
    m_pool_method = 'sum',             # whether to mean or sum pool for output node representation
    soft_edges = False,                # extra GLU on the edges, purportedly helps stabilize the network in updated version of the paper
    coor_weights_clamp_value = None    # clamping of the coordinate updates, again, for stabilization purposes
)
```

--------------------------------

### Define EGNN Sparse Network Model

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/examples/egnn_test.ipynb

Initializes the EGNN_Sparse_Network with specific layer, feature, position, and embedding dimensions. Ensure all required parameters are correctly set for your task.

```python
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)
```

--------------------------------

### Create and Use EGNN_Sparse_Network

Source: https://context7.com/lucidrains/egnn-pytorch/llms.txt

Instantiate and perform a forward pass with the EGNN_Sparse_Network for processing graph data with sparse connectivity. Requires concatenated input of coordinates, features, and token IDs.

```python
import torch
from egnn_pytorch import EGNN_Sparse_Network

# Create complete sparse EGNN network
net = EGNN_Sparse_Network(
    n_layers=5,               # number of MPNN layers
    feats_dim=128,            # feature dimension
    pos_dim=3,                # coordinate dimension
    edge_attr_dim=4,          # edge attribute dimension
    m_dim=16,                 # hidden message dimension
    fourier_features=2,       # fourier encoding features
    soft_edge=0,              # soft edge gating
    embedding_nums=[21],      # vocabulary sizes for embeddings
    embedding_dims=[16],      # embedding dimensions
    edge_embedding_nums=[],   # edge vocabulary sizes
    edge_embedding_dims=[],   # edge embedding dimensions
    update_coors=True,        # update coordinates
    update_feats=True,        # update features
    norm_feats=True,          # normalize features
    norm_coors=False,         # normalize coordinates
    dropout=0.,               # dropout rate
    coor_weights_clamp_value=None,  # clamp weights
    aggr="add",               # aggregation method
    recalc=0                  # recalculate edges every N layers
)

# Input: concatenated [coordinates, features, token_ids]
n_nodes = 100
x = torch.randn(n_nodes, 3 + 127)  # coordinates + features
token_ids = torch.randint(0, 21, (n_nodes, 1)).float()
x = torch.cat([x, token_ids], dim=-1)  # append token for embedding

# Sparse connectivity
edge_index = torch.randint(0, n_nodes, (2, 500))
edge_attr = torch.randn(500, 4)
batch = torch.zeros(n_nodes, dtype=torch.long)

# Forward pass
x_out = net(x, edge_index, batch, edge_attr)
# x_out: (n_nodes, pos_dim + feats_dim + embedding_dim - 1)
```

--------------------------------

### Mount Google Drive for Data Access

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/examples/egnn_test.ipynb

Mounts the Google Drive directory to access stored datasets.

```python
# get CASP7 data copied from drive

from google.colab import drive
drive.mount('/content/drive/')
```

--------------------------------

### Clone Geometric Vector Perceptron Repository

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/examples/egnn_test.ipynb

Clones the geometric-vector-perceptron repository, which may contain related utilities or models.

```bash
!git clone https://github.com/hypnopump/geometric-vector-perceptron
```

--------------------------------

### Prepare Protein Data for EGNN

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/examples/egnn_test.ipynb

Configures encoding parameters and processes protein data for input into an EGNN model.

```python
NEEDED_INFO = {"cutoffs": [], # "15_closest"
               "bond_scales": [1, 2],
               "aa_pos_scales": [2,4,8,16,32,64,128],
               "atom_pos_scales": [1,2,4,8,16,32],
               "dist2ca_norm_scales": [1,2,4],
               "bb_norms_atoms": [0.5], # will encode 3 vectors with this
               # nn-degree connection
               "adj_degree": 2
              }
# get model sizes from encoded protein
seq, true_coords, angles, padding_seq, mask, id = train_examples_storer[-1] 
NEEDED_INFO["seq"] = seq[:-padding_seq or None]
NEEDED_INFO["covalent_bond"] = prot_covalent_bond(seq)


# encode as needed
encoded = encode_whole_protein(seq, true_coords, padding_seq, needed_info=NEEDED_INFO, free_mem=True)
x, edge_index, edge_attrs, embedd_info = encoded
# add position coords
cloud_mask = scn_cloud_mask(seq)
if padding_seq:
    cloud_mask[-padding_seq:] = 0.
cloud_mask = cloud_mask.bool()
flat_cloud_mask = rearrange(cloud_mask, 'l c -> (l c)')

x = torch.cat( [ true_coords[flat_cloud_mask], x ], dim=-1 )

### adjust for egnn: 
embedd_info["bond_n_scalars"] -= 2*len(NEEDED_INFO["bond_scales"])+1
embedd_info["bond_n_vectors"] = 0
embedd_info
```

--------------------------------

### EGNN_Sparse Layer for PyTorch Geometric

Source: https://context7.com/lucidrains/egnn-pytorch/llms.txt

Initialize the `EGNN_Sparse` layer, compatible with PyTorch Geometric, for efficient sparse graph processing. Input features should be concatenated coordinates and node features.

```python
import torch
from egnn_pytorch import EGNN_Sparse

# Create sparse EGNN layer (requires torch_geometric)
layer = EGNN_Sparse(
    feats_dim=128,          # feature dimension
    pos_dim=3,              # coordinate dimension
    edge_attr_dim=4,        # edge attribute dimension
    m_dim=16,               # hidden message dimension
    fourier_features=0,     # fourier encoding features
    soft_edge=0,            # soft edge gating
    norm_feats=False,       # normalize features
    norm_coors=False,       # normalize coordinates
    update_feats=True,      # update features
    update_coors=True,      # update coordinates
    dropout=0.,             # dropout rate
    coor_weights_clamp_value=None,  # clamp weights
    aggr="add"              # aggregation: 'add', 'mean', 'max'
)

# Input format: concatenated [coordinates, features]
n_nodes = 100
x = torch.randn(n_nodes, 3 + 128)  # (n_nodes, pos_dim + feats_dim)

# Sparse edge index (COO format)
edge_index = torch.randint(0, n_nodes, (2, 500))  # (2, n_edges)
edge_attr = torch.randn(500, 4)  # (n_edges, edge_attr_dim)
batch = torch.zeros(n_nodes, dtype=torch.long)  # batch assignment

# Forward pass
x_out = layer(x, edge_index, edge_attr=edge_attr, batch=batch)
# x_out: (n_nodes, pos_dim + feats_dim)
```

--------------------------------

### Protein Training Loop Implementation

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/examples/egnn_test.ipynb

This loop handles data loading, coordinate masking, edge filtering, model prediction, and backpropagation for protein structure training.

```python
max_seq_len = MAX_LEN # 120
iteration   = 0
n_per_iter  = MAX_PROTS # 400
for ep in range( 100 ): # 1*n_per_iter
    # delete useless data from prev iter - but not in last one
    if ep > 0:
        del true_coords, angles, edge_index, edge_attrs
        del target_coords, pred_coords, base_coords
        del encoded, target_aligned, pred_aligned
        gc.collect()
        
        
    # get model sizes from encoded protein
    seq, true_coords, angles, padding_seq, mask, pid = get_prot(dataloader_=dataloaders_,
                                                           vocab_=VOCAB, 
                                                           min_len=500,
                                                           max_len=MAX_LEN, 
                                                           verbose=0)
    NEEDED_INFO["seq"] = seq[:-padding_seq or None]
    NEEDED_INFO["covalent_bond"] = prot_covalent_bond(seq)
        
    # pass to device
    true_coords = true_coords.to(device) # .double()
    angles = angles.to(device)           # .double()

    # encode as needed
    masked_coords = true_coords + noise * torch.randn_like(true_coords) # (*2-1)
    encoded = encode_whole_protein(seq, true_coords, padding_seq, needed_info=NEEDED_INFO, free_mem=True)
    x, edge_index, edge_attrs, embedd_info = encoded
    # add position coords - better mask accounting for missing atoms
    cloud_mask_naive = scn_cloud_mask(seq).bool()
    cloud_mask = scn_cloud_mask(seq, coords=true_coords).bool()
    if padding_seq:
      cloud_mask[-padding_seq:] = 0.
    # cloud is all points, chain is all for which we have labels
    chain_mask = mask.unsqueeze(-1) * cloud_mask
    flat_chain_mask = rearrange(chain_mask, 'l c -> (l c)')
    flat_cloud_mask = rearrange(cloud_mask, 'l c -> (l c)')
    # slice useless norm and vector embeddings
    masked_coords = masked_coords[flat_cloud_mask]

    #############
    # MASK EDGES AND NODES ACCOUNTING FOR SCN MISSING ATOMS
    #############
    # NODES
    x = torch.cat( [ masked_coords, x[:, -2:][cloud_mask[cloud_mask_naive]] ], dim=-1 )
    # EDGES: delete all edges with masked-out atoms

    # pick all current indexes and turn them to 1.
    to_mask_edges = torch.zeros(edge_index.amax()+1, edge_index.amax()+1).to(edge_index.device)
    to_mask_edges[edge_index[0], edge_index[1]] = 1.
    # delete erased bonds
    masked_out_atoms = (-1*(cloud_mask[cloud_mask_naive].float() - 1)).bool()
    to_mask_edges[masked_out_atoms] *= 0.
    to_mask_edges = to_mask_edges * to_mask_edges.t()
    # get mask for the edge_attrs
    attr_mask = to_mask_edges[edge_index[0], edge_index[1]].bool()
    edge_attrs = edge_attrs[attr_mask, :]
    # delete unwanted rows and cols
    wanted = to_mask_edges.sum(dim=-1).bool()
    edge_index = (to_mask_edges[wanted, :][:, wanted]).nonzero().t()
    #############
    # continue
    #############
    edge_attrs = edge_attrs[:, -1:]
    batch = torch.tensor([0 for i in range(x.shape[0])], device=device).long()
    
    if torch.amax(edge_index) >= x.shape[0]:
      print("wtf, breaking, debug, index out of bounds")
      break

    # predict
    preds = model.forward(x, edge_index, batch=batch, edge_attr=edge_attrs,
                          recalc_edge=None, verbose = False)
    
    # MEASURE ERROR - format pred and target
    target_coords = true_coords[flat_cloud_mask].clone()
    pred_coords   = preds[:, :3]
    base_coords   = x[:, :3]
    
    # option 2: loss is RMSD on reconstructed coords  // align - sometimes svc fails - idk why
    try:
        pred_aligned, target_aligned = kabsch_torch(pred_coords.t(), target_coords.t()) # (3, N)

        loss = ( (pred_aligned.t() - target_aligned.t())[flat_chain_mask[flat_cloud_mask]]**2 ).mean() 
    except:
        pred_aligned, target_aligned = None, None
        print("svd failed convergence, ep:", ep)
        loss = ( (pred_coords - target_coords)[flat_chain_mask[flat_cloud_mask]]**2 ).mean()
    # measure error
    loss_base = ((base_coords - target_coords)**2).mean() 
    # not aligned: # loss = ((pred_coords - target_coords)**2).mean()**0.5 
    
    # back pass
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    # records / prints
    iteration += 1
    epoch_losses.append( loss.item() )
    baseline_losses.append( loss_base.item() )

    n_print = 10
    if iteration % n_print == 1:
        tic = time.time()
        print("BATCH: {0} / {1}, loss: {2}, baseline_loss: {3}, time: {4}".format(iteration, n_per_iter,
                                                                                  np.mean(epoch_losses[-n_print:]),
                                                                                  baseline_losses[-1],
                                                                                  tic-tac))
        tac = time.time()
        if iteration % n_per_iter == 1:
            print("---------------------------------")
```

--------------------------------

### Clone EGNN PyTorch Repository

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/examples/egnn_test.ipynb

Clones the main EGNN PyTorch repository from GitHub. This repository contains the core implementation.

```bash
!git clone https://github.com/hypnopump/egnn-pytorch
```

--------------------------------

### Initialize EGNN Layer with Edge Features

Source: https://context7.com/lucidrains/egnn-pytorch/llms.txt

Configure an EGNN layer to incorporate edge features, which are concatenated with node pair features and distance information. This allows for richer message passing between connected nodes.

```python
import torch
from egnn_pytorch import EGNN

# Create EGNN layer with edge features
layer = EGNN(dim=512, edge_dim=4)

# Inputs
feats = torch.randn(1, 16, 512)     # node features
coors = torch.randn(1, 16, 3)        # node coordinates
edges = torch.randn(1, 16, 16, 4)   # edge features (dense adjacency)

# Forward pass with edge features
feats_out, coors_out = layer(feats, coors, edges=edges)
# feats_out: (1, 16, 512), coors_out: (1, 16, 3)
```

--------------------------------

### Reconstruct and Visualize Protein Structure

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/examples/egnn_test.ipynb

Uses StructureBuilder to reconstruct a protein structure from coordinates and visualize it using 3Dmol.

```python
sb = sidechainnet.StructureBuilder( torch.tensor([AAS.index(x) for x in seq[:-padding_seq or None]]),
                                    crd=rearrange(input_rebuilt, 'l c d -> (l c) d') ) # pred_rebuilt
sb.to_3Dmol()
```

--------------------------------

### Initialize and Use EGNN Layer

Source: https://context7.com/lucidrains/egnn-pytorch/llms.txt

Create a single EGNN layer for equivariant message passing on node features and coordinates. Ensure correct dimensions for features, edges, and messages. This layer updates both features and coordinates by default.

```python
import torch
from egnn_pytorch import EGNN

# Create a single EGNN layer
layer = EGNN(
    dim=512,              # feature dimension
    edge_dim=0,           # edge feature dimension (0 = no edge features)
    m_dim=16,             # hidden message dimension
    fourier_features=0,   # fourier encoding of distances
    num_nearest_neighbors=0,  # k-nearest neighbors (0 = all nodes)
    dropout=0.0,          # dropout rate
    norm_feats=False,     # layer normalize features
    norm_coors=False,     # normalize coordinates
    update_feats=True,    # whether to update features
    update_coors=True,    # whether to update coordinates
    only_sparse_neighbors=False,  # use adjacency matrix for neighbors
    valid_radius=float('inf'),    # valid radius for neighbors
    m_pool_method='sum',  # pooling method ('sum' or 'mean')
    soft_edges=False,     # apply gating to edges
    coor_weights_clamp_value=None  # clamp coordinate updates
)

# Input: batch of 16 nodes with 512-dim features and 3D coordinates
feats = torch.randn(1, 16, 512)  # (batch, nodes, features)
coors = torch.randn(1, 16, 3)    # (batch, nodes, 3D coordinates)

# Forward pass - both features and coordinates are updated
feats_out, coors_out = layer(feats, coors)
# feats_out: (1, 16, 512), coors_out: (1, 16, 3)
```

--------------------------------

### Check NVIDIA GPU Status

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/examples/egnn_test.ipynb

Checks the status and details of the NVIDIA GPU available in the environment. Useful for confirming GPU availability and configuration.

```bash
!nvidia-smi
```

--------------------------------

### Rebuild 3D Structures for Visualization

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/examples/egnn_test.ipynb

Reconstructs 3D coordinate arrays for target, predicted, and input structures, applying masks to ensure correct placement. Used for visualization purposes.

```python
scaffs  = {"cloud_mask": scn_cloud_mask(seq[:-padding_seq or None]).bool()}
wrapper = torch.zeros(*scaffs["cloud_mask"].shape, 3).cpu()
# rebuild target
target_rebuilt = wrapper.clone()
target_rebuilt[scaffs["cloud_mask"].cpu()] = target_aligned.t().cpu()
# rebuild encoded-decoded
pre_target_rebuilt = wrapper.clone()
pre_target_rebuilt[scaffs["cloud_mask"].cpu()] = target_aligned.t().cpu()
# build input coords (w/out sidechain)
input_rebuilt = wrapper.clone()
input_rebuilt[scaffs["cloud_mask"].cpu()] = base_coords.cpu()
# build predicted
pred_rebuilt = wrapper.clone()
pred_rebuilt[scaffs["cloud_mask"].cpu()] = pred_aligned.t().cpu()
```

--------------------------------

### EGNN Network with Continuous Edges

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/README.md

Initialize an EGNN network to handle continuous edge features. The `edge_dim` parameter must be set to match the dimensionality of the `continuous_edges` tensor.

```python
import torch
from egnn_pytorch import EGNN_Network

net = EGNN_Network(
    num_tokens = 21,
    dim = 32,
    depth = 3,
    edge_dim = 4,
    num_nearest_neighbors = 3
)

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

continuous_edges = torch.randn(1, 1024, 1024, 4)

# naive adjacency matrix
# assuming the sequence is connected as a chain, with at most 2 neighbors - (1024, 1024)
i = torch.arange(1024)
adj_mat = (i[:, None] >= (i[None, :] - 1)) & (i[:, None] <= (i[None, :] + 1))

feats_out, coors_out = net(feats, coors, edges = continuous_edges, mask = mask, adj_mat = adj_mat) # (1, 1024, 32), (1, 1024, 3)
```

--------------------------------

### Generate Adjacency Matrix and Edge Attributes

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/examples/egnn_test.ipynb

Constructs an adjacency matrix based on bond data and converts it to sparse edge indices and attributes. Requires a pre-defined GVP_DATA dictionary and an nth_deg_adjacency helper function.

```python
adj_mat = torch.zeros(idxs.amax()+14, idxs.amax()+14)
    for i,idx in enumerate(idxs):
        # bond with next aa
        extra = []
        if i < idxs.shape[0]-1:
            extra = [[2, (idxs[i+1]-idx).item()]]

        bonds = idx + torch.tensor( GVP_DATA[seq[i]]['bonds'] + extra ).long().t() 
        adj_mat[bonds[0], bonds[1]] = 1.
    # convert to undirected
    adj_mat = adj_mat + adj_mat.t()
    # do N_th degree adjacency
    adj_mat, attr_mat = nth_deg_adjacency(adj_mat, n=adj_degree, sparse=True)

    edge_idxs = attr_mat.nonzero().t().long()
    edge_attrs = attr_mat[edge_idxs[0], edge_idxs[1]]
    return edge_idxs, edge_attrs
```

--------------------------------

### EGNN Network for Stable Training with High Neighbors

Source: https://context7.com/lucidrains/egnn-pytorch/llms.txt

Configure EGNN_Network for stable training when using a high number of neighbors by enabling coordinate normalization (`norm_coors=True`) and clamping coordinate update weights (`coor_weights_clamp_value`).

```python
import torch
from egnn_pytorch import EGNN_Network

# Create stable network for high neighbor count
net = EGNN_Network(
    num_tokens=21,
    dim=32,
    depth=3,
    num_nearest_neighbors=32,    # high neighbor count
    norm_coors=True,             # normalize relative coordinates
    coor_weights_clamp_value=2.  # clamp coordinate update weights
)

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

# Stable forward pass with many neighbors
feats_out, coors_out = net(feats, coors, mask=mask)
# feats_out: (1, 1024, 32), coors_out: (1, 1024, 3)
```

--------------------------------

### Cite E(n) Equivariant Graph Neural Networks

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/README.md

BibTeX citation for the EGNN paper.

```bibtex
@misc{satorras2021en,
    title 	= {E(n) Equivariant Graph Neural Networks}, 
    author 	= {Victor Garcia Satorras and Emiel Hoogeboom and Max Welling},
    year 	= {2021},
    eprint 	= {2102.09844},
    archivePrefix = {arXiv},
    primaryClass = {cs.LG}
}
```

--------------------------------

### Pack Scalars and Vectors for Bonds

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/examples/egnn_test.ipynb

Packs scalar and vector information for covalent bonds. This snippet is part of a larger process for encoding bond features, including norms and attributes.

```python
bond_buckets[native_idxs[0], native_idxs[1]] = native_attrs
    bond_attrs = bond_buckets[whole_bond_idxs[0] , whole_bond_idxs[1]]
    # pack scalars and vectors - extra token for covalent bonds
    bond_n_vectors = 1
    bond_n_scalars = (2 * len(needed_info["bond_scales"]) + 1) + 1 # last one is an embedd of size 1+len(cutoffs)
    whole_bond_enc = torch.cat([bond_vecs, # 1 vector - no need of reverse - we do 2x bonds (symmetry)
                                # scalars
                                bond_norms_enc, # 2 * len(scales) 
                                (bond_attrs-1).unsqueeze(-1) # 1 
                               ], dim=-1) 
    # free gpu mem
    if free_mem:
        del bond_buckets, bond_norms_enc, bond_vecs, dist_mat,
            close_bond_idxs, native_bond_idxs
        if closest: 
            del masked_dist_mat, sorted_col_idxs, sorted_row_idxs

    embedd_info = {"bond_n_vectors": bond_n_vectors, 
                   "bond_n_scalars": bond_n_scalars, 
                   "bond_embedding_nums": [ len(needed_info["cutoffs"]) + needed_info["adj_degree"] ]} # extra one for covalent (default)

    return whole_bond_idxs, whole_bond_enc, embedd_info
```

--------------------------------

### Calculate Mean Baseline Loss

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/examples/egnn_test.ipynb

Calculates the average baseline loss from a list of recorded baseline losses. This is typically done after training to establish a reference point.

```python
np.mean(baseline_losses)
```

--------------------------------

### Identify Protein Covalent Bonds

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/examples/egnn_test.ipynb

Returns the indices of covalent bonds for a protein sequence. Requires a protein sequence and optionally a cloud mask. Outputs edge indices and their attributes.

```python
def prot_covalent_bond(seq, adj_degree=1, cloud_mask=None):
    """ Returns the idxs of covalent bonds for a protein.
        Inputs 
        * seq: str. Protein sequence in 1-letter AA code.
        * cloud_mask: mask selecting the present atoms.
        Outputs: edge_idxs
    """
    # create or infer cloud_mask
    if cloud_mask is None: 
        cloud_mask = scn_cloud_mask(seq).bool()
    device, precise = cloud_mask.device, cloud_mask.type()
    # get starting poses for every aa
    scaff = torch.zeros_like(cloud_mask)
    scaff[:, 0] = 1
    idxs = scaff[cloud_mask].nonzero().view(-1)
    # get poses + idxs from the dict with GVP_DATA - return all edges
    adj_mat = torch.zeros(idxs.amax()+14, idxs.amax()+14)
    attr_mat = torch.zeros_like(adj_mat)
    for i,idx in enumerate(idxs):
        # bond with next aa
        extra = []
        if i < idxs.shape[0]-1:
            extra = [[2, (idxs[i+1]-idx).item()]]

        bonds = idx + torch.tensor( GVP_DATA[seq[i]]['bonds'] + extra ).long().t() 
        adj_mat[bonds[0], bonds[1]] = 1.

    # convert to undirected
    adj_mat = adj_mat + adj_mat.t()
    # do N_th degree adjacency
    for i in range(adj_degree):
        if i == 0:
            attr_mat += adj_mat
            continue

        adj_mat = (adj_mat @ adj_mat).bool().float()
        attr_mat[ (adj_mat - attr_mat.bool().float()).bool() ] += i+1

    edge_idxs = attr_mat.nonzero().t().long()
    edge_attrs = attr_mat[edge_idxs[0], edge_idxs[1]]

    return edge_idxs, edge_attrs
```

--------------------------------

### Define Data Constraints

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/examples/egnn_test.ipynb

Sets the length and count limits for processing protein data.

```python
# MIN_LEN = 70
MAX_LEN = 100000
MIN_LEN = 1000
MAX_PROTS = 360
```

--------------------------------

### EGNN Network with Continuous Edge Features

Source: https://context7.com/lucidrains/egnn-pytorch/llms.txt

Enable EGNN_Network to process continuous edge features by specifying `edge_dim`. This allows richer edge representations based on attributes like bond types or distances.

```python
import torch
from egnn_pytorch import EGNN_Network

# Create network with edge feature support
net = EGNN_Network(
    num_tokens=21,
    dim=32,
    depth=3,
    edge_dim=4,             # dimension of continuous edge features
    num_nearest_neighbors=3
)

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

# Continuous edge features (e.g., bond properties)
continuous_edges = torch.randn(1, 1024, 1024, 4)

# Optional adjacency matrix
i = torch.arange(1024)
adj_mat = (i[:, None] >= (i[None, :] - 1)) & (i[:, None] <= (i[None, :] + 1))

# Forward pass with edge features
feats_out, coors_out = net(feats, coors, edges=continuous_edges, mask=mask, adj_mat=adj_mat)
# feats_out: (1, 1024, 32), coors_out: (1, 1024, 3)
```

--------------------------------

### EGNN Layer Usage with Edges

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/README.md

Utilize EGNN layers with edge information. Ensure that the `edge_dim` parameter matches the dimensionality of the provided edge features.

```python
import torch
from egnn_pytorch import EGNN

layer1 = EGNN(dim = 512, edge_dim = 4)
layer2 = EGNN(dim = 512, edge_dim = 4)

feats = torch.randn(1, 16, 512)
coors = torch.randn(1, 16, 3)
edges = torch.randn(1, 16, 16, 4)

feats, coors = layer1(feats, coors, edges)
feats, coors = layer2(feats, coors, edges) # (1, 16, 512), (1, 16, 3)
```

--------------------------------

### EGNN Network with Sparse Neighbors

Source: https://github.com/lucidrains/egnn-pytorch/blob/main/README.md

Configure an EGNN network to only attend to sparse neighbors using an adjacency matrix. This is useful for graph structures where explicit neighbor connections are defined.

```python
import torch
from egnn_pytorch import EGNN_Network

net = EGNN_Network(
    num_tokens = 21,
    dim = 32,
    depth = 3,
    only_sparse_neighbors = True
)

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

# naive adjacency matrix
# assuming the sequence is connected as a chain, with at most 2 neighbors - (1024, 1024)
i = torch.arange(1024)
adj_mat = (i[:, None] >= (i[None, :] - 1)) & (i[:, None] <= (i[None, :] + 1))

feats_out, coors_out = net(feats, coors, mask = mask, adj_mat = adj_mat) # (1, 1024, 32), (1, 1024, 3)
```

--------------------------------

### Generate Z-axis Rotation Matrix

Source: https://context7.com/lucidrains/egnn-pytorch/llms.txt

Generates a 3D rotation matrix for rotation around the Z-axis using a specified angle in radians. Useful for data augmentation.

```python
import torch
from egnn_pytorch.utils import rot_z

# Single axis rotations
gamma = torch.tensor(0.5)  # radians
R_z = rot_z(gamma)  # rotation around z-axis
# R_z: (3, 3) rotation matrix
```