atlasfold 1.0.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- atlasfold-1.0.0/LICENSE +21 -0
- atlasfold-1.0.0/PKG-INFO +253 -0
- atlasfold-1.0.0/README.md +217 -0
- atlasfold-1.0.0/pyproject.toml +51 -0
- atlasfold-1.0.0/setup.cfg +4 -0
- atlasfold-1.0.0/src/atlasfold/__init__.py +0 -0
- atlasfold-1.0.0/src/atlasfold/cli/__init__.py +67 -0
- atlasfold-1.0.0/src/atlasfold/cli/__main__.py +6 -0
- atlasfold-1.0.0/src/atlasfold/cli/monomer.py +496 -0
- atlasfold-1.0.0/src/atlasfold/cli/multigpu.py +125 -0
- atlasfold-1.0.0/src/atlasfold/cli/multimer.py +513 -0
- atlasfold-1.0.0/src/atlasfold/common/__init__.py +0 -0
- atlasfold-1.0.0/src/atlasfold/common/ccd.py +287 -0
- atlasfold-1.0.0/src/atlasfold/common/featurize.py +308 -0
- atlasfold-1.0.0/src/atlasfold/common/file_io.py +327 -0
- atlasfold-1.0.0/src/atlasfold/common/metadata.py +232 -0
- atlasfold-1.0.0/src/atlasfold/common/process.py +70 -0
- atlasfold-1.0.0/src/atlasfold/common/protein.py +307 -0
- atlasfold-1.0.0/src/atlasfold/common/residue_constants.py +383 -0
- atlasfold-1.0.0/src/atlasfold/common/templates.py +286 -0
- atlasfold-1.0.0/src/atlasfold/configs/__init__.py +0 -0
- atlasfold-1.0.0/src/atlasfold/configs/atlasfold.py +43 -0
- atlasfold-1.0.0/src/atlasfold/configs/atlasfold_multimer.py +54 -0
- atlasfold-1.0.0/src/atlasfold/data/__init__.py +0 -0
- atlasfold-1.0.0/src/atlasfold/data/cif_factory.py +244 -0
- atlasfold-1.0.0/src/atlasfold/data/fasta.py +53 -0
- atlasfold-1.0.0/src/atlasfold/data/mmseq2.py +174 -0
- atlasfold-1.0.0/src/atlasfold/model/__init__.py +3 -0
- atlasfold-1.0.0/src/atlasfold/model/model.py +558 -0
- atlasfold-1.0.0/src/atlasfold/model/model_multimer.py +654 -0
- atlasfold-1.0.0/src/atlasfold/model/network/__init__.py +0 -0
- atlasfold-1.0.0/src/atlasfold/model/network/atom_attention.py +211 -0
- atlasfold-1.0.0/src/atlasfold/model/network/attention.py +294 -0
- atlasfold-1.0.0/src/atlasfold/model/network/block.py +293 -0
- atlasfold-1.0.0/src/atlasfold/model/network/confidence_head.py +515 -0
- atlasfold-1.0.0/src/atlasfold/model/network/diffusion_head.py +450 -0
- atlasfold-1.0.0/src/atlasfold/model/network/diffusion_transformer.py +513 -0
- atlasfold-1.0.0/src/atlasfold/model/network/distogram_head.py +50 -0
- atlasfold-1.0.0/src/atlasfold/model/network/misc.py +138 -0
- atlasfold-1.0.0/src/atlasfold/model/network/primitives/__init__.py +26 -0
- atlasfold-1.0.0/src/atlasfold/model/network/primitives/activation.py +15 -0
- atlasfold-1.0.0/src/atlasfold/model/network/primitives/dropout.py +45 -0
- atlasfold-1.0.0/src/atlasfold/model/network/primitives/initialize.py +82 -0
- atlasfold-1.0.0/src/atlasfold/model/network/primitives/linear.py +110 -0
- atlasfold-1.0.0/src/atlasfold/model/network/primitives/normalization.py +77 -0
- atlasfold-1.0.0/src/atlasfold/model/network/primitives/transition.py +28 -0
- atlasfold-1.0.0/src/atlasfold/model/network/primitives/triangle_update.py +515 -0
- atlasfold-1.0.0/src/atlasfold/model/network/rel_pos_encoding.py +141 -0
- atlasfold-1.0.0/src/atlasfold/model/network/template.py +239 -0
- atlasfold-1.0.0/src/atlasfold/model/network/trunk.py +168 -0
- atlasfold-1.0.0/src/atlasfold/model/utils/__init__.py +0 -0
- atlasfold-1.0.0/src/atlasfold/model/utils/confidence_metrics.py +392 -0
- atlasfold-1.0.0/src/atlasfold/pretrained.py +156 -0
- atlasfold-1.0.0/src/atlasfold/runner.py +484 -0
- atlasfold-1.0.0/src/atlasfold/runner_multimer.py +579 -0
- atlasfold-1.0.0/src/atlasfold/train/__init__.py +0 -0
- atlasfold-1.0.0/src/atlasfold/train/config.py +74 -0
- atlasfold-1.0.0/src/atlasfold/train/losses/__init__.py +1 -0
- atlasfold-1.0.0/src/atlasfold/train/losses/confidence.py +386 -0
- atlasfold-1.0.0/src/atlasfold/train/losses/diffusion.py +206 -0
- atlasfold-1.0.0/src/atlasfold/train/losses/distogram.py +62 -0
- atlasfold-1.0.0/src/atlasfold/train/monomer/__init__.py +0 -0
- atlasfold-1.0.0/src/atlasfold/train/monomer/cropper.py +179 -0
- atlasfold-1.0.0/src/atlasfold/train/monomer/datamodule.py +149 -0
- atlasfold-1.0.0/src/atlasfold/train/monomer/dataset.py +552 -0
- atlasfold-1.0.0/src/atlasfold/train/monomer/model_train.py +363 -0
- atlasfold-1.0.0/src/atlasfold/train/monomer/structure_alignment.py +106 -0
- atlasfold-1.0.0/src/atlasfold/train/monomer/train_module.py +681 -0
- atlasfold-1.0.0/src/atlasfold/train/monomer/validation_metrics.py +74 -0
- atlasfold-1.0.0/src/atlasfold/train/multimer/__init__.py +0 -0
- atlasfold-1.0.0/src/atlasfold/train/multimer/cropper.py +341 -0
- atlasfold-1.0.0/src/atlasfold/train/multimer/datamodule.py +168 -0
- atlasfold-1.0.0/src/atlasfold/train/multimer/dataset.py +1135 -0
- atlasfold-1.0.0/src/atlasfold/train/multimer/model_train.py +369 -0
- atlasfold-1.0.0/src/atlasfold/train/multimer/train_alignment.py +472 -0
- atlasfold-1.0.0/src/atlasfold/train/multimer/train_module.py +460 -0
- atlasfold-1.0.0/src/atlasfold/train/multimer/validation_metrics.py +432 -0
- atlasfold-1.0.0/src/atlasfold/train/utils/__init__.py +0 -0
- atlasfold-1.0.0/src/atlasfold/train/utils/dl_sampler.py +106 -0
- atlasfold-1.0.0/src/atlasfold/train/utils/ema.py +111 -0
- atlasfold-1.0.0/src/atlasfold/train/utils/gradient_logging.py +25 -0
- atlasfold-1.0.0/src/atlasfold/train/utils/kernels/__init__.py +0 -0
- atlasfold-1.0.0/src/atlasfold/train/utils/kernels/cdist.py +328 -0
- atlasfold-1.0.0/src/atlasfold/train/utils/lr_scheduler.py +54 -0
- atlasfold-1.0.0/src/atlasfold/utils/__init__.py +0 -0
- atlasfold-1.0.0/src/atlasfold/utils/checkpointing.py +99 -0
- atlasfold-1.0.0/src/atlasfold/utils/geometry/__init__.py +0 -0
- atlasfold-1.0.0/src/atlasfold/utils/geometry/metrics.py +184 -0
- atlasfold-1.0.0/src/atlasfold/utils/geometry/random_augment.py +271 -0
- atlasfold-1.0.0/src/atlasfold/utils/geometry/rigid_align.py +525 -0
- atlasfold-1.0.0/src/atlasfold/utils/kernel.py +36 -0
- atlasfold-1.0.0/src/atlasfold/utils/misc.py +82 -0
- atlasfold-1.0.0/src/atlasfold/utils/rasa.py +236 -0
- atlasfold-1.0.0/src/atlasfold/utils/torch_utils.py +162 -0
- atlasfold-1.0.0/src/atlasfold.egg-info/PKG-INFO +253 -0
- atlasfold-1.0.0/src/atlasfold.egg-info/SOURCES.txt +107 -0
- atlasfold-1.0.0/src/atlasfold.egg-info/dependency_links.txt +1 -0
- atlasfold-1.0.0/src/atlasfold.egg-info/entry_points.txt +2 -0
- atlasfold-1.0.0/src/atlasfold.egg-info/requires.txt +23 -0
- atlasfold-1.0.0/src/atlasfold.egg-info/top_level.txt +2 -0
- atlasfold-1.0.0/src/atlaslm/__init__.py +2 -0
- atlasfold-1.0.0/src/atlaslm/alphabet.py +37 -0
- atlasfold-1.0.0/src/atlaslm/layers/__init__.py +2 -0
- atlasfold-1.0.0/src/atlaslm/layers/attention.py +113 -0
- atlasfold-1.0.0/src/atlaslm/layers/regression_head.py +17 -0
- atlasfold-1.0.0/src/atlaslm/layers/rotary.py +104 -0
- atlasfold-1.0.0/src/atlaslm/layers/transformer_stack.py +154 -0
- atlasfold-1.0.0/src/atlaslm/model.py +285 -0
- atlasfold-1.0.0/src/atlaslm/pretrained.py +152 -0
atlasfold-1.0.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Seonghwan Seo
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
atlasfold-1.0.0/PKG-INFO
ADDED
|
@@ -0,0 +1,253 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: atlasfold
|
|
3
|
+
Version: 1.0.0
|
|
4
|
+
Summary: Protein language model-based protein folding and co-folding
|
|
5
|
+
Author-email: SeonghwanSeo <shwan0106@gmail.com>
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/SeonghwanSeo/atlasfold
|
|
8
|
+
Project-URL: Documentation, https://github.com/SeonghwanSeo/atlasfold#readme
|
|
9
|
+
Project-URL: Repository, https://github.com/SeonghwanSeo/atlasfold
|
|
10
|
+
Project-URL: Issues, https://github.com/SeonghwanSeo/atlasfold/issues
|
|
11
|
+
Keywords: protein,protein-language-model,protein-folding,cofolding
|
|
12
|
+
Requires-Python: >=3.10
|
|
13
|
+
Description-Content-Type: text/markdown
|
|
14
|
+
License-File: LICENSE
|
|
15
|
+
Requires-Dist: numpy
|
|
16
|
+
Requires-Dist: torch
|
|
17
|
+
Requires-Dist: huggingface_hub
|
|
18
|
+
Provides-Extra: fold
|
|
19
|
+
Requires-Dist: scipy; extra == "fold"
|
|
20
|
+
Requires-Dist: einops; extra == "fold"
|
|
21
|
+
Requires-Dist: gemmi; extra == "fold"
|
|
22
|
+
Requires-Dist: omegaconf; extra == "fold"
|
|
23
|
+
Requires-Dist: numba; extra == "fold"
|
|
24
|
+
Provides-Extra: cuequiv
|
|
25
|
+
Requires-Dist: cuequivariance_torch>=0.6.0; extra == "cuequiv"
|
|
26
|
+
Requires-Dist: cuequivariance_ops_cu12>=0.6.0; extra == "cuequiv"
|
|
27
|
+
Requires-Dist: cuequivariance_ops_torch_cu12>=0.6.0; extra == "cuequiv"
|
|
28
|
+
Provides-Extra: train
|
|
29
|
+
Requires-Dist: msgpack; extra == "train"
|
|
30
|
+
Requires-Dist: lightning; extra == "train"
|
|
31
|
+
Requires-Dist: wandb; extra == "train"
|
|
32
|
+
Requires-Dist: lmdb; extra == "train"
|
|
33
|
+
Requires-Dist: pandas; extra == "train"
|
|
34
|
+
Requires-Dist: zstandard; extra == "train"
|
|
35
|
+
Dynamic: license-file
|
|
36
|
+
|
|
37
|
+
# AtlasFold
|
|
38
|
+
|
|
39
|
+
[\[Paper\]](https://www.biorxiv.org/content/10.64898/2026.09.04.749352v2), [\[PDF\]](docs/atlasfold.pdf)
|
|
40
|
+
|
|
41
|
+
The **Atlas family** is a collection of open protein models for sequence representation and structure prediction.
|
|
42
|
+
AtlasLM is a protein language model (PLM), while AtlasFold and AtlasFold-M are trainable PLM-based models for protein folding and co-folding, respectively.
|
|
43
|
+
|
|
44
|
+

|
|
45
|
+
|
|
46
|
+
AtlasFold achieves state-of-the-art accuracy among protein language model-based folding methods. AtlasFold and AtlasFold-M predict structures without an MSA search.
|
|
47
|
+
This repository provides pretrained models, training and inference code, staged training configurations, and preprocessing workflows for monomer and multimer folding.
|
|
48
|
+
|
|
49
|
+
## Installation
|
|
50
|
+
|
|
51
|
+
AtlasFold requires Python 3.10 or later. A CUDA GPU is recommended for structure prediction.
|
|
52
|
+
|
|
53
|
+
Install the inference dependencies from PyPI:
|
|
54
|
+
|
|
55
|
+
```bash
|
|
56
|
+
pip install "atlasfold[fold,cuequiv]"
|
|
57
|
+
```
|
|
58
|
+
|
|
59
|
+
To use AtlasLM as a standalone protein language model:
|
|
60
|
+
|
|
61
|
+
```bash
|
|
62
|
+
pip install atlasfold
|
|
63
|
+
```
|
|
64
|
+
|
|
65
|
+
To install the latest development version from GitHub:
|
|
66
|
+
|
|
67
|
+
```bash
|
|
68
|
+
git clone https://github.com/SeonghwanSeo/atlasfold.git
|
|
69
|
+
cd atlasfold
|
|
70
|
+
pip install -e ".[fold,cuequiv]"
|
|
71
|
+
```
|
|
72
|
+
|
|
73
|
+
AtlasFold supports [cuEquivariance](https://docs.nvidia.com/cuda/cuequivariance/) kernels for faster inference. For systems without compatible NVIDIA CUDA hardware, install `atlasfold[fold]` instead.
|
|
74
|
+
|
|
75
|
+
## Running your first prediction
|
|
76
|
+
|
|
77
|
+
For a monomer, save one protein sequence per FASTA record:
|
|
78
|
+
|
|
79
|
+
```text
|
|
80
|
+
>protein_a
|
|
81
|
+
MKTAYIAKQRQISFVKSHFSRQDILDLWIYHTQGYFPD
|
|
82
|
+
>protein_b
|
|
83
|
+
FNPVGVAFKGNNGKYLSRIHRSGIDYTEFAKDNTD
|
|
84
|
+
```
|
|
85
|
+
|
|
86
|
+
Then run:
|
|
87
|
+
|
|
88
|
+
```bash
|
|
89
|
+
atlasfold monomer --input-fasta monomer.fasta --out-dir predictions/monomer/
|
|
90
|
+
```
|
|
91
|
+
|
|
92
|
+
For a protein complex, use one FASTA record per complex and separate chains with `:`:
|
|
93
|
+
|
|
94
|
+
```text
|
|
95
|
+
>complex_a
|
|
96
|
+
MKTAYIAKQRQISFVKSHFS:GGHVDHGKSTTTGHLIYK
|
|
97
|
+
>complex_b
|
|
98
|
+
MKEGFYWIQHNGRVQVAYYTHGVTEDLETGQTIIGVWHLTQGDDICHNGEAEILAGPLEPPI:MKEGFYWIQHNGRVQVA
|
|
99
|
+
YYTHGVTEDLETGQTIIGVWHLTQGDDICHNGEAEILAGPLEPPI
|
|
100
|
+
```
|
|
101
|
+
|
|
102
|
+
```bash
|
|
103
|
+
atlasfold multimer --input-fasta multimer.fasta --out-dir predictions/multimer/
|
|
104
|
+
```
|
|
105
|
+
|
|
106
|
+
Template-assisted inference for AtlasFold-M is not supported by the current runner or CLI.
|
|
107
|
+
|
|
108
|
+
The repository entry point provides the same interface:
|
|
109
|
+
|
|
110
|
+
```bash
|
|
111
|
+
python run_atlasfold.py monomer --input-fasta monomer.fasta --out-dir predictions/monomer/
|
|
112
|
+
python run_atlasfold.py multimer --input-fasta multimer.fasta --out-dir predictions/multimer/
|
|
113
|
+
```
|
|
114
|
+
|
|
115
|
+
Both AtlasFold and AtlasFold-M support batched inference with multiple FASTA records, enabling high-throughput structure prediction.
|
|
116
|
+
|
|
117
|
+
Run `atlasfold monomer --help` or `atlasfold multimer --help` for all options, and see the [inference guide](docs/inference.md) for batching, sampling, confidence values, and output formats.
|
|
118
|
+
|
|
119
|
+
## Performance and GPU memory
|
|
120
|
+
|
|
121
|
+
Peak GPU memory grows with total residue length when generating five diffusion samples.
|
|
122
|
+
|
|
123
|
+
| Total residues | AtlasFold | AtlasFold-M |
|
|
124
|
+
| ---: | ---: | ---: |
|
|
125
|
+
| 256 | 7.22 GiB | 7.31 GiB |
|
|
126
|
+
| 512 | 8.80 GiB | 9.07 GiB |
|
|
127
|
+
| 1,024 | 15.02 GiB | 16.05 GiB |
|
|
128
|
+
| 1,536 | 25.36 GiB | 27.63 GiB |
|
|
129
|
+
| 2,048 | 39.81 GiB | 43.83 GiB |
|
|
130
|
+
|
|
131
|
+
For multi-target workloads, batched inference substantially increases throughput. With five diffusion samples and up to 4,096 residues processed per batch:
|
|
132
|
+
|
|
133
|
+
| Workload | Unbatched | Batched | Throughput gain |
|
|
134
|
+
| --- | ---: | ---: | ---: |
|
|
135
|
+
| AtlasFold, 64-residue monomers | 0.775 sequences/s | 9.663 sequences/s | 12.5× |
|
|
136
|
+
| AtlasFold-M, 256-residue complexes | 0.149 complexes/s | 0.314 complexes/s | 2.1× |
|
|
137
|
+
|
|
138
|
+
Measurements were collected on a single NVIDIA B200 using PyTorch 2.10.0, CUDA 12.8, and cuEquivariance 0.10.0. See the [performance guide](docs/performance.md) for the complete memory and runtime measurements.
|
|
139
|
+
|
|
140
|
+
## Python API
|
|
141
|
+
|
|
142
|
+
### Structure prediction
|
|
143
|
+
|
|
144
|
+
```python
|
|
145
|
+
from atlasfold.pretrained import get_runner, load_model
|
|
146
|
+
|
|
147
|
+
model = load_model("atlasfold", device="cuda")
|
|
148
|
+
runner = get_runner(model)
|
|
149
|
+
result = runner.fold(
|
|
150
|
+
"protein_a",
|
|
151
|
+
"MKTAYIAKQRQISFVKSHFSRQDILDLWIYHTQGYFPD",
|
|
152
|
+
num_samples=5,
|
|
153
|
+
)
|
|
154
|
+
|
|
155
|
+
print(result.best.avg_plddt, result.best.ptm)
|
|
156
|
+
with open("protein_a.pdb", "w") as handle:
|
|
157
|
+
handle.write(result.best.to_pdb())
|
|
158
|
+
with open("protein_a.cif", "w") as handle:
|
|
159
|
+
handle.write(result.best.to_mmcif())
|
|
160
|
+
```
|
|
161
|
+
|
|
162
|
+
Use `load_model("atlasfold-m", device="cuda")` for a complex and pass a list of chain sequences to `runner.fold()`. The Python API defaults to CPU if `device` is omitted; the CLI selects CUDA automatically when available.
|
|
163
|
+
|
|
164
|
+
### Protein language model
|
|
165
|
+
|
|
166
|
+
```python
|
|
167
|
+
import torch
|
|
168
|
+
|
|
169
|
+
from atlaslm import load_model
|
|
170
|
+
|
|
171
|
+
model = load_model("atlaslm-3b", device="cuda", dtype=torch.bfloat16)
|
|
172
|
+
output = model.embed_sequences(
|
|
173
|
+
["MKTAYIAKQRQISFVKSHFSRQDILDLWIYHTQGYFPD"],
|
|
174
|
+
return_hidden_states=True,
|
|
175
|
+
)
|
|
176
|
+
|
|
177
|
+
print(output.embeddings.shape)
|
|
178
|
+
print(len(output.hidden_states))
|
|
179
|
+
```
|
|
180
|
+
|
|
181
|
+
Pass `return_attentions=True` to return attention maps. Attention tensors grow quadratically with sequence length and can require substantially more memory.
|
|
182
|
+
|
|
183
|
+
## Available models
|
|
184
|
+
|
|
185
|
+
| Model | Parameters | Weight download | Use | Weights |
|
|
186
|
+
| --- | ---: | ---: | --- | --- |
|
|
187
|
+
| AtlasLM-600M | 575M | 1.07 GiB | Protein sequence representations | [`SeonghwanSeo/atlaslm-600m-base`](https://huggingface.co/SeonghwanSeo/atlaslm-600m-base) |
|
|
188
|
+
| AtlasLM-3B | 3.06B | 5.71 GiB | Protein sequence representations and AtlasFold backbone | [`SeonghwanSeo/atlaslm-3b-base`](https://huggingface.co/SeonghwanSeo/atlaslm-3b-base) |
|
|
189
|
+
| AtlasFold | 215M + AtlasLM-3B | 0.80 GiB + AtlasLM-3B | Monomer structure prediction | [`SeonghwanSeo/atlasfold-260703`](https://huggingface.co/SeonghwanSeo/atlasfold-260703) |
|
|
190
|
+
| AtlasFold-M | 220M + AtlasLM-3B | 0.82 GiB + AtlasLM-3B | Protein-complex structure prediction | [`SeonghwanSeo/atlasfold-m-260725`](https://huggingface.co/SeonghwanSeo/atlasfold-m-260725) |
|
|
191
|
+
|
|
192
|
+
AtlasLM-600M is deprecated now that ESMC-600M is available for commercial use. AtlasLM-3B is the recommended AtlasLM checkpoint.
|
|
193
|
+
|
|
194
|
+
## Evaluation
|
|
195
|
+
|
|
196
|
+
Evaluation protocols and results for AtlasFold and AtlasFold-M are provided in the [benchmark documentation](docs/benchmarks.md). The associated prediction structures and evaluation artifacts are available from the release folder below.
|
|
197
|
+
|
|
198
|
+
## Checkpoints, data, and benchmark artifacts
|
|
199
|
+
|
|
200
|
+
Large release artifacts are hosted in the [AtlasFold Google Drive folder](https://drive.google.com/drive/folders/1xjSBmbCFqghWj8xKYuEIDISKoQkCX45I?usp=sharing):
|
|
201
|
+
|
|
202
|
+
| Artifact | Contents |
|
|
203
|
+
| --- | --- |
|
|
204
|
+
| [Intermediate checkpoints](https://drive.google.com/drive/folders/1TDz2Ng4-zYfpTxlqg5wD3MkzwJyL2gTm) | AtlasLM pretraining checkpoints and AtlasFold/AtlasFold-M staged training checkpoints |
|
|
205
|
+
| [Structural datasets](https://drive.google.com/drive/folders/1EiRTKSUL3iD_MQ_0qmj5Sb-2KMh-3kmS) | Processed monomer and multimer training and validation data |
|
|
206
|
+
| [Benchmark artifacts](https://drive.google.com/drive/folders/1KjhQe4yvLMSBEJdxXZ6oC9Wi-pw5a439) | CAMEO22, CASP14, CASP15 and FoldBench results |
|
|
207
|
+
|
|
208
|
+
|
|
209
|
+
## Training
|
|
210
|
+
|
|
211
|
+
Install the training dependencies with `pip install -e ".[fold,train,cuequiv]"`. AtlasFold monomer training uses four progressively longer crop stages, and AtlasFold-M fine-tuning uses three stages initialized from the monomer model. See the [training guide](docs/training.md) for data setup, released intermediate checkpoints, complete commands, and configuration overrides, and the [data guide](docs/data.md) for the released dataset layout and provenance.
|
|
212
|
+
|
|
213
|
+
## Citation
|
|
214
|
+
|
|
215
|
+
```bibtex
|
|
216
|
+
@article{seo2026atlasfold,
|
|
217
|
+
author = {Seo, Seonghwan and Kim, Hyeongwoo and Moon, Seokhyun and Kim, Woo Youn and {Team KAIST}},
|
|
218
|
+
title = {AtlasFold: Protein structure prediction with metagenomic-scale language models},
|
|
219
|
+
year = {2026},
|
|
220
|
+
doi = {10.64898/2026.09.04.749352},
|
|
221
|
+
URL = {https://www.biorxiv.org/content/10.64898/2026.09.04.749352v2},
|
|
222
|
+
journal = {bioRxiv}
|
|
223
|
+
}
|
|
224
|
+
```
|
|
225
|
+
|
|
226
|
+
## Acknowledgements
|
|
227
|
+
|
|
228
|
+
This project was developed as part of the **K-Fold** initiative supported by the Ministry of Science and ICT (MSIT) of the Republic of Korea. The K-Fold project for biomolecular complex prediction is currently under active development with numerous contributors at KAIST and will be released soon!
|
|
229
|
+
|
|
230
|
+
---
|
|
231
|
+
|
|
232
|
+
I would like to thank [Dr. Hyeongwoo Kim](https://scholar.google.com/citations?user=YpiY1q8AAAAJ&hl=en&oi=ao), [Dr. Seokhyun Moon](https://scholar.google.com/citations?hl=en&user=U1j8Ip8AAAAJ), and [Prof. Woo Youn Kim](https://scholar.google.com/citations?user=elJ5KrcAAAAJ&hl=en) for their guidance and support during the development of AtlasFold.
|
|
233
|
+
|
|
234
|
+
This project is built upon the pioneering works of Google DeepMind, Meta AI, OpenFold Consortium, and EvolutionaryScale in the fields of biomolecular language modeling and structure prediction. I am deeply grateful to the open-source community for advancing the fields of biomolecular language modeling and structure prediction.
|
|
235
|
+
|
|
236
|
+
**Foundations of AtlasLM:**
|
|
237
|
+
|
|
238
|
+
- **ESM2**: Lin, Zeming, et al. "Evolutionary-scale prediction of atomic-level protein structure with a language model." Science 379.6637 (2023): 1123-1130.
|
|
239
|
+
- **ESM3**: Hayes, Thomas, et al. "Simulating 500 million years of evolution with a language model." Science 387.6736 (2025): 850-858.
|
|
240
|
+
- **ESMC**: ESM Team. "ESM Cambrian: Revealing the mysteries of proteins with unsupervised learning." EvolutionaryScale Website, December 4, 2024. https://evolutionaryscale.ai/blog/esm-cambrian.
|
|
241
|
+
|
|
242
|
+
**Foundations of AtlasFold:**
|
|
243
|
+
|
|
244
|
+
- **AlphaFold2**: Jumper, John, et al. "Highly accurate protein structure prediction with AlphaFold." Nature 596.7873 (2021): 583-589.
|
|
245
|
+
- **OpenFold**: Ahdritz, Gustaf, et al. "OpenFold: retraining AlphaFold2 yields new insights into its learning mechanisms and capacity for generalization." Nature Methods 21.8 (2024): 1514-1524.
|
|
246
|
+
- **ESMFold**: Lin, Zeming, et al. "Evolutionary-scale prediction of atomic-level protein structure with a language model." Science 379.6637 (2023): 1123-1130.
|
|
247
|
+
- **AlphaFold3**: Abramson, Josh, et al. "Accurate structure prediction of biomolecular interactions with AlphaFold 3." Nature 630.8016 (2024): 493-500.
|
|
248
|
+
- **SimpleFold**: Wang, Yuyang, et al. "SimpleFold: Folding proteins is simpler than you think." arXiv preprint arXiv:2509.18480 (2025).
|
|
249
|
+
- **ESMFold2**: Candido, Salvatore, et al. "Language Modeling Materializes a World Model of Protein Biology." [bioRxiv preprint](https://www.biorxiv.org/content/10.64898/2026.06.03.729735) (2026).
|
|
250
|
+
|
|
251
|
+
## License
|
|
252
|
+
|
|
253
|
+
The source code, model weights, and released datasets are licensed under the [MIT License](LICENSE).
|
|
@@ -0,0 +1,217 @@
|
|
|
1
|
+
# AtlasFold
|
|
2
|
+
|
|
3
|
+
[\[Paper\]](https://www.biorxiv.org/content/10.64898/2026.09.04.749352v2), [\[PDF\]](docs/atlasfold.pdf)
|
|
4
|
+
|
|
5
|
+
The **Atlas family** is a collection of open protein models for sequence representation and structure prediction.
|
|
6
|
+
AtlasLM is a protein language model (PLM), while AtlasFold and AtlasFold-M are trainable PLM-based models for protein folding and co-folding, respectively.
|
|
7
|
+
|
|
8
|
+

|
|
9
|
+
|
|
10
|
+
AtlasFold achieves state-of-the-art accuracy among protein language model-based folding methods. AtlasFold and AtlasFold-M predict structures without an MSA search.
|
|
11
|
+
This repository provides pretrained models, training and inference code, staged training configurations, and preprocessing workflows for monomer and multimer folding.
|
|
12
|
+
|
|
13
|
+
## Installation
|
|
14
|
+
|
|
15
|
+
AtlasFold requires Python 3.10 or later. A CUDA GPU is recommended for structure prediction.
|
|
16
|
+
|
|
17
|
+
Install the inference dependencies from PyPI:
|
|
18
|
+
|
|
19
|
+
```bash
|
|
20
|
+
pip install "atlasfold[fold,cuequiv]"
|
|
21
|
+
```
|
|
22
|
+
|
|
23
|
+
To use AtlasLM as a standalone protein language model:
|
|
24
|
+
|
|
25
|
+
```bash
|
|
26
|
+
pip install atlasfold
|
|
27
|
+
```
|
|
28
|
+
|
|
29
|
+
To install the latest development version from GitHub:
|
|
30
|
+
|
|
31
|
+
```bash
|
|
32
|
+
git clone https://github.com/SeonghwanSeo/atlasfold.git
|
|
33
|
+
cd atlasfold
|
|
34
|
+
pip install -e ".[fold,cuequiv]"
|
|
35
|
+
```
|
|
36
|
+
|
|
37
|
+
AtlasFold supports [cuEquivariance](https://docs.nvidia.com/cuda/cuequivariance/) kernels for faster inference. For systems without compatible NVIDIA CUDA hardware, install `atlasfold[fold]` instead.
|
|
38
|
+
|
|
39
|
+
## Running your first prediction
|
|
40
|
+
|
|
41
|
+
For a monomer, save one protein sequence per FASTA record:
|
|
42
|
+
|
|
43
|
+
```text
|
|
44
|
+
>protein_a
|
|
45
|
+
MKTAYIAKQRQISFVKSHFSRQDILDLWIYHTQGYFPD
|
|
46
|
+
>protein_b
|
|
47
|
+
FNPVGVAFKGNNGKYLSRIHRSGIDYTEFAKDNTD
|
|
48
|
+
```
|
|
49
|
+
|
|
50
|
+
Then run:
|
|
51
|
+
|
|
52
|
+
```bash
|
|
53
|
+
atlasfold monomer --input-fasta monomer.fasta --out-dir predictions/monomer/
|
|
54
|
+
```
|
|
55
|
+
|
|
56
|
+
For a protein complex, use one FASTA record per complex and separate chains with `:`:
|
|
57
|
+
|
|
58
|
+
```text
|
|
59
|
+
>complex_a
|
|
60
|
+
MKTAYIAKQRQISFVKSHFS:GGHVDHGKSTTTGHLIYK
|
|
61
|
+
>complex_b
|
|
62
|
+
MKEGFYWIQHNGRVQVAYYTHGVTEDLETGQTIIGVWHLTQGDDICHNGEAEILAGPLEPPI:MKEGFYWIQHNGRVQVA
|
|
63
|
+
YYTHGVTEDLETGQTIIGVWHLTQGDDICHNGEAEILAGPLEPPI
|
|
64
|
+
```
|
|
65
|
+
|
|
66
|
+
```bash
|
|
67
|
+
atlasfold multimer --input-fasta multimer.fasta --out-dir predictions/multimer/
|
|
68
|
+
```
|
|
69
|
+
|
|
70
|
+
Template-assisted inference for AtlasFold-M is not supported by the current runner or CLI.
|
|
71
|
+
|
|
72
|
+
The repository entry point provides the same interface:
|
|
73
|
+
|
|
74
|
+
```bash
|
|
75
|
+
python run_atlasfold.py monomer --input-fasta monomer.fasta --out-dir predictions/monomer/
|
|
76
|
+
python run_atlasfold.py multimer --input-fasta multimer.fasta --out-dir predictions/multimer/
|
|
77
|
+
```
|
|
78
|
+
|
|
79
|
+
Both AtlasFold and AtlasFold-M support batched inference with multiple FASTA records, enabling high-throughput structure prediction.
|
|
80
|
+
|
|
81
|
+
Run `atlasfold monomer --help` or `atlasfold multimer --help` for all options, and see the [inference guide](docs/inference.md) for batching, sampling, confidence values, and output formats.
|
|
82
|
+
|
|
83
|
+
## Performance and GPU memory
|
|
84
|
+
|
|
85
|
+
Peak GPU memory grows with total residue length when generating five diffusion samples.
|
|
86
|
+
|
|
87
|
+
| Total residues | AtlasFold | AtlasFold-M |
|
|
88
|
+
| ---: | ---: | ---: |
|
|
89
|
+
| 256 | 7.22 GiB | 7.31 GiB |
|
|
90
|
+
| 512 | 8.80 GiB | 9.07 GiB |
|
|
91
|
+
| 1,024 | 15.02 GiB | 16.05 GiB |
|
|
92
|
+
| 1,536 | 25.36 GiB | 27.63 GiB |
|
|
93
|
+
| 2,048 | 39.81 GiB | 43.83 GiB |
|
|
94
|
+
|
|
95
|
+
For multi-target workloads, batched inference substantially increases throughput. With five diffusion samples and up to 4,096 residues processed per batch:
|
|
96
|
+
|
|
97
|
+
| Workload | Unbatched | Batched | Throughput gain |
|
|
98
|
+
| --- | ---: | ---: | ---: |
|
|
99
|
+
| AtlasFold, 64-residue monomers | 0.775 sequences/s | 9.663 sequences/s | 12.5× |
|
|
100
|
+
| AtlasFold-M, 256-residue complexes | 0.149 complexes/s | 0.314 complexes/s | 2.1× |
|
|
101
|
+
|
|
102
|
+
Measurements were collected on a single NVIDIA B200 using PyTorch 2.10.0, CUDA 12.8, and cuEquivariance 0.10.0. See the [performance guide](docs/performance.md) for the complete memory and runtime measurements.
|
|
103
|
+
|
|
104
|
+
## Python API
|
|
105
|
+
|
|
106
|
+
### Structure prediction
|
|
107
|
+
|
|
108
|
+
```python
|
|
109
|
+
from atlasfold.pretrained import get_runner, load_model
|
|
110
|
+
|
|
111
|
+
model = load_model("atlasfold", device="cuda")
|
|
112
|
+
runner = get_runner(model)
|
|
113
|
+
result = runner.fold(
|
|
114
|
+
"protein_a",
|
|
115
|
+
"MKTAYIAKQRQISFVKSHFSRQDILDLWIYHTQGYFPD",
|
|
116
|
+
num_samples=5,
|
|
117
|
+
)
|
|
118
|
+
|
|
119
|
+
print(result.best.avg_plddt, result.best.ptm)
|
|
120
|
+
with open("protein_a.pdb", "w") as handle:
|
|
121
|
+
handle.write(result.best.to_pdb())
|
|
122
|
+
with open("protein_a.cif", "w") as handle:
|
|
123
|
+
handle.write(result.best.to_mmcif())
|
|
124
|
+
```
|
|
125
|
+
|
|
126
|
+
Use `load_model("atlasfold-m", device="cuda")` for a complex and pass a list of chain sequences to `runner.fold()`. The Python API defaults to CPU if `device` is omitted; the CLI selects CUDA automatically when available.
|
|
127
|
+
|
|
128
|
+
### Protein language model
|
|
129
|
+
|
|
130
|
+
```python
|
|
131
|
+
import torch
|
|
132
|
+
|
|
133
|
+
from atlaslm import load_model
|
|
134
|
+
|
|
135
|
+
model = load_model("atlaslm-3b", device="cuda", dtype=torch.bfloat16)
|
|
136
|
+
output = model.embed_sequences(
|
|
137
|
+
["MKTAYIAKQRQISFVKSHFSRQDILDLWIYHTQGYFPD"],
|
|
138
|
+
return_hidden_states=True,
|
|
139
|
+
)
|
|
140
|
+
|
|
141
|
+
print(output.embeddings.shape)
|
|
142
|
+
print(len(output.hidden_states))
|
|
143
|
+
```
|
|
144
|
+
|
|
145
|
+
Pass `return_attentions=True` to return attention maps. Attention tensors grow quadratically with sequence length and can require substantially more memory.
|
|
146
|
+
|
|
147
|
+
## Available models
|
|
148
|
+
|
|
149
|
+
| Model | Parameters | Weight download | Use | Weights |
|
|
150
|
+
| --- | ---: | ---: | --- | --- |
|
|
151
|
+
| AtlasLM-600M | 575M | 1.07 GiB | Protein sequence representations | [`SeonghwanSeo/atlaslm-600m-base`](https://huggingface.co/SeonghwanSeo/atlaslm-600m-base) |
|
|
152
|
+
| AtlasLM-3B | 3.06B | 5.71 GiB | Protein sequence representations and AtlasFold backbone | [`SeonghwanSeo/atlaslm-3b-base`](https://huggingface.co/SeonghwanSeo/atlaslm-3b-base) |
|
|
153
|
+
| AtlasFold | 215M + AtlasLM-3B | 0.80 GiB + AtlasLM-3B | Monomer structure prediction | [`SeonghwanSeo/atlasfold-260703`](https://huggingface.co/SeonghwanSeo/atlasfold-260703) |
|
|
154
|
+
| AtlasFold-M | 220M + AtlasLM-3B | 0.82 GiB + AtlasLM-3B | Protein-complex structure prediction | [`SeonghwanSeo/atlasfold-m-260725`](https://huggingface.co/SeonghwanSeo/atlasfold-m-260725) |
|
|
155
|
+
|
|
156
|
+
AtlasLM-600M is deprecated now that ESMC-600M is available for commercial use. AtlasLM-3B is the recommended AtlasLM checkpoint.
|
|
157
|
+
|
|
158
|
+
## Evaluation
|
|
159
|
+
|
|
160
|
+
Evaluation protocols and results for AtlasFold and AtlasFold-M are provided in the [benchmark documentation](docs/benchmarks.md). The associated prediction structures and evaluation artifacts are available from the release folder below.
|
|
161
|
+
|
|
162
|
+
## Checkpoints, data, and benchmark artifacts
|
|
163
|
+
|
|
164
|
+
Large release artifacts are hosted in the [AtlasFold Google Drive folder](https://drive.google.com/drive/folders/1xjSBmbCFqghWj8xKYuEIDISKoQkCX45I?usp=sharing):
|
|
165
|
+
|
|
166
|
+
| Artifact | Contents |
|
|
167
|
+
| --- | --- |
|
|
168
|
+
| [Intermediate checkpoints](https://drive.google.com/drive/folders/1TDz2Ng4-zYfpTxlqg5wD3MkzwJyL2gTm) | AtlasLM pretraining checkpoints and AtlasFold/AtlasFold-M staged training checkpoints |
|
|
169
|
+
| [Structural datasets](https://drive.google.com/drive/folders/1EiRTKSUL3iD_MQ_0qmj5Sb-2KMh-3kmS) | Processed monomer and multimer training and validation data |
|
|
170
|
+
| [Benchmark artifacts](https://drive.google.com/drive/folders/1KjhQe4yvLMSBEJdxXZ6oC9Wi-pw5a439) | CAMEO22, CASP14, CASP15 and FoldBench results |
|
|
171
|
+
|
|
172
|
+
|
|
173
|
+
## Training
|
|
174
|
+
|
|
175
|
+
Install the training dependencies with `pip install -e ".[fold,train,cuequiv]"`. AtlasFold monomer training uses four progressively longer crop stages, and AtlasFold-M fine-tuning uses three stages initialized from the monomer model. See the [training guide](docs/training.md) for data setup, released intermediate checkpoints, complete commands, and configuration overrides, and the [data guide](docs/data.md) for the released dataset layout and provenance.
|
|
176
|
+
|
|
177
|
+
## Citation
|
|
178
|
+
|
|
179
|
+
```bibtex
|
|
180
|
+
@article{seo2026atlasfold,
|
|
181
|
+
author = {Seo, Seonghwan and Kim, Hyeongwoo and Moon, Seokhyun and Kim, Woo Youn and {Team KAIST}},
|
|
182
|
+
title = {AtlasFold: Protein structure prediction with metagenomic-scale language models},
|
|
183
|
+
year = {2026},
|
|
184
|
+
doi = {10.64898/2026.09.04.749352},
|
|
185
|
+
URL = {https://www.biorxiv.org/content/10.64898/2026.09.04.749352v2},
|
|
186
|
+
journal = {bioRxiv}
|
|
187
|
+
}
|
|
188
|
+
```
|
|
189
|
+
|
|
190
|
+
## Acknowledgements
|
|
191
|
+
|
|
192
|
+
This project was developed as part of the **K-Fold** initiative supported by the Ministry of Science and ICT (MSIT) of the Republic of Korea. The K-Fold project for biomolecular complex prediction is currently under active development with numerous contributors at KAIST and will be released soon!
|
|
193
|
+
|
|
194
|
+
---
|
|
195
|
+
|
|
196
|
+
I would like to thank [Dr. Hyeongwoo Kim](https://scholar.google.com/citations?user=YpiY1q8AAAAJ&hl=en&oi=ao), [Dr. Seokhyun Moon](https://scholar.google.com/citations?hl=en&user=U1j8Ip8AAAAJ), and [Prof. Woo Youn Kim](https://scholar.google.com/citations?user=elJ5KrcAAAAJ&hl=en) for their guidance and support during the development of AtlasFold.
|
|
197
|
+
|
|
198
|
+
This project is built upon the pioneering works of Google DeepMind, Meta AI, OpenFold Consortium, and EvolutionaryScale in the fields of biomolecular language modeling and structure prediction. I am deeply grateful to the open-source community for advancing the fields of biomolecular language modeling and structure prediction.
|
|
199
|
+
|
|
200
|
+
**Foundations of AtlasLM:**
|
|
201
|
+
|
|
202
|
+
- **ESM2**: Lin, Zeming, et al. "Evolutionary-scale prediction of atomic-level protein structure with a language model." Science 379.6637 (2023): 1123-1130.
|
|
203
|
+
- **ESM3**: Hayes, Thomas, et al. "Simulating 500 million years of evolution with a language model." Science 387.6736 (2025): 850-858.
|
|
204
|
+
- **ESMC**: ESM Team. "ESM Cambrian: Revealing the mysteries of proteins with unsupervised learning." EvolutionaryScale Website, December 4, 2024. https://evolutionaryscale.ai/blog/esm-cambrian.
|
|
205
|
+
|
|
206
|
+
**Foundations of AtlasFold:**
|
|
207
|
+
|
|
208
|
+
- **AlphaFold2**: Jumper, John, et al. "Highly accurate protein structure prediction with AlphaFold." Nature 596.7873 (2021): 583-589.
|
|
209
|
+
- **OpenFold**: Ahdritz, Gustaf, et al. "OpenFold: retraining AlphaFold2 yields new insights into its learning mechanisms and capacity for generalization." Nature Methods 21.8 (2024): 1514-1524.
|
|
210
|
+
- **ESMFold**: Lin, Zeming, et al. "Evolutionary-scale prediction of atomic-level protein structure with a language model." Science 379.6637 (2023): 1123-1130.
|
|
211
|
+
- **AlphaFold3**: Abramson, Josh, et al. "Accurate structure prediction of biomolecular interactions with AlphaFold 3." Nature 630.8016 (2024): 493-500.
|
|
212
|
+
- **SimpleFold**: Wang, Yuyang, et al. "SimpleFold: Folding proteins is simpler than you think." arXiv preprint arXiv:2509.18480 (2025).
|
|
213
|
+
- **ESMFold2**: Candido, Salvatore, et al. "Language Modeling Materializes a World Model of Protein Biology." [bioRxiv preprint](https://www.biorxiv.org/content/10.64898/2026.06.03.729735) (2026).
|
|
214
|
+
|
|
215
|
+
## License
|
|
216
|
+
|
|
217
|
+
The source code, model weights, and released datasets are licensed under the [MIT License](LICENSE).
|
|
@@ -0,0 +1,51 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=77.0.0"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "atlasfold"
|
|
7
|
+
version = "1.0.0"
|
|
8
|
+
description = "Protein language model-based protein folding and co-folding"
|
|
9
|
+
authors = [
|
|
10
|
+
{ name = "SeonghwanSeo", email = "shwan0106@gmail.com" }
|
|
11
|
+
]
|
|
12
|
+
readme = "README.md"
|
|
13
|
+
requires-python = ">=3.10"
|
|
14
|
+
license = "MIT"
|
|
15
|
+
keywords = ["protein", "protein-language-model", "protein-folding", "cofolding"]
|
|
16
|
+
dependencies = [
|
|
17
|
+
"numpy",
|
|
18
|
+
"torch",
|
|
19
|
+
"huggingface_hub",
|
|
20
|
+
]
|
|
21
|
+
|
|
22
|
+
[project.urls]
|
|
23
|
+
Homepage = "https://github.com/SeonghwanSeo/atlasfold"
|
|
24
|
+
Documentation = "https://github.com/SeonghwanSeo/atlasfold#readme"
|
|
25
|
+
Repository = "https://github.com/SeonghwanSeo/atlasfold"
|
|
26
|
+
Issues = "https://github.com/SeonghwanSeo/atlasfold/issues"
|
|
27
|
+
|
|
28
|
+
[project.scripts]
|
|
29
|
+
atlasfold = "atlasfold.cli:main"
|
|
30
|
+
|
|
31
|
+
[project.optional-dependencies]
|
|
32
|
+
fold = [
|
|
33
|
+
"scipy",
|
|
34
|
+
"einops",
|
|
35
|
+
"gemmi",
|
|
36
|
+
"omegaconf",
|
|
37
|
+
"numba", # for SASA computation
|
|
38
|
+
]
|
|
39
|
+
cuequiv = [
|
|
40
|
+
"cuequivariance_torch>=0.6.0",
|
|
41
|
+
"cuequivariance_ops_cu12>=0.6.0",
|
|
42
|
+
"cuequivariance_ops_torch_cu12>=0.6.0",
|
|
43
|
+
]
|
|
44
|
+
train = [
|
|
45
|
+
"msgpack",
|
|
46
|
+
"lightning", # for training loop
|
|
47
|
+
"wandb", # for experiment tracking
|
|
48
|
+
"lmdb", # for training dataset storage
|
|
49
|
+
"pandas",
|
|
50
|
+
"zstandard"
|
|
51
|
+
]
|
|
File without changes
|
|
@@ -0,0 +1,67 @@
|
|
|
1
|
+
"""Unified AtlasFold command-line interface."""
|
|
2
|
+
|
|
3
|
+
import argparse
|
|
4
|
+
import sys
|
|
5
|
+
from collections.abc import Sequence
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
def create_parser() -> argparse.ArgumentParser:
|
|
9
|
+
"""Create the top-level parser used to select an inference pipeline."""
|
|
10
|
+
parser = argparse.ArgumentParser(description="Run AtlasFold inference.")
|
|
11
|
+
commands = parser.add_subparsers(
|
|
12
|
+
dest="model",
|
|
13
|
+
metavar="{monomer,multimer}",
|
|
14
|
+
required=True,
|
|
15
|
+
)
|
|
16
|
+
commands.add_parser(
|
|
17
|
+
"monomer",
|
|
18
|
+
add_help=False,
|
|
19
|
+
help="Run AtlasFold monomer inference.",
|
|
20
|
+
)
|
|
21
|
+
commands.add_parser(
|
|
22
|
+
"multimer",
|
|
23
|
+
add_help=False,
|
|
24
|
+
help="Run AtlasFold-Multimer inference.",
|
|
25
|
+
)
|
|
26
|
+
return parser
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def _parse_command(argv: Sequence[str] | None) -> tuple[str, list[str]]:
|
|
30
|
+
values = list(sys.argv[1:] if argv is None else argv)
|
|
31
|
+
if not values or values[0] in {"-h", "--help"}:
|
|
32
|
+
create_parser().parse_args(values)
|
|
33
|
+
|
|
34
|
+
command = values[0]
|
|
35
|
+
if command not in {"monomer", "multimer"}:
|
|
36
|
+
create_parser().parse_args([command])
|
|
37
|
+
|
|
38
|
+
return command, values[1:]
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
def main(argv: Sequence[str] | None = None) -> None:
|
|
42
|
+
"""Run the selected AtlasFold inference pipeline."""
|
|
43
|
+
command, command_argv = _parse_command(argv)
|
|
44
|
+
prog = f"{create_parser().prog} {command}"
|
|
45
|
+
|
|
46
|
+
if command == "monomer":
|
|
47
|
+
from atlasfold.cli import monomer
|
|
48
|
+
|
|
49
|
+
parser = monomer.create_parser(prog=prog)
|
|
50
|
+
args = parser.parse_args(command_argv)
|
|
51
|
+
monomer.run(args)
|
|
52
|
+
return
|
|
53
|
+
|
|
54
|
+
elif command == "multimer":
|
|
55
|
+
from atlasfold.cli import multimer
|
|
56
|
+
|
|
57
|
+
parser = multimer.create_parser(prog=prog)
|
|
58
|
+
args = parser.parse_args(command_argv)
|
|
59
|
+
multimer.run(args)
|
|
60
|
+
return
|
|
61
|
+
|
|
62
|
+
else:
|
|
63
|
+
raise ValueError(f"Unsupported model: {command}")
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
if __name__ == "__main__":
|
|
67
|
+
main()
|