Orb

Orb: Orbital Materials' machine-learned potentials, through orb-models.

The codec calls Orb's models directly in the calling process (model.predict on the graph Orb's adapter builds), as Orb's ASE calculator does internally; install it next to qccompute with pip install orb-models. orb-models and fairchem-core need separate environments. Weights download on first use, and each model is loaded once per process and reused.

Inputs map to Orb as follows:

• model.method is a name from orb_models.forcefield.pretrained .ORB_PRETRAINED_MODELS: orb-v3-conservative-omol (trained on OMol25, smooth conservative forces; the usual choice for molecules), orb-v3-direct-omol (faster, direct non-conservative forces), orbmol-v2, or the materials models (orb-v3-conservative-inf-omat, ...). Models have no basis set.
• Models trained on OMol25 get structure.charge and structure.multiplicity through atoms.info["charge"] and atoms.info["spin"]. The others describe neutral closed-shell systems only, so the codec raises for anything else.
• keywords: device ("cpu", "cuda"; CUDA when available by default) and precision ("float32-high" by default, "float32-highest", or "float64"). Nothing else is accepted. A model on a GPU declares one for qccompute, which gives each calculation of a list its own; an explicit index such as "cuda:1" opts out of that.
• cmdline_args and files are not accepted.

Calculation types: energy and gradient. Orb provides no Hessian; use geomeTRIC with Orb gradients for optimizations and transition states. Orb predicts energies in eV; the codec converts them, and the forces, to Hartree and bohr.

Orb does not move the molecule, so gradients are in the input's frame. No input field carries atom indices.

Energies and gradients of several structures can be computed as one batch: qccompute.compute does this for a list of inputs that share a model, keywords, and calculation type, evaluating them in one pass instead of one call each. Hessians are computed one structure at a time.

Examples

Each script runs as written. Copy it, delete the options you don't need, and run it with python script.py. They live in the repository's examples/ folder.

Single point: energy, gradient

examples/orb/single_point.py

"""Orb single point: energy or gradient from an Orb potential.

Requires: ``pip install qccompute orb-models``, in an environment without
``fairchem-core``. Orb is a Python library, so nothing has to be on PATH: the model
runs inside this Python process, on the GPU when PyTorch sees one. Weights download on
first use; later calculations in the same process reuse the loaded model.

How the codec talks to Orb (https://github.com/orbital-materials/orb-models):

- ``model.method`` is a name from ``orb_models.forcefield.pretrained
  .ORB_PRETRAINED_MODELS``: "orb-v3-conservative-omol" (OMol25, conservative forces;
  the usual choice for molecules), "orb-v3-direct-omol" (faster, non-conservative
  forces), "orbmol-v2", or the materials models. There is no basis set.
- OMol25 models read charge and multiplicity from ``structure``; the others describe
  neutral closed-shell systems only.
- ``precision`` is "float32-high" by default, or "float32-highest" or "float64".
- Orb has no Hessian. Use geomeTRIC with Orb gradients for optimizations and
  transition states (examples/geometric/).
- Energies come back in Hartree and gradients in Hartree/bohr, converted from Orb's
  eV and angstrom.
"""

from qccompute import compute
from qcdata import CalcType, Model, ProgramInput, Structure

calctype = CalcType.gradient  # or CalcType.energy

# Coordinates are in bohr (Structure.from_xyz / Structure.open read angstrom files).
structure = Structure(
    symbols=["O", "H", "H"],
    geometry=[[0.0, 0.0, 0.0], [0.524, 1.687, 0.48], [1.146, -0.45, -1.354]],
    charge=0,
    multiplicity=1,
)

prog_input = ProgramInput(
    program="orb",
    calctype=calctype,
    structure=structure,
    model=Model(method="orb-v3-conservative-omol"),  # or "orb-v3-direct-omol", ...
    keywords={
        "device": "cpu",  # or "cuda"; leave out to use a GPU when present
        # "precision": "float64",  # or "float32-highest"
    },
)

output = compute(prog_input)

results = output.results
print("Energy (Hartree):", results.energy)
if results.gradient is not None:
    print("Gradient (Hartree/bohr):\n", results.gradient)
print("orb-models version:", results.provenance.program_version)

Many structures at once: batches

examples/orb/many_calculations.py

"""Orb for many structures at once: a bond scan evaluated in batches.

Requires: what examples/orb/single_point.py requires, and qccompute 0.21.0 or newer.

How batching works:

- Pass ``compute`` a list of inputs. Energy and gradient inputs with the same program,
  model, and keywords are evaluated in batches, one model call for up to
  ``get_codec("orb").max_batch`` (32) structures, in this Python process. On a GPU
  that is many times faster than one call per structure.
- Each structure's results are what computing it alone gives, to the model's precision.
  Structures in one batch may differ in size, composition, and charge.
- Inputs with different models or keywords go in separate batches. Hessians are
  computed one structure at a time.
- A batch that runs out of GPU memory is split in half until it fits. A batch that
  fails otherwise is computed again one structure at a time, so each error belongs to
  its own input.
- The model stays loaded in this process, so the next ``compute`` call (the next step
  of a path optimization, say) does not load it again.
- ``output.execution.extras`` records ``batch_size`` and the structure's
  ``batch_index``; ``logs`` and ``wall_time`` are the whole batch's.
- Orb compiles its model with ``torch.compile`` (orb-models' default). The first calls
  at new input sizes take seconds while PyTorch optimizes for them; after a few sizes
  it stops recompiling and calls are fast. So the one-at-a-time timing below includes
  that setup for single structures, and a script making only a handful of Orb calls
  pays it without benefiting. Optimizations, scans, and paths make many calls and
  come out well ahead.
"""

import math
import time

from qccompute import compute
from qcdata import CalcType, Model, ProgramInput, Structure


def water(r: float) -> Structure:
    """Water with one O-H bond stretched to r bohr, the other at 1.81 bohr."""
    angle = math.radians(104.5)
    return Structure(
        symbols=["O", "H", "H"],
        geometry=[
            [0.0, 0.0, 0.0],
            [r, 0.0, 0.0],
            [1.81 * math.cos(angle), 1.81 * math.sin(angle), 0.0],
        ],
    )


distances = [1.4 + 0.1 * i for i in range(20)]  # bohr
inputs = [
    ProgramInput(
        program="orb",
        calctype=CalcType.gradient,  # or CalcType.energy
        structure=water(r),
        model=Model(method="orb-v3-conservative-omol"),
        keywords={
            # "precision": "float64",  # GPU default float32-high differs in the last digits
            # "device": "cpu",  # or "cuda"; leave out to use a GPU when present
        },
    )
    for r in distances
]

compute(inputs)  # load the model and warm up, so the timings compare evaluation only

start = time.perf_counter()
outputs = compute(inputs)  # one batch of 20
batched = time.perf_counter() - start

start = time.perf_counter()
alone = [compute(inp) for inp in inputs]  # 20 separate calls, for comparison
separate = time.perf_counter() - start

print("r (bohr)  energy (Hartree)  batch_index/batch_size")
for r, output in zip(distances, outputs, strict=True):
    extras = output.execution.extras
    print(
        f"{r:8.2f}  {output.results.energy:16.8f}  "
        f"{extras['batch_index']}/{extras['batch_size']}"
    )
largest = max(
    abs(a.results.energy - b.results.energy)
    for a, b in zip(alone, outputs, strict=True)
)
print(f"Batched: {batched:.3f} s; one at a time: {separate:.3f} s")
print(f"Largest energy difference from computing alone: {largest:.1e} Hartree")