You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
* Add dispersion models and update internal check script
* Update model with charge+spin conditioning. Add model stubs. Still need to add weights
* Add paths. Handle charge and spin in atoms.info
* Final working implementation
* Update readmes
* Address comments
* Address comments
---------
Co-authored-by: ben rhodes <benrhodes@bens-MacBook-Pro.local>
|**orb-v1**|[link](https://orbitalmaterials-public-models.s3.us-west-1.amazonaws.com/forcefields/orbff-v1-20240827.ckpt)| 92897eda08609425ee001955c7885139 | Yes | No | Full dataset pretraining, MPTraj + Alexandria finetuning |
7
-
|**orb-mptraj-only-v1**|[link](https://orbitalmaterials-public-models.s3.us-west-1.amazonaws.com/forcefields/orbff-mptraj-only-v1-20240827.ckpt)| ff42a2bc1e1f50b5f3ee2a20b83cf3a2 | Yes | No | MPTraj pretraining and finetuning only |
8
-
|**orb-d3-v1**|[link](https://orbitalmaterials-public-models.s3.us-west-1.amazonaws.com/forcefields/orb-d3-v1-20240902.ckpt)| 470c7d3482ead3bc97cd4b46382d5e47 | No | Yes | Full dataset pretraining, MPTraj + Alexandria finetuning, integrated D3 corrections |
9
-
|**orb-d3-sm-v1**|[link](https://orbitalmaterials-public-models.s3.us-west-1.amazonaws.com/forcefields/orb-d3-sm-v1-20240902.ckpt)| 64fe91603e46ad5fa695525e3f1e9397 | No | Yes | First 10 layers of a pretrained model finetuned on mptrj + alexandria with D3 corrections |
10
-
|**orb-d3-xs-v1**|[link](https://orbitalmaterials-public-models.s3.us-west-1.amazonaws.com/forcefields/orb-d3-xs-v1-20240902.ckpt)| 79d042f9f16c4407795426a75498fbb7 | No | Yes | First 5 layers of a pretrained model finetuned on mptrj + alexandria with D3 corrections |
11
-
12
-
13
-
14
-
### Matbench Discovery Results
15
-
16
-
17
-
### orb-v1: Full dataset pretraining, MPtraj + Alexandria finetuning
We provide several pretrained models that can be used to calculate energies, forces & stresses of atomic systems. All models are provided in the `orb_models.forcefield.pretrained` module.
4
+
5
+
### OrbMol Models
6
+
7
+
These models are a continuation of the `orb-v3` series, but are trained on the [Open Molecules 2025 (OMol25)](https://arxiv.org/pdf/2505.08762) dataset—over 100M high-accuracy DFT calculations (ωB97M-V/def2-TZVPD) on diverse molecular systems including metal complexes, biomolecules, and electrolytes.
8
+
9
+
There are two options:
10
+
*`orb-v3-conserative-omol`
11
+
*`orb-v3-direct-omol`
12
+
13
+
See below for more explanation of this naming convention. Both models have `inf` neighbors, ensuring a continuous PES.
14
+
15
+
### [V3 Models](https://arxiv.org/abs/2504.06231)
16
+
V3 models use the following naming convention: ```orb-v3-X-Y-Z``` where:
17
+
-`X`: Model type - `direct` or `conservative`. Conservative models compute forces and stress via backpropagation, which is a physically motivated choice that appears necessary for certain types of simulation such as NVE Molecular dynamics. Conservative models are signficantly slower and use more memory than their direct counterparts.
18
+
19
+
-`Y`: Maximum neighbors per atom: `20` or `inf`. A finite cutoff of `20` induces discontinuties in the PES, which can lead to significant inaccuracies for certain types of highly sensitive calculations (e.g. calculations involving Hessians). However, finite cutoffs reduce the amount of edge processing in the network, reducing latency and memory use.
20
+
21
+
-`Z`: Training dataset - `omat` or `mpa`. Both of these dataset consist of small bulk crystal structures. We find that models trained on such data can generalise reasonably well to non-periodic systems (organic molecules) or partially periodic systems (slabs), but caution is advised in these scenarios.
22
+
23
+
#### Features:
24
+
- Model compilation using PyTorch 2.6.0+, enabling faster inference while maintaining support for dynamic graph sizes
25
+
- Wider architecture (1024 vs 512) with fewer layers (5 vs 15) compared to v2, resulting in 2-3x faster performance with similar parameter count
26
+
- Two variants available: direct models and conservative models (forces/stress computed via backpropagation)
27
+
- Trained on the larger, more diverse OMat24 dataset
- Enhanced stability through Huber loss and a ZBL pair repulsion term added to forces
30
+
- Models available with both unlimited neighbors and 20-neighbor maximum configurations
31
+
- New confidence head providing intrinsic uncertainty estimates for predictions
32
+
33
+
#### Advice / Caveats
34
+
- Consider using `orb-v3-conservative-120-omat` for initial testing, specifying `precision='float32-highest'` when loading the model. This is the most computational expensive but accurate configuration. If this level of accuracy meets your needs, then other models and precisions can be investigated to improve speed and system-size scalability.
35
+
- We do not advise using the `-mpa` models unless they are required for compatability with benchmarks (for example, Matbench Discovery). They are generally less performant.
36
+
- Orb-v3 models are **compiled** by default and use Pytorch's dynamic batching, which means that they do not need to recompile as graph sizes change. However, the first call to the model will be slower, as the graph is compiled by torch.
37
+
38
+
### [V2 Models](https://arxiv.org/abs/2410.22570)
39
+
40
+
-`orb-v2` - trained on [MPTraj](https://figshare.com/articles/dataset/Materials_Project_Trjectory_MPtrj_Dataset/23713842?file=41619375) + [Alexandria](https://alexandria.icams.rub.de/).
41
+
-`orb-mptraj-only-v2` - trained on the MPTraj dataset only to reproduce our second Matbench Discovery result. We do not recommend using this model for general use.
42
+
-`orb-d3-v2` - trained on MPTraj + Alexandria with integrated D3 corrections. In general, we recommend using this model, particularly for systems where dispersion interactions are important. This model was trained to predict D3-corrected targets and hence is the same speed as `orb-v2`. Incorporating D3 into the model like this is substantially faster than using analytical D3 corrections.
43
+
-`orb-d3-{sm,xs}-v2` - Smaller versions of `orb-d3-v2`. The `sm` model has 10 layers, whilst the `xs` model has 5 layers.
44
+
45
+
#### Features
46
+
- v2 models use a smoothed cosine distance cutoff for the attention mechanism, ensuring a continuous PES.
47
+
- The force predictions now have net zero forces, meaning they are much more stable for MD simulations.
48
+
- The models are generally more accurate (Increase in 2-3% on the matbench discovery dataset).
49
+
50
+
### [V1 Models](https://arxiv.org/abs/2410.22570)
51
+
52
+
Our initial release. These models were state of the art performance on the matbench discovery dataset at time of release, but have since been superceeded and removed.
Copy file name to clipboardExpand all lines: README.md
+35-62Lines changed: 35 additions & 62 deletions
Display the source diff
Display the rich diff
Original file line number
Diff line number
Diff line change
@@ -27,61 +27,21 @@ Alternatively, you can use Docker to run orb-models; [see instructions below](#d
27
27
28
28
### Updates
29
29
30
-
**April 2025**: We have released the [Orb-v3 set of potentials](https://arxiv.org/abs/2504.06231). These models improve substantially over Orb-v2, in particular:
30
+
**August 2025**: Release of the OrbMol potentials (blog post forthcoming).
31
31
32
-
- Model compilation using PyTorch 2.6.0+, enabling faster inference while maintaining support for dynamic graph sizes
33
-
- Wider architecture (1024 vs 512) with fewer layers (5 vs 15) compared to v2, resulting in 2-3x faster performance with similar parameter count
34
-
- Two variants available: direct models and conservative models (forces/stress computed via backpropagation)
35
-
- Trained on the larger, more diverse OMat24 dataset
- Enhanced stability through Huber loss and a ZBL pair repulsion term added to forces
38
-
- Models available with both unlimited neighbors and 20-neighbor maximum configurations
39
-
- New confidence head providing intrinsic uncertainty estimates for predictions
32
+
* Trained on the [Open Molecules 2025 (OMol25)](https://arxiv.org/pdf/2505.08762) dataset—over 100M high-accuracy DFT calculations (ωB97M-V/def2-TZVPD) on diverse molecular systems including metal complexes, biomolecules, and electrolytes.
33
+
* Architecturally similar to the highly-performant Orb-v3 models, but now explicit total charges and spins can be passed as input.
34
+
* To get started with these models, see: [How to specify total charge and spin for OrbMol](#how-to-specify-total-charge-and-spin-for-orbmol).
40
35
36
+
**April 2025**: Release of the [Orb-v3 set of potentials](https://arxiv.org/abs/2504.06231).
41
37
42
-
**Oct 2024**: We have released the [Orb-v2 set of potentials](https://arxiv.org/abs/2410.22570). These models have two major changes:
43
-
- v2 models use a smoothed cosine distance cutoff for the attention mechanism. This is a more physically motivated cutoff that is better suited for MPNNs.
44
-
- The force predictions now have net zero forces, meaning they are much more stable for MD simulations.
45
-
- The models are generally more accurate (Increase in 2-3% on the matbench discovery dataset).
38
+
**Oct 2024**: Release of the [Orb-v2 set of potentials](https://arxiv.org/abs/2410.22570).
46
39
47
-
These models are substantially better for all use cases, so we have removed the v1 models from the new orb-models package. To load the v1 models, please install the v0.3.2 version of orb-models.
40
+
**Sept 2024**: Release of v1 models - state of the art performance on the matbench discovery dataset.
48
41
49
-
**Sept 2024**: v1 models released - state of the art performance on the matbench discovery dataset.
50
42
51
-
52
-
### Pretrained models
53
-
54
-
We provide several pretrained models that can be used to calculate energies, forces & stresses of atomic systems. All models are provided in the `orb_models.forcefield.pretrained` module.
55
-
56
-
#### V3 Models
57
-
V3 models use the following naming convention:
58
-
59
-
```orb-v3-X-Y-Z```
60
-
61
-
where:
62
-
-`X`: Model type (`direct` or `conservative`) - determines how forces/stress are computed
63
-
-`Y`: Maximum neighbors per atom (`20` or `inf`)
64
-
-`Z`: Training dataset (`omat` or `mpa`)
65
-
66
-
For example, `orb-v3-conservative-inf-omat` is a model that:
67
-
- Computes forces/stress as gradients of energy
68
-
- Has effectively infinite neighbors (120 in practice)
69
-
- Was trained on the OMat24 dataset
70
-
71
-
72
-
Orb-v3 models are **compiled** by default and use Pytorch's dynamic batching, which means that they do not need to recompile as graph sizes change. However, the first call to the model will be slower, as the graph is compiled by torch.
73
-
74
-
75
-
**We suggest using models trained on OMAT24**, as these models are more performant and the data they are trained on uses newer pseudopotentials in VASP (PBE54 vs PBE52)*. `-mpa` models should be used if compatability with benchmarks (for example, Matbench Discovery) is required.
76
-
77
-
#### V2 Models
78
-
79
-
-`orb-v2` - trained on [MPTraj](https://figshare.com/articles/dataset/Materials_Project_Trjectory_MPtrj_Dataset/23713842?file=41619375) + [Alexandria](https://alexandria.icams.rub.de/).
80
-
-`orb-mptraj-only-v2` - trained on the MPTraj dataset only to reproduce our second Matbench Discovery result. We do not recommend using this model for general use.
81
-
-`orb-d3-v2` - trained on MPTraj + Alexandria with integrated D3 corrections. In general, we recommend using this model, particularly for systems where dispersion interactions are important. This model was trained to predict D3-corrected targets and hence is the same speed as `orb-v2`. Incorporating D3 into the model like this is substantially faster than using analytical D3 corrections.
82
-
-`orb-d3-{sm,xs}-v2` - Smaller versions of `orb-d3-v2`. The `sm` model has 10 layers, whilst the `xs` model has 5 layers.
83
-
84
-
For more information on the models, please see the [MODELS.md](MODELS.md) file.
43
+
### Available models
44
+
See [MODELS.md](MODELS.md) for a full list of available models along with guidance.
Or you can use it to run MD simulations. The script, an example input xyz file and a Colab notebook demonstration are available in the [examples directory.](./examples) This should work with any input, simply modify the input_file and cell_size parameters. We recommend using constant volume simulations.
162
120
121
+
#### How to specify total charge and spin for OrbMol
122
+
123
+
The OrbMol models *require* total charge and spin to be specified. This can be done by setting them in `atoms.info` dictionary.
124
+
125
+
```python
126
+
import ase
127
+
from ase.build import molecule
128
+
from orb_models.forcefield import atomic_system, pretrained
129
+
from orb_models.forcefield.base import batch_graphs
130
+
131
+
device ="cpu"# or device="cuda"
132
+
orbff = pretrained.orb_v3_conservative_omol(
133
+
device=device,
134
+
precision="float32-high", # or "float32-highest" / "float64
After the model is finetuned, checkpoints will, by default, be saved to the ckpts folder in the directory you ran the finetuning script from.
199
+
Where base_model is an element of `orb_models.forcefield.pretrained.ORB_PRETRAINED_MODELS.keys()`.
227
200
228
-
You can use the new model and load the checkpoint by:
201
+
After the model is finetuned, checkpoints will, by default, be saved to the ckpts folder in the directory you ran the finetuning script from. You can use the new model and load the checkpoint by:
0 commit comments