orb-models

Orb-v3 models follow the naming pattern orb-v3-X-Y-Z. Choose your configuration based on the following components:

X: Model Type

• direct: Computes forces and stress directly. Faster and uses less memory.
• conservative: Computes forces and stress via backpropagation. Physically motivated for simulations like NVE Molecular Dynamics, but significantly slower and more memory-intensive.

Y: Maximum Neighbors per Atom

• 20: Finite cutoff. Reduces latency and memory usage but induces discontinuities in the Potential Energy Surface (PES), which can impact sensitive calculations like Hessians.
• inf: Unlimited neighbors. Ensures a continuous PES.

Z: Training Dataset

• omat: Trained on the OMat24 dataset (bulk crystals). Generally generalizes well to organic molecules or slabs.
• mpa: Trained on the MPA dataset. Generally less performant; only use if required for specific benchmarks like Matbench Discovery.

The finetune.py script requires training data to be stored in an ASE SQLite database format. Each entry in the database must be an ASE Atoms object with specific properties attached via a SinglePointCalculator.

Required Properties

• Positions: Atomic positions (stored in the Atoms object).
• Atomic numbers: Element types (stored in the Atoms object).
• Cell: Unit cell vectors (for periodic systems).
• Energy: Total energy in eV.
• Forces: Forces on each atom in eV/Å, shape (n_atoms, 3).
• Stress (optional): Stress tensor in Voigt notation (6 components), in eV/ų.

Stress Tensor Format

If providing stress, use Voigt notation [σ_xx, σ_yy, σ_zz, σ_yz, σ_xz, σ_xy]. To convert a 3x3 tensor:

stress_voigt = [
    stress_3x3[0, 0],  # σ_xx
    stress_3x3[1, 1],  # σ_yy
    stress_3x3[2, 2],  # σ_zz
    stress_3x3[1, 2],  # σ_yz
    stress_3x3[0, 2],  # σ_xz
    stress_3x3[0, 1],  # σ_xy
]

import ase
import ase.db
from ase import Atoms
from ase.calculators.singlepoint import SinglePointCalculator

# Create a database file
db = ase.db.connect('my_training_data.db')

for structure in my_structures:
    atoms = Atoms(
        symbols=structure['symbols'],
        positions=structure['positions'],
        cell=structure['cell'],
        pbc=True
    )
    
    calc = SinglePointCalculator(
        atoms=atoms,
        energy=structure['energy'],
        forces=structure['forces'],
        stress=structure['stress']
    )
    atoms.calc = calc
    db.write(atoms)

For high-performance batched optimization and MD simulations, use the OrbTorchSimModel with TorchSim. This requires the torch-sim-atomistic package.

Workflow:

1. Convert a list of ASE atoms to a TorchSim state using ts.io.atoms_to_state().
2. Initialize OrbTorchSimModel(orbff, atoms_adapter).
3. Use ts.optimize() for geometry relaxation or run the model directly on the state for MD.

import ase
import torch
import torch_sim as ts
from ase.build import bulk
from orb_models.forcefield import pretrained
from orb_models.forcefield.inference.orb_torchsim import OrbTorchSimModel

device = "cpu"
orbff, atoms_adapter = pretrained.orb_v3_conservative_inf_omat(
  device=device,
  precision="float32-high",
)

atoms1 = bulk('Cu', 'fcc', a=3.58, cubic=True)
atoms2 = bulk('Si', 'diamond', a=5.43, cubic=True)
atoms_list = [atoms1, atoms2]

# Convert to TorchSim state
ts_state = ts.io.atoms_to_state(atoms_list, device, dtype=torch.get_default_dtype())

ts_model = OrbTorchSimModel(orbff, atoms_adapter)

# Optimize
relaxed_state = ts.optimize(
    system=ts_state,
    convergence_fn=ts.generate_force_convergence_fn(force_tol=0.01, include_cell_forces=False),
    model=ts_model,
    optimizer=ts.Optimizer["fire"],
    max_steps=100,
    steps_between_swaps=10,
)
results = ts_model(relaxed_state)
print(results["energy"])

You can use orb-models directly by loading a pretrained model and an atoms adapter. This workflow involves converting ASE atoms to a graph, predicting properties (energy, forces, stress), and optionally converting the results back to ASE atoms.

Key steps:

1. Load model and adapter via pretrained.<model_name>().
2. Convert ASE atoms to a graph using atoms_adapter.from_ase_atoms().
3. Predict using orbff.predict(graph).
4. Convert back to ASE using graph.to_ase_atoms().

import ase
from ase.build import bulk
from orb_models.forcefield import pretrained

device = "cpu"  # or device="cuda"
orbff, atoms_adapter = pretrained.orb_v3_conservative_inf_omat(
  device=device,
  precision="float32-high",   # or "float32-highest" / "float64"
)
atoms = bulk('Cu', 'fcc', a=3.58, cubic=True)
graph = atoms_adapter.from_ase_atoms(atoms, device=device)

result = orbff.predict(graph, split=False)

# Convert to ASE atoms
atoms = graph.to_ase_atoms(
    energy=result["energy"],
    forces=result["forces"],
    stress=result["stress"]
)