Garden's Almanac of Matter Models OrbMol-v2
May 2026
Paper acs.digitellinc.com β
Code GitHub β
Weights Orbital-Materials/orbmol-v2
License Apache-2.0 β
| Checkpoint | Params | Polaris | Sophia | Perlmutter | Delta | Frontier | Della |
|---|---|---|---|---|---|---|---|
| orbmol-v2 | β | β | β | β | β | β | |
| orbmol:custom | β | β | β | β | β | β |
β verified, last 30 days β installed, not recently verified β not installed
No cluster has installed any checkpoint of this model yet.
Running this model
1# from a job or interactive session on a supported cluster:
2from rootstock import RootstockCalculator
3
4# this checkpoint accepts charge and spin
5atoms.info["charge"] = -1
6atoms.info["spin"] = 2
7
8with RootstockCalculator(
9 cluster=YOUR_CLUSTER_ID, # eg, "frontier"
10 checkpoint="orbmol-v2",
11 device="cuda",
12) as calc:
13 # model is now running in subprocess on compute node
14 atoms.calc = calc
15 atoms.get_potential_energy()
Environments
Rootstock runs each model family inside an isolated Python environment defined by a single file. This environment file includes the specific dependencies needed, plus a setup() function that loads the model and returns an ASE calculator. These files are usually almost identical for a given model family, but because of cluster-specific quirks (eg, an old CUDA driver) the dependencies and setup code can vary a bit.
orbmol_env.py
1# /// script
2# requires-python = ">=3.12,<3.13"
3# dependencies = [
4# # 0.7.0 (2026-05-26) introduces orbmol_v2 and the orbmol-v1-* aliases
5# # (orbmol-v1-conservative == orb-v3-conservative-omol). Upstream still
6# # pins dm-tree==0.1.8 (no cp313 wheel), hence the <3.13 cap.
7# "orb-models>=0.7,<0.8",
8# "ase>=3.25",
9# "torch>=2.8,<3.0",
10# "cached_path==1.8.10",
11# # torch's ROCm wheels depend on this; it lives only on the ROCm
12# # index, so it must be a direct dep for [tool.uv.sources] to route it.
13# "pytorch-triton-rocm",
14# ]
15#
16# [tool.uv.sources]
17# torch = { index = "pytorch-rocm" }
18# pytorch-triton-rocm = { index = "pytorch-rocm" }
19#
20# [[tool.uv.index]]
21# name = "pytorch-rocm"
22# url = "https://download.pytorch.org/whl/rocm6.4"
23# explicit = true
24# ///
25"""OrbMol env (ROCm) β Orbital Materials' molecular potentials (OMol25/OPoly26).
26
27orbmol-v1-conservative is upstream's alias for orb-v3-conservative-omol;
28orbmol-v2 (2026-05) adds learnable long-range electrostatics (CoulombModule).
29
30ROCm caveats (orb-models 0.7 hard-depends on nvalchemi-toolkit-ops, NVIDIA
31Warp/CUDA kernels β pure-python wheel, so it *installs* fine here):
32
33- Neighbor lists: the 0.7 default edge_method is knn_alchemi (Warp/CUDA);
34 torch ROCm reports device.type == "cuda", so the default would try to
35 launch CUDA kernels on an MI250X. knn_scipy is rejected upstream off-CPU,
36 so we pin edge_method="knn_brute_force" (pure torch cdist/topk, runs on
37 the device; O(N^2) but these are molecular systems) β override via setup
38 kwarg if a ROCm-capable path appears. Upstream deprecates it in favor of
39 knn_alchemi, which is not an option here.
40- torch.compile defaults off (compile=False): the triton-ROCm inductor
41 path is unvalidated on MI250X; pass compile=True to opt in.
42- orbmol-v2 electrostatics: non-periodic systems use a pure-torch direct
43 Coulomb sum (fine on ROCm). Periodic systems go through nvalchemiops
44 Particle Mesh Ewald β expect failure on Frontier. Molecules only.
45
46Charge/spin: these are OMol-style conditioned models β ORBCalculator raises if
47atoms.info lacks "charge"/"spin" (spin = multiplicity). Set them per structure
48via atoms.info; absent that, we default to neutral singlet (charge=0, spin=1).
49"""
50
51import os
52import shutil
53import urllib.request
54from pathlib import Path
55
56CHECKPOINTS = {
57 "orbmol-v1-conservative": "orbmol-v1-conservative",
58 "orbmol-v2": "orbmol-v2",
59 # Your own fine-tuned OrbMol weights: pair with weights= (loaded via
60 # setup_from_path); pass arch= to pick the base architecture.
61 "orbmol:custom": None,
62}
63
64
65def _default_weights_url(load_fn) -> str:
66 """The upstream URL baked into the loader's ``weights_path`` default."""
67 import inspect
68
69 default = inspect.signature(load_fn).parameters["weights_path"].default
70 if not isinstance(default, str) or not default.startswith(("http://", "https://")):
71 raise RuntimeError(
72 f"{load_fn.__name__} has no URL default for weights_path "
73 f"(got {default!r}); update this env file for the installed orb-models"
74 )
75 return default
76
77
78def _local_weights_path(url: str) -> Path:
79 """Where the checkpoint lives in the shared model cache."""
80 cache = Path(os.environ.get("XDG_CACHE_HOME") or Path.home() / ".cache")
81 return cache / "orb" / os.path.basename(url)
82
83
84def _fetch(url: str, dest: Path) -> None:
85 """Download ``url`` to ``dest`` atomically (tmp file + rename)."""
86 dest.parent.mkdir(parents=True, exist_ok=True)
87 tmp = dest.with_name(f"{dest.name}.tmp.{os.getpid()}")
88 try:
89 with urllib.request.urlopen(url) as resp, open(tmp, "wb") as out:
90 shutil.copyfileobj(resp, out)
91 os.replace(tmp, dest)
92 finally:
93 tmp.unlink(missing_ok=True)
94
95
96def _make_calculator(orbff, atoms_adapter, device, edge_method, **kwargs):
97 from ase.calculators.calculator import all_changes
98 from orb_models.forcefield.inference.calculator import ORBCalculator
99
100 class OrbMolCalculator(ORBCalculator):
101 """ORBCalculator that defaults missing charge/spin to neutral singlet."""
102
103 def calculate(self, atoms=None, properties=None, system_changes=all_changes):
104 if atoms is not None:
105 atoms.info.setdefault("charge", 0)
106 atoms.info.setdefault("spin", 1)
107 super().calculate(atoms, properties, system_changes)
108
109 return OrbMolCalculator(
110 orbff, atoms_adapter, device=device, edge_method=edge_method, **kwargs
111 )
112
113
114def setup(
115 checkpoint: str,
116 device: str = "cuda",
117 precision: str = "float32-high",
118 compile: bool = False,
119 edge_method: str = "knn_brute_force",
120 **kwargs,
121):
122 # Extra **kwargs go to ORBCalculator (e.g. max_num_neighbors=,
123 # half_supercell=); precision/compile go to the checkpoint loader.
124 import torch
125 from orb_models.forcefield import pretrained
126
127 fn_name = CHECKPOINTS[checkpoint].replace("-", "_")
128 load_fn = getattr(pretrained, fn_name)
129
130 # orb-models resolves its default weights URL through `cached_path`, which
131 # write-locks its cache dir even on warm hits β EACCES for anyone who can't
132 # write the shared install. Handed a *local* path instead, cached_path
133 # returns it without locking. So the weights are pre-fetched into the
134 # shared model cache at `rootstock add` time (maintainer, cache writable)
135 # and every later serve loads that file.
136 url = _default_weights_url(load_fn)
137 weights = _local_weights_path(url)
138 if not weights.exists():
139 _fetch(url, weights)
140
141 orbff, atoms_adapter = load_fn(
142 weights_path=str(weights),
143 device=torch.device(device),
144 precision=precision,
145 compile=compile,
146 )
147 return _make_calculator(
148 orbff, atoms_adapter, torch.device(device), edge_method, **kwargs
149 )
150
151
152def setup_from_path(
153 path: str,
154 device: str = "cuda",
155 arch: str = "orbmol-v2",
156 precision: str = "float32-high",
157 compile: bool = False,
158 edge_method: str = "knn_brute_force",
159 **kwargs,
160):
161 # Custom checkpoints (`:custom` ids with user weights). A weights file
162 # doesn't say which architecture produced it, so `arch` names the
163 # pretrained loader to instantiate β pass the right one at call time
164 # (setup_kwargs={"arch": ...} / --kwarg arch=...). Handing the loader a
165 # local path also means no network and no cached_path locking (see setup()).
166 import torch
167 from orb_models.forcefield import pretrained
168
169 fn_name = arch.replace("-", "_").replace(":", "_")
170 try:
171 load_fn = getattr(pretrained, fn_name)
172 except AttributeError:
173 raise ValueError(
174 f"unknown OrbMol architecture {arch!r}; expected a loader name from "
175 f"orb_models.forcefield.pretrained, e.g. orbmol-v2, orbmol-v1-conservative"
176 ) from None
177
178 orbff, atoms_adapter = load_fn(
179 weights_path=path,
180 device=torch.device(device),
181 precision=precision,
182 compile=compile,
183 )
184 return _make_calculator(
185 orbff, atoms_adapter, torch.device(device), edge_method, **kwargs
186 )
187
Built on Frontier: 2026-08-12
Couldn't load the current environments from Rootstock.