### Install En-transformer

Source: https://github.com/lucidrains/en-transformer/blob/main/README.md

Install the En-transformer library using pip.

```bash
$ pip install En-transformer
```

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

### Install sidechainnet

Source: https://github.com/lucidrains/en-transformer/blob/main/README.md

Install the sidechainnet library, which is a dependency for running the protein backbone coordinate denoising toy task.

```bash
$ pip install sidechainnet
```

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

### Protein Backbone Denoising Example

Source: https://context7.com/lucidrains/en-transformer/llms.txt

A complete training example for denoising protein backbone coordinates using SidechainNet data. Demonstrates practical usage for structural biology applications.

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

torch.set_default_dtype(torch.float64)

# Initialize model
transformer = EnTransformer(
    num_tokens = 21,      # 20 amino acids + unknown
    dim = 32,
    dim_head = 64,
    heads = 4,
    depth = 4,
    rel_pos_emb = True,   # sequence has inherent order
    neighbors = 16
).cuda()

# Load protein data
data = scn.load(
    casp_version = 12,
    thinning = 30,
    with_pytorch = 'dataloaders',
    batch_size = 1,
    dynamic_batching = False
)


```

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

### Initialize EnTransformer with Features and Coordinates

Source: https://github.com/lucidrains/en-transformer/blob/main/README.md

Initialize the EnTransformer model with specified dimensions, depth, heads, and edge features. This configuration is suitable for general use with pre-computed features and coordinates.

```python
import torch
from en_transformer import EnTransformer

model = EnTransformer(
    dim = 512,
    depth = 4,                       # depth
    dim_head = 64,                   # dimension per head
    heads = 8,                       # number of heads
    edge_dim = 4,                    # dimension of edge feature
    neighbors = 64,                  # only do attention between coordinates N nearest neighbors - set to 0 to turn off
    talking_heads = True,            # use Shazeer's talking heads https://arxiv.org/abs/2003.02436
    checkpoint = True,               # use checkpointing so one can increase depth at little memory cost (and increase neighbors attended to)
    use_cross_product = True,        # use cross product vectors (idea by @MattMcPartlon)
    num_global_linear_attn_heads = 4 # if your number of neighbors above is low, you can assign a certain number of attention heads to weakly attend globally to all other nodes through linear attention (https://arxiv.org/abs/1812.01243)
)

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

mask = torch.ones(1, 1024).bool()

feats, coors = model(feats, coors, edges = edges, mask = mask)  # (1, 1024, 512), (1, 1024, 3)
```

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

### Verifying Equivariance Properties

Source: https://context7.com/lucidrains/en-transformer/llms.txt

Demonstrates how to verify E(n)-equivariance properties of the EnTransformer model. Ensures features remain invariant and coordinates transform equivariantly under rotation and translation.

```python
import torch
from en_transformer.utils import rot

torch.set_default_dtype(torch.float64)  # Use double precision for accuracy

model = EnTransformer(
    dim = 512,
    depth = 1,
    edge_dim = 4,
    rel_pos_emb = True
)

# Random rotation and translation
R = rot(*torch.rand(3))
T = torch.randn(1, 1, 3)

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

# Apply transformation before model
feats1, coors1 = model(feats, coors @ R + T, edges)

# Apply model then transform
feats2, coors2 = model(feats, coors, edges)

# Verify equivariance
assert torch.allclose(feats1, feats2, atol=1e-6), 'Features are invariant'
assert torch.allclose(coors1, coors2 @ R + T, atol=1e-6), 'Coordinates are equivariant'

print("Equivariance verified!")
```

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

### EnTransformer - Basic Usage with Continuous Features

Source: https://context7.com/lucidrains/en-transformer/llms.txt

Initialize and use the EnTransformer model with continuous node features and edge features. Ensure input tensors for features, coordinates, edges, and a padding mask are provided.

```python
import torch
from en_transformer import EnTransformer

# Basic usage with continuous features
model = EnTransformer(
    dim = 512,                        # feature dimension
    depth = 4,                        # number of transformer layers
    dim_head = 64,                    # dimension per attention head
    heads = 8,                        # number of attention heads
    edge_dim = 4,                     # edge feature dimension
    neighbors = 64,                   # attend to N nearest neighbors (0 = all)
    talking_heads = True,             # use talking heads attention
    checkpoint = True,                # gradient checkpointing for memory
    use_cross_product = True,         # cross product for coordinate updates
    num_global_linear_attn_heads = 4  # global linear attention heads
)

# Input tensors
feats = torch.randn(1, 1024, 512)     # (batch, nodes, dim)
coors = torch.randn(1, 1024, 3)       # (batch, nodes, 3D coordinates)
edges = torch.randn(1, 1024, 1024, 4) # (batch, nodes, nodes, edge_dim)
mask = torch.ones(1, 1024).bool()     # (batch, nodes) padding mask

# Forward pass
feats_out, coors_out = model(feats, coors, edges, mask=mask)
# Output: feats_out (1, 1024, 512), coors_out (1, 1024, 3)
```

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

### Initialize EnTransformer with Continuous Edges and Sparse Neighbors

Source: https://github.com/lucidrains/en-transformer/blob/main/README.md

Initialize the EnTransformer to handle continuous edge information while attending only to sparse neighbors. This configuration is useful when explicit edge features are provided.

```python
import torch
from en_transformer import EnTransformer
from en_transformer.utils import rot

model = EnTransformer(
    dim = 512,
    depth = 1,
    heads = 4,
    dim_head = 32,
    edge_dim = 4,
    num_nearest_neighbors = 0,
    only_sparse_neighbors = True
)

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

i = torch.arange(feats.shape[1])
adj_mat = (i[:, None] <= (i[None, :] + 1)) & (i[:, None] >= (i[None, :] - 1))

feats1, coors1 = model(feats, coors, adj_mat = adj_mat, edges = edges)
```

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

### Training Loop for En Transformer

Source: https://context7.com/lucidrains/en-transformer/llms.txt

This snippet demonstrates a typical training loop for the En Transformer model. It includes data loading, coordinate and sequence extraction, noise addition, denoising, loss computation, and optimizer steps. Ensure data is moved to CUDA and appropriate types are used.

```python
optim = Adam(transformer.parameters(), lr=1e-3)

for batch in data['train']:
    seqs = batch.seqs.cuda().argmax(dim=-1)
    coords = batch.crds.cuda().type(torch.float64)
    masks = batch.msks.cuda().bool()

    # Extract backbone atoms (N, CA, C)
    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)
    mask = repeat(masks, 'b n -> b (n c)', c=3)

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

    # Compute loss only on valid atoms
    loss = F.mse_loss(denoised_coords[mask], coords[mask])

    loss.backward()
    optim.step()
    optim.zero_grad()

    print(f'Loss: {loss.item():.4f}')
```

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

### Initialize EnTransformer with Token Embeddings

Source: https://github.com/lucidrains/en-transformer/blob/main/README.md

Initialize the EnTransformer model for cases where atomic and bond types are provided as tokens. This configuration automatically handles token embeddings.

```python
import torch
from en_transformer import EnTransformer

model = EnTransformer(
    num_tokens = 10,       # number of unique nodes, say atoms
    rel_pos_emb = True,    # set this to true if your sequence is not an unordered set. it will accelerate convergence
    num_edge_tokens = 5,   # number of unique edges, say bond types
    dim = 128,
    edge_dim = 16,
    depth = 3,
    heads = 4,
    dim_head = 32,
    neighbors = 8
)

atoms = torch.randint(0, 10, (1, 16))    # 10 different types of atoms
bonds = torch.randint(0, 5, (1, 16, 16)) # 5 different types of bonds (n x n)
coors = torch.randn(1, 16, 3)            # atomic spatial coordinates

feats_out, coors_out = model(atoms, coors, edges = bonds) # (1, 16, 512), (1, 16, 3)
```

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

### Initialize EnTransformer for Sparse Neighbors with Adjacency Matrix

Source: https://github.com/lucidrains/en-transformer/blob/main/README.md

Initialize the EnTransformer to attend only to sparse neighbors defined by an adjacency matrix. Requires setting `only_sparse_neighbors` to true and passing the adjacency matrix.

```python
import torch
from en_transformer import EnTransformer

model = EnTransformer(
    num_tokens = 10,
    dim = 512,
    depth = 1,
    heads = 4,
    dim_head = 32,
    neighbors = 0,
    only_sparse_neighbors = True,    # must be set to true
    num_adj_degrees = 2,             # the number of degrees to derive from 1st degree neighbors passed in
    adj_dim = 8                      # whether to pass the adjacency degree information as an edge embedding
)

atoms = torch.randint(0, 10, (1, 16))
coors = torch.randn(1, 16, 3)

# naively assume a single chain of atoms
i = torch.arange(atoms.shape[1])
adj_mat = (i[:, None] <= (i[None, :] + 1)) & (i[:, None] >= (i[None, :] - 1))

# adjacency matrix must be passed in
feats_out, coors_out = model(atoms, coors, adj_mat = adj_mat) # (1, 16, 512), (1, 16, 3)
```

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

### Rotation Utility Functions

Source: https://context7.com/lucidrains/en-transformer/llms.txt

Provides SO(3) rotation matrix functions for testing equivariance. Generates rotation matrices using Euler angles. Use for applying rotations to coordinates.

```python
import torch
from en_transformer.utils import rot, rot_y, rot_z

# Generate rotation matrices
alpha, beta, gamma = torch.rand(3)

# Rotation around Z-axis
Rz = rot_z(gamma)  # (3, 3) rotation matrix

# Rotation around Y-axis
Ry = rot_y(beta)   # (3, 3) rotation matrix

# Combined rotation (Z-Y-Z Euler angles)
R = rot(alpha, beta, gamma)  # (3, 3) rotation matrix

# Apply rotation to coordinates
coors = torch.randn(1, 16, 3)
rotated_coors = coors @ R  # (1, 16, 3) rotated coordinates
```

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

### EnTransformer Model Initialization and Forward Pass

Source: https://context7.com/lucidrains/en-transformer/llms.txt

The EnTransformer class is the primary interface for building E(n)-equivariant models. It processes node features and 3D coordinates while maintaining equivariance.

```APIDOC
## EnTransformer Class

### Description
Initializes and executes the EnTransformer model for processing node features and 3D coordinates.

### Parameters
- **dim** (int) - Feature dimension
- **depth** (int) - Number of transformer layers
- **dim_head** (int) - Dimension per attention head
- **heads** (int) - Number of attention heads
- **edge_dim** (int) - Edge feature dimension
- **neighbors** (int) - Number of nearest neighbors to attend to (0 = all)
- **talking_heads** (bool) - Whether to use talking heads attention
- **checkpoint** (bool) - Whether to use gradient checkpointing
- **use_cross_product** (bool) - Whether to use cross product for coordinate updates
- **num_global_linear_attn_heads** (int) - Number of global linear attention heads

### Request Example
model = EnTransformer(dim=512, depth=4, heads=8)
feats_out, coors_out = model(feats, coors, edges, mask=mask)

### Response
- **feats_out** (Tensor) - Updated node features
- **coors_out** (Tensor) - Refined 3D coordinates
```

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

### EnTransformer - Usage with Token Embeddings

Source: https://context7.com/lucidrains/en-transformer/llms.txt

Configure EnTransformer to handle discrete inputs like atom and bond types by setting `num_tokens` and `num_edge_tokens`. The model will manage embeddings internally, accepting integer tensors for nodes and edges.

```python
import torch
from en_transformer import EnTransformer

# Model with automatic embeddings for atoms and bonds
model = EnTransformer(
    num_tokens = 10,       # number of unique node types (e.g., atom types)
    num_edge_tokens = 5,   # number of unique edge types (e.g., bond types)
    rel_pos_emb = True,    # relative position embedding for ordered sequences
    dim = 128,
    edge_dim = 16,
    depth = 3,
    heads = 4,
    dim_head = 32,
    neighbors = 8
)

# Integer inputs for atoms and bonds
atoms = torch.randint(0, 10, (1, 16))     # (batch, nodes) - atom type indices
bonds = torch.randint(0, 5, (1, 16, 16))  # (batch, nodes, nodes) - bond types
coors = torch.randn(1, 16, 3)             # (batch, nodes, 3D coordinates)

# Forward pass with token inputs
feats_out, coors_out = model(atoms, coors, edges=bonds)
# Output: feats_out (1, 16, 128), coors_out (1, 16, 3)
```

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

### EnTransformer with Token Embeddings

Source: https://context7.com/lucidrains/en-transformer/llms.txt

Configures the model to handle discrete inputs like atom types and bond types using internal embeddings.

```APIDOC
## EnTransformer Token Embeddings

### Description
Enables the model to process integer tensor inputs for nodes and edges by specifying token counts.

### Parameters
- **num_tokens** (int) - Number of unique node types
- **num_edge_tokens** (int) - Number of unique edge types
- **rel_pos_emb** (bool) - Whether to use relative position embeddings

### Request Example
model = EnTransformer(num_tokens=10, num_edge_tokens=5, dim=128)
feats_out, coors_out = model(atoms, coors, edges=bonds)
```

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

### EquivariantAttention Module

Source: https://context7.com/lucidrains/en-transformer/llms.txt

The core E(n)-equivariant attention mechanism. Computes attention weights based on distances and features, then updates node features (invariant) and coordinates (equivariant). Use for standalone attention layers.

```python
import torch
from en_transformer.en_transformer import EquivariantAttention

# Standalone attention layer
attention = EquivariantAttention(
    dim = 256,
    dim_head = 64,
    heads = 4,
    edge_dim = 8,
    coors_hidden_dim = 16,
    neighbors = 16,
    only_sparse_neighbors = False,
    valid_neighbor_radius = float('inf'),
    rel_pos_emb = True,
    norm_rel_coors = True,
    use_cross_product = True,
    talking_heads = True,
    dropout = 0.1,
    num_global_linear_attn_heads = 2
)

feats = torch.randn(1, 32, 256)
coors = torch.randn(1, 32, 3)
edges = torch.randn(1, 32, 32, 8)
mask = torch.ones(1, 32).bool()

# Forward pass through attention
feats_delta, coors_delta = attention(feats, coors, edges=edges, mask=mask)
# Output: feature updates and coordinate updates
```

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

### EnTransformer with Sparse Adjacency Matrix

Source: https://context7.com/lucidrains/en-transformer/llms.txt

Configures the model to use sparse neighbor attention based on a provided adjacency matrix.

```APIDOC
## EnTransformer Sparse Attention

### Description
Restricts attention to neighbors defined by an adjacency matrix, useful for molecular graphs.

### Parameters
- **only_sparse_neighbors** (bool) - If True, only attend to adjacency-defined neighbors
- **num_adj_degrees** (int) - Number of degrees to include in neighbor search
- **adj_dim** (int) - Adjacency degree embedding dimension

### Request Example
model = EnTransformer(only_sparse_neighbors=True, num_adj_degrees=2)
feats_out, coors_out = model(atoms, coors, adj_mat=adj_mat)
```

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

### EnTransformer - Sparse Adjacency Matrix Attention

Source: https://context7.com/lucidrains/en-transformer/llms.txt

Utilize sparse neighbor attention by setting `only_sparse_neighbors=True` and providing an adjacency matrix (`adj_mat`). This is suitable for molecular graphs where attention is restricted to connected nodes, with optional multi-hop neighbor inclusion via `num_adj_degrees`.

```python
import torch
from en_transformer import EnTransformer

# Model with sparse neighbor attention
model = EnTransformer(
    num_tokens = 10,
    dim = 512,
    depth = 1,
    heads = 4,
    dim_head = 32,
    neighbors = 0,
    only_sparse_neighbors = True,  # only attend to adjacency-defined neighbors
    num_adj_degrees = 2,           # include 1st and 2nd degree neighbors
    adj_dim = 8                    # adjacency degree embedding dimension
)

atoms = torch.randint(0, 10, (1, 16))
coors = torch.randn(1, 16, 3)

# Create adjacency matrix (chain of atoms example)
n = atoms.shape[1]
i = torch.arange(n)
adj_mat = (i[:, None] <= (i[None, :] + 1)) & (i[:, None] >= (i[None, :] - 1))

# Forward pass with adjacency matrix
feats_out, coors_out = model(atoms, coors, adj_mat=adj_mat)
```

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

### EnTransformer with Continuous Edges and Sparse Adjacency

Source: https://context7.com/lucidrains/en-transformer/llms.txt

Model combining continuous edge features with sparse adjacency attention. Use when arbitrary edge information needs to be passed while restricting attention to specific node pairs.

```python
import torch
from en_transformer import EnTransformer

# Model combining continuous edges with sparse attention
model = EnTransformer(
    dim = 512,
    depth = 1,
    heads = 4,
    dim_head = 32,
    edge_dim = 4,
    num_nearest_neighbors = 0,
    only_sparse_neighbors = True
)

feats = torch.randn(1, 16, 512)
coors = torch.randn(1, 16, 3)
edges = torch.randn(1, 16, 16, 4)  # continuous edge features

# Create chain adjacency matrix
n = feats.shape[1]
i = torch.arange(n)
adj_mat = (i[:, None] <= (i[None, :] + 1)) & (i[:, None] >= (i[None, :] - 1))

# Forward pass with edges and adjacency
feats_out, coors_out = model(feats, coors, adj_mat=adj_mat, edges=edges)
# Output: feats_out (1, 16, 512), coors_out (1, 16, 3)
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.