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.
Files changed (109) hide show
  1. atlasfold-1.0.0/LICENSE +21 -0
  2. atlasfold-1.0.0/PKG-INFO +253 -0
  3. atlasfold-1.0.0/README.md +217 -0
  4. atlasfold-1.0.0/pyproject.toml +51 -0
  5. atlasfold-1.0.0/setup.cfg +4 -0
  6. atlasfold-1.0.0/src/atlasfold/__init__.py +0 -0
  7. atlasfold-1.0.0/src/atlasfold/cli/__init__.py +67 -0
  8. atlasfold-1.0.0/src/atlasfold/cli/__main__.py +6 -0
  9. atlasfold-1.0.0/src/atlasfold/cli/monomer.py +496 -0
  10. atlasfold-1.0.0/src/atlasfold/cli/multigpu.py +125 -0
  11. atlasfold-1.0.0/src/atlasfold/cli/multimer.py +513 -0
  12. atlasfold-1.0.0/src/atlasfold/common/__init__.py +0 -0
  13. atlasfold-1.0.0/src/atlasfold/common/ccd.py +287 -0
  14. atlasfold-1.0.0/src/atlasfold/common/featurize.py +308 -0
  15. atlasfold-1.0.0/src/atlasfold/common/file_io.py +327 -0
  16. atlasfold-1.0.0/src/atlasfold/common/metadata.py +232 -0
  17. atlasfold-1.0.0/src/atlasfold/common/process.py +70 -0
  18. atlasfold-1.0.0/src/atlasfold/common/protein.py +307 -0
  19. atlasfold-1.0.0/src/atlasfold/common/residue_constants.py +383 -0
  20. atlasfold-1.0.0/src/atlasfold/common/templates.py +286 -0
  21. atlasfold-1.0.0/src/atlasfold/configs/__init__.py +0 -0
  22. atlasfold-1.0.0/src/atlasfold/configs/atlasfold.py +43 -0
  23. atlasfold-1.0.0/src/atlasfold/configs/atlasfold_multimer.py +54 -0
  24. atlasfold-1.0.0/src/atlasfold/data/__init__.py +0 -0
  25. atlasfold-1.0.0/src/atlasfold/data/cif_factory.py +244 -0
  26. atlasfold-1.0.0/src/atlasfold/data/fasta.py +53 -0
  27. atlasfold-1.0.0/src/atlasfold/data/mmseq2.py +174 -0
  28. atlasfold-1.0.0/src/atlasfold/model/__init__.py +3 -0
  29. atlasfold-1.0.0/src/atlasfold/model/model.py +558 -0
  30. atlasfold-1.0.0/src/atlasfold/model/model_multimer.py +654 -0
  31. atlasfold-1.0.0/src/atlasfold/model/network/__init__.py +0 -0
  32. atlasfold-1.0.0/src/atlasfold/model/network/atom_attention.py +211 -0
  33. atlasfold-1.0.0/src/atlasfold/model/network/attention.py +294 -0
  34. atlasfold-1.0.0/src/atlasfold/model/network/block.py +293 -0
  35. atlasfold-1.0.0/src/atlasfold/model/network/confidence_head.py +515 -0
  36. atlasfold-1.0.0/src/atlasfold/model/network/diffusion_head.py +450 -0
  37. atlasfold-1.0.0/src/atlasfold/model/network/diffusion_transformer.py +513 -0
  38. atlasfold-1.0.0/src/atlasfold/model/network/distogram_head.py +50 -0
  39. atlasfold-1.0.0/src/atlasfold/model/network/misc.py +138 -0
  40. atlasfold-1.0.0/src/atlasfold/model/network/primitives/__init__.py +26 -0
  41. atlasfold-1.0.0/src/atlasfold/model/network/primitives/activation.py +15 -0
  42. atlasfold-1.0.0/src/atlasfold/model/network/primitives/dropout.py +45 -0
  43. atlasfold-1.0.0/src/atlasfold/model/network/primitives/initialize.py +82 -0
  44. atlasfold-1.0.0/src/atlasfold/model/network/primitives/linear.py +110 -0
  45. atlasfold-1.0.0/src/atlasfold/model/network/primitives/normalization.py +77 -0
  46. atlasfold-1.0.0/src/atlasfold/model/network/primitives/transition.py +28 -0
  47. atlasfold-1.0.0/src/atlasfold/model/network/primitives/triangle_update.py +515 -0
  48. atlasfold-1.0.0/src/atlasfold/model/network/rel_pos_encoding.py +141 -0
  49. atlasfold-1.0.0/src/atlasfold/model/network/template.py +239 -0
  50. atlasfold-1.0.0/src/atlasfold/model/network/trunk.py +168 -0
  51. atlasfold-1.0.0/src/atlasfold/model/utils/__init__.py +0 -0
  52. atlasfold-1.0.0/src/atlasfold/model/utils/confidence_metrics.py +392 -0
  53. atlasfold-1.0.0/src/atlasfold/pretrained.py +156 -0
  54. atlasfold-1.0.0/src/atlasfold/runner.py +484 -0
  55. atlasfold-1.0.0/src/atlasfold/runner_multimer.py +579 -0
  56. atlasfold-1.0.0/src/atlasfold/train/__init__.py +0 -0
  57. atlasfold-1.0.0/src/atlasfold/train/config.py +74 -0
  58. atlasfold-1.0.0/src/atlasfold/train/losses/__init__.py +1 -0
  59. atlasfold-1.0.0/src/atlasfold/train/losses/confidence.py +386 -0
  60. atlasfold-1.0.0/src/atlasfold/train/losses/diffusion.py +206 -0
  61. atlasfold-1.0.0/src/atlasfold/train/losses/distogram.py +62 -0
  62. atlasfold-1.0.0/src/atlasfold/train/monomer/__init__.py +0 -0
  63. atlasfold-1.0.0/src/atlasfold/train/monomer/cropper.py +179 -0
  64. atlasfold-1.0.0/src/atlasfold/train/monomer/datamodule.py +149 -0
  65. atlasfold-1.0.0/src/atlasfold/train/monomer/dataset.py +552 -0
  66. atlasfold-1.0.0/src/atlasfold/train/monomer/model_train.py +363 -0
  67. atlasfold-1.0.0/src/atlasfold/train/monomer/structure_alignment.py +106 -0
  68. atlasfold-1.0.0/src/atlasfold/train/monomer/train_module.py +681 -0
  69. atlasfold-1.0.0/src/atlasfold/train/monomer/validation_metrics.py +74 -0
  70. atlasfold-1.0.0/src/atlasfold/train/multimer/__init__.py +0 -0
  71. atlasfold-1.0.0/src/atlasfold/train/multimer/cropper.py +341 -0
  72. atlasfold-1.0.0/src/atlasfold/train/multimer/datamodule.py +168 -0
  73. atlasfold-1.0.0/src/atlasfold/train/multimer/dataset.py +1135 -0
  74. atlasfold-1.0.0/src/atlasfold/train/multimer/model_train.py +369 -0
  75. atlasfold-1.0.0/src/atlasfold/train/multimer/train_alignment.py +472 -0
  76. atlasfold-1.0.0/src/atlasfold/train/multimer/train_module.py +460 -0
  77. atlasfold-1.0.0/src/atlasfold/train/multimer/validation_metrics.py +432 -0
  78. atlasfold-1.0.0/src/atlasfold/train/utils/__init__.py +0 -0
  79. atlasfold-1.0.0/src/atlasfold/train/utils/dl_sampler.py +106 -0
  80. atlasfold-1.0.0/src/atlasfold/train/utils/ema.py +111 -0
  81. atlasfold-1.0.0/src/atlasfold/train/utils/gradient_logging.py +25 -0
  82. atlasfold-1.0.0/src/atlasfold/train/utils/kernels/__init__.py +0 -0
  83. atlasfold-1.0.0/src/atlasfold/train/utils/kernels/cdist.py +328 -0
  84. atlasfold-1.0.0/src/atlasfold/train/utils/lr_scheduler.py +54 -0
  85. atlasfold-1.0.0/src/atlasfold/utils/__init__.py +0 -0
  86. atlasfold-1.0.0/src/atlasfold/utils/checkpointing.py +99 -0
  87. atlasfold-1.0.0/src/atlasfold/utils/geometry/__init__.py +0 -0
  88. atlasfold-1.0.0/src/atlasfold/utils/geometry/metrics.py +184 -0
  89. atlasfold-1.0.0/src/atlasfold/utils/geometry/random_augment.py +271 -0
  90. atlasfold-1.0.0/src/atlasfold/utils/geometry/rigid_align.py +525 -0
  91. atlasfold-1.0.0/src/atlasfold/utils/kernel.py +36 -0
  92. atlasfold-1.0.0/src/atlasfold/utils/misc.py +82 -0
  93. atlasfold-1.0.0/src/atlasfold/utils/rasa.py +236 -0
  94. atlasfold-1.0.0/src/atlasfold/utils/torch_utils.py +162 -0
  95. atlasfold-1.0.0/src/atlasfold.egg-info/PKG-INFO +253 -0
  96. atlasfold-1.0.0/src/atlasfold.egg-info/SOURCES.txt +107 -0
  97. atlasfold-1.0.0/src/atlasfold.egg-info/dependency_links.txt +1 -0
  98. atlasfold-1.0.0/src/atlasfold.egg-info/entry_points.txt +2 -0
  99. atlasfold-1.0.0/src/atlasfold.egg-info/requires.txt +23 -0
  100. atlasfold-1.0.0/src/atlasfold.egg-info/top_level.txt +2 -0
  101. atlasfold-1.0.0/src/atlaslm/__init__.py +2 -0
  102. atlasfold-1.0.0/src/atlaslm/alphabet.py +37 -0
  103. atlasfold-1.0.0/src/atlaslm/layers/__init__.py +2 -0
  104. atlasfold-1.0.0/src/atlaslm/layers/attention.py +113 -0
  105. atlasfold-1.0.0/src/atlaslm/layers/regression_head.py +17 -0
  106. atlasfold-1.0.0/src/atlaslm/layers/rotary.py +104 -0
  107. atlasfold-1.0.0/src/atlaslm/layers/transformer_stack.py +154 -0
  108. atlasfold-1.0.0/src/atlaslm/model.py +285 -0
  109. atlasfold-1.0.0/src/atlaslm/pretrained.py +152 -0
@@ -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.
@@ -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
+ ![AtlasFold predictions (blue) overlaid with experimental structures (gray) for CASP15 target T1183 (8IFX_B) and complex 8OI4.](docs/images/atlasfold-predictions.png)
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
+ ![AtlasFold predictions (blue) overlaid with experimental structures (gray) for CASP15 target T1183 (8IFX_B) and complex 8OI4.](docs/images/atlasfold-predictions.png)
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
+ ]
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
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()
@@ -0,0 +1,6 @@
1
+ """Allow running the AtlasFold CLI with ``python -m atlasfold.cli``."""
2
+
3
+ from atlasfold.cli import main
4
+
5
+ if __name__ == "__main__":
6
+ main()