# proteindesignlab/slae architecture

> Architecture diagram of the GitHub repository proteindesignlab/slae, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-11-10). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/proteindesignlab/slae
- Repository: https://github.com/proteindesignlab/slae

## Overview

SLAE is a pretrained, all-atom protein-structure representation and reconstruction framework. Its principal workflow loads a structure into batched protein data, featurizes atoms and local graph edges, encodes the graph into atom/residue/graph embeddings, and optionally decodes residue embeddings into full-atom coordinates for PDB output. The graph also captures the dataset, tensor-conversion, and neural-network building-block layers. Some implementation details are unsampled; those relationships are shown only where the README or excerpts support them.

Main entry points: [`SLAE/inference.py`](https://github.com/proteindesignlab/slae/blob/main/SLAE/inference.py) (Inference API), [`SLAE/datasets/datamodule.py`](https://github.com/proteindesignlab/slae/blob/main/SLAE/datasets/datamodule.py) (Data module), [`SLAE/datasets/dataset.py`](https://github.com/proteindesignlab/slae/blob/main/SLAE/datasets/dataset.py) (Structure dataset), [`SLAE/datasets/dataloader.py`](https://github.com/proteindesignlab/slae/blob/main/SLAE/datasets/dataloader.py) (Protein loader)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_input["Structure input"]
  node_dataset[("Structure dataset<br/>[dataset.py]")]
  node_datamodule["Data module<br/>[datamodule.py]"]
  node_loader["Protein loader<br/>[dataloader.py]"]
end

subgraph group_representation["Graph representation"]
  node_atom_tensor["Atom tensors<br/>[atom_tensor.py]"]
  node_fa_repr["Atom representation"]
  node_featurizer["Graph featurizer"]
end

subgraph group_models["Autoencoder models"]
  node_encoder["Protein encoder<br/>[encoder.py]"]
  node_nn_blocks["Equivariant blocks"]
  node_decoder["All-atom decoder<br/>[decoder.py]"]
end

subgraph group_output["Inference output"]
  node_inference["Inference API<br/>[inference.py]"]
  node_coordinates["Reconstructed coordinates"]
  node_pdb_writer["PDB writer<br/>[write_pdb.py]"]
end

node_user(("Researcher"))

node_user -->|"requests encode"| node_inference
node_datamodule -->|"creates datasets"| node_dataset
node_datamodule -->|"provides batches"| node_loader
node_dataset -->|"converts atoms"| node_atom_tensor
node_loader -->|"batches data for"| node_featurizer
node_featurizer -->|"transforms representation"| node_fa_repr
node_fa_repr -->|"extracts atoms"| node_atom_tensor
node_inference -->|"runs encode"| node_encoder
node_encoder -->|"uses building blocks"| node_nn_blocks
node_inference -->|"runs decode"| node_decoder
node_decoder -->|"reconstructs coordinates"| node_coordinates
node_coordinates -->|"writes structure"| node_pdb_writer
node_user -->|"selects structure data"| node_datamodule
node_inference -->|"returns embeddings and coordinates"| node_coordinates

click node_dataset "https://github.com/proteindesignlab/slae/blob/main/SLAE/datasets/dataset.py"
click node_datamodule "https://github.com/proteindesignlab/slae/blob/main/SLAE/datasets/datamodule.py"
click node_loader "https://github.com/proteindesignlab/slae/blob/main/SLAE/datasets/dataloader.py"
click node_atom_tensor "https://github.com/proteindesignlab/slae/blob/main/SLAE/io/atom_tensor.py"
click node_fa_repr "https://github.com/proteindesignlab/slae/blob/main/SLAE/features/fa_representation.py"
click node_featurizer "https://github.com/proteindesignlab/slae/blob/main/SLAE/features/graph_featurizer.py"
click node_encoder "https://github.com/proteindesignlab/slae/blob/main/SLAE/model/encoder.py"
click node_nn_blocks "https://github.com/proteindesignlab/slae/tree/main/SLAE/nn"
click node_inference "https://github.com/proteindesignlab/slae/blob/main/SLAE/inference.py"
click node_decoder "https://github.com/proteindesignlab/slae/blob/main/SLAE/model/decoder.py"
click node_pdb_writer "https://github.com/proteindesignlab/slae/blob/main/SLAE/io/write_pdb.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_dataset,node_datamodule,node_loader toneBlue
class node_atom_tensor,node_fa_repr,node_featurizer toneAmber
class node_encoder,node_nn_blocks,node_decoder toneMint
class node_inference,node_coordinates,node_pdb_writer toneRose
class node_user toneIndigo
```

## Components

### Structure input

- [Structure dataset](https://github.com/proteindesignlab/slae/blob/main/SLAE/datasets/dataset.py): `SLAE/datasets/dataset.py`
- [Data module](https://github.com/proteindesignlab/slae/blob/main/SLAE/datasets/datamodule.py): `SLAE/datasets/datamodule.py`
- [Protein loader](https://github.com/proteindesignlab/slae/blob/main/SLAE/datasets/dataloader.py): `SLAE/datasets/dataloader.py`

### Graph representation

- [Atom tensors](https://github.com/proteindesignlab/slae/blob/main/SLAE/io/atom_tensor.py): `SLAE/io/atom_tensor.py`
- [Atom representation](https://github.com/proteindesignlab/slae/blob/main/SLAE/features/fa_representation.py): `SLAE/features/fa_representation.py`
- [Graph featurizer](https://github.com/proteindesignlab/slae/blob/main/SLAE/features/graph_featurizer.py): `SLAE/features/graph_featurizer.py`

### Autoencoder models

- [Protein encoder](https://github.com/proteindesignlab/slae/blob/main/SLAE/model/encoder.py): `SLAE/model/encoder.py`
- [Equivariant blocks](https://github.com/proteindesignlab/slae/tree/main/SLAE/nn): `SLAE/nn`
- [All-atom decoder](https://github.com/proteindesignlab/slae/blob/main/SLAE/model/decoder.py): `SLAE/model/decoder.py`

### Inference output

- [Inference API](https://github.com/proteindesignlab/slae/blob/main/SLAE/inference.py): `SLAE/inference.py`
- **Reconstructed coordinates**
- [PDB writer](https://github.com/proteindesignlab/slae/blob/main/SLAE/io/write_pdb.py): `SLAE/io/write_pdb.py`

### Other components

- **Researcher**

## Connections

- Researcher → Inference API: requests encode (evidence: `README.md`)
- Data module → Structure dataset: creates datasets (evidence: `SLAE/datasets/datamodule.py`)
- Data module → Protein loader: provides batches (evidence: `SLAE/datasets/datamodule.py`)
- Structure dataset → Atom tensors: converts atoms (evidence: `SLAE/datasets/dataset.py`)
- Protein loader → Graph featurizer: batches data for (evidence: `README.md`)
- Graph featurizer → Atom representation: transforms representation (evidence: `SLAE/features/graph_featurizer.py`)
- Atom representation → Atom tensors: extracts atoms (evidence: `SLAE/features/fa_representation.py`)
- Inference API → Protein encoder: runs encode (evidence: `README.md`)
- Protein encoder → Equivariant blocks: uses building blocks (evidence: `SLAE/model/encoder.py`)
- Inference API → All-atom decoder: runs decode (evidence: `README.md`)
- All-atom decoder → Reconstructed coordinates: reconstructs coordinates (evidence: `README.md`)
- Reconstructed coordinates → PDB writer: writes structure (evidence: `README.md`)
- Researcher → Data module: selects structure data (evidence: `README.md`)
- Inference API → Reconstructed coordinates: returns embeddings and coordinates (evidence: `README.md`)