Ported changes from the internal repo

Author: vsimkus

MODELS.md

@@ -7,8 +7,8 @@ We provide several pretrained models that can be used to calculate energies, for
 These models are a continuation of the [`orb-v3`](#v3-models) series 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. Note: The training data does not contain periodic systems and these models have not been carefully tested on periodic systems. 
 
 There are two options:
-* `orb-v3-conservative-omol`
-* `orb-v3-direct-omol`
+* `orbmol-v1-conservative`
+* `orbmol-v1-direct`
 
 See below for more explanation of this naming convention. Both models have `inf` neighbors, ensuring a continuous PES.
 
@@ -18,7 +18,7 @@ See below for more explanation of this naming convention. Both models have `inf`
 
 OrbMol-v2 extends the OrbMol architecture with **learnable per-atom electrostatics**: a `LatentChargeHead` predicts per-atom partial charges (constrained to sum to the system total charge), a `LatentSpinHead` predicts per-atom spins (constrained to sum to 2S = `spin_multiplicity − 1`), and a `CoulombModule` adds a long-range Coulomb energy on top of the GNN — direct bare-1/r Coulomb sum for non-periodic systems, Particle Mesh Ewald via `nvalchemiops` for periodic systems. The energy head (`ChargeConditionedEnergyHead`) is conditioned on the predicted charges and spins per-atom.
 
-Trained on OMol25 (ωB97M-V/def2-TZVPD); supports both periodic and non-periodic systems. Stress is enabled via `model.enable_stress()` if needed.
+Trained on OMol25 and OPoly26 (ωB97M-V/def2-TZVPD); supports both periodic and non-periodic systems. Stress is enabled via `model.enable_stress()` if needed.
 
 ```python
 from orb_models.forcefield.pretrained import orbmol_v2

README.md

@@ -21,17 +21,18 @@ Alternatively, you can use Docker to run orb-models; [see instructions below](#d
 
 ### Updates
 
+**May 2026**: Release of OrbMol-v2 — extends the OrbMol architecture with learnable per-atom electrostatics:
+
+* New `LatentChargeHead` and `LatentSpinHead` predict per-atom charges and spins (constrained to sum to the system total charge / 2S = `spin_multiplicity − 1`), and a `CoulombModule` adds long-range Coulomb energy on top of the GNN — direct Coulomb sum for non-periodic systems, Particle Mesh Ewald via `nvalchemiops` for periodic.
+* The energy head (`ChargeConditionedEnergyHead`) is conditioned on the predicted charges and spins per atom.
+* Trained on OMol25 and OPoly26 (ωB97M-V/def2-TZVPD); load with `pretrained.orbmol_v2(device="cuda")`.
+
 **February 2026**: Improved GPU-accelerated graph construction with [ALCHEMI Toolkit-Ops](https://github.com/NVIDIA/nvalchemi-toolkit-ops) and batched simulation with [TorchSim](https://github.com/TorchSim/torch-sim):
 
 * Alchemi-based graph construction (GPU-accelerated, up to 12x faster for large single systems, and sub-linear batch scaling delivering >100x graph construction speed-up for large batches of small systems)
 * TorchSim wrapper for batched optimisation and simulation, see [usage with TorchSim](#usage-with-torchsim)
 * Alchemi-based D3 dispersion correction module, see [D3 correction](#d3-correction)
 
-**May 2026**: Release of OrbMol-v2 — extends the OrbMol architecture with learnable per-atom electrostatics:
-
-* New `LatentChargeHead` and `LatentSpinHead` predict per-atom charges and spins (constrained to sum to the system total charge / 2S = `spin_multiplicity − 1`), and a `CoulombModule` adds long-range Coulomb energy on top of the GNN — bare 1/r direct sum for non-periodic systems, Particle Mesh Ewald via `nvalchemiops` for periodic.
-* The energy head (`ChargeConditionedEnergyHead`) is conditioned on the predicted charges and spins per atom.
-* Trained on OMol25 (ωB97M-V/def2-TZVPD); load with `pretrained.orbmol_v2(device="cuda")`.
 
 **August 2025**: Release of the [OrbMol potentials](https://www.orbitalindustries.com/posts/orbmol-extending-orb-to-molecular-systems):
 
@@ -195,7 +196,7 @@ from ase.build import molecule
 from orb_models.forcefield import pretrained
 
 device = "cpu"  # or device="cuda"
-orbff, atoms_adapter = pretrained.orb_v3_conservative_omol(
+orbff, atoms_adapter = pretrained.orbmol_v1_conservative(
   device=device,
   precision="float32-high",   # or "float32-highest" / "float64
 )

orb_models/common/models/base.py

@@ -36,6 +36,9 @@ def loss(self, batch: T) -> ModelOutput:
         """Encodes to latents before message passing."""
         raise NotImplementedError()
 
+    def prepare_for_inference(self) -> None:
+        """Hook called before inference. Override to enable inference-only features."""
+
 
 class RegressorModelMixin[T: AbstractAtomBatch](ModelMixin[T]):
     """Model Mixin for our regression models."""

orb_models/forcefield/models/conservative_regressor.py

@@ -40,7 +40,8 @@ class ConservativeForcefieldRegressor(base.RegressorModelMixin[AtomGraphs]):
                 - "grad_forces"
                 - "grad_stress"
             which weight the gradient based losses of forces/stress respectively.
-        distill_direct_heads: Whether to distill the direct heads into the conservative heads.
+        coulomb_module: Optional CoulombModule for long-range electrostatics.
+            When present, a latent_charges head must also be in heads.
         **kwargs: Additional kwargs, used for backwards compatibility of deprecated arguments.
     """
 
@@ -55,12 +56,10 @@ def __init__(
         heads: Mapping[str, ForcefieldHead | ConfidenceHead],
         model: MoleculeGNS,
         loss_weights: dict[str, float] | None = None,
-        distill_direct_heads: bool = False,
         online_normalisation: bool = True,
         level_of_theory: str | None = None,
         forces_loss_type: Literal["mae", "mse", "huber_0.01", "condhuber_0.01"] = "condhuber_0.01",
         pair_repulsion: bool = False,
-        pair_repulsion_node_aggregation: str = "mean",
         has_stress: bool = True,
         coulomb_module: CoulombModule | None = None,
         **kwargs,
@@ -85,7 +84,6 @@ def __init__(
         _validate_heads_and_loss_weights(heads, nongrad_loss_weights)
 
         self.loss_weights = loss_weights
-        self.distill_direct_heads = distill_direct_heads
         self.forces_loss_type = forces_loss_type
 
         self.model = model
@@ -95,11 +93,7 @@ def __init__(
 
         self.pair_repulsion = pair_repulsion
         if self.pair_repulsion:
-            self.pair_repulsion_fn = ZBLBasis(
-                p=6,
-                compute_gradients=False,
-                node_aggregation=pair_repulsion_node_aggregation,
-            )
+            self.pair_repulsion_fn = ZBLBasis(p=6, compute_gradients=False, node_aggregation="sum")
 
         self.coulomb_module = coulomb_module
         if self.coulomb_module is not None:
@@ -115,11 +109,21 @@ def __init__(
         self.forces_target = PROPERTIES[self.forces_name]
         self.grad_forces_name = f"{self.grad_prefix}_{self.forces_name}"
 
-        # Stress names are always derived (from level_of_theory); has_stress toggles computation
-        self.stress_name: str = f"stress-{level_of_theory}" if level_of_theory else "stress"
-        self.stress_target: PropertyDefinition = PROPERTIES[self.stress_name]
-        self.grad_stress_name: str = f"{self.grad_prefix}_{self.stress_name}"
+        # Stress is optional since only periodic systems have it
         self.has_stress = has_stress
+        if self.has_stress:
+            self.stress_name: str | None = (
+                f"stress-{level_of_theory}" if level_of_theory else "stress"
+            )
+            self.stress_target: PropertyDefinition | None = PROPERTIES[self.stress_name]
+            self.grad_stress_name: str | None = f"{self.grad_prefix}_{self.stress_name}"
+        else:
+            self.stress_name = None
+            self.stress_target = None
+            self.grad_stress_name = None
+        assert self.has_stress == (self.grad_stress_name is not None), (
+            "grad_stress_name must be set if has_stress is True"
+        )
 
         self.grad_rotation_name = "rotational_grad"
 
@@ -129,8 +133,17 @@ def __init__(
                 self.extra_properties.append(heads[name].target.fullname)
 
     def enable_stress(self) -> None:
-        """Enable stress computation."""
+        """Enable stress computation. No-op if already enabled."""
+        if self.has_stress:
+            return
         self.has_stress = True
+        self.stress_name = "stress"
+        self.stress_target = PROPERTIES["stress"]
+        self.grad_stress_name = f"{self.grad_prefix}_{self.stress_name}"
+
+    def prepare_for_inference(self) -> None:
+        """Enable stress for inference — always available via autograd."""
+        self.enable_stress()
 
     def disable_stress(self) -> None:
         """Disable stress computation."""
@@ -146,11 +159,11 @@ def properties(self):
             self.grad_rotation_name,
         ]
         if self.has_stress:
+            assert self.grad_stress_name is not None, (
+                "grad_stress_name must be set if has_stress is True"
+            )
             props.append(self.grad_stress_name)
-        for name in self.extra_properties:
-            if not self.has_stress and "stress" in name:
-                continue
-            props.append(name)
+        props.extend(self.extra_properties)
         return props
 
     def forward(self, batch: AtomGraphs) -> dict[str, torch.Tensor]:
@@ -167,54 +180,43 @@ def forward(self, batch: AtomGraphs) -> dict[str, torch.Tensor]:
         node_features = out["node_features"]
 
         # Predict per-atom charges/spins BEFORE energy head so they can
-        # be used as conditioning features (ChargeConditionedEnergyHead) and by CoulombModule.
+        # be used as conditioning features in ChargeConditionedEnergyHead and CoulombModule.
         latent_charges = None
         if "latent_charges" in self.heads:
             latent_charges = self.heads["latent_charges"](node_features, batch)
+
         latent_spins = None
         if "latent_spins" in self.heads:
             latent_spins = self.heads["latent_spins"](node_features, batch)
 
         energy_head = self.heads[self.energy_name]
-        is_charge_conditioned = isinstance(energy_head, ChargeConditionedEnergyHead)
-        if is_charge_conditioned:
-            # ChargeConditionedEnergyHead.forward returns physical interaction energy directly.
-            assert latent_charges is not None, (
-                "ChargeConditionedEnergyHead requires a 'latent_charges' head"
-            )
-            raw_energy = energy_head(
+        energy_head = cast(ForcefieldHead, energy_head)
+        if isinstance(energy_head, ChargeConditionedEnergyHead):
+            interaction_energy = energy_head(
                 node_features,
                 batch,
                 per_atom_charges=latent_charges,
                 per_atom_spins=latent_spins,
             )
-            if self.pair_repulsion:
-                raw_energy = raw_energy + self.pair_repulsion_fn(batch)["energy"]
         else:
-            energy_head = cast(ForcefieldHead, energy_head)
-            base_energy = energy_head(node_features, batch)
-            raw_energy = energy_head.denormalize(base_energy, batch)
-            if self.pair_repulsion:
-                raw_energy = raw_energy + self.pair_repulsion_fn(batch)["energy"]
+            assert latent_spins is None, "Latent spins are predicted but not used."
+            interaction_energy = energy_head(node_features, batch)
+        if self.pair_repulsion:
+            interaction_energy += self.pair_repulsion_fn(batch)["energy"]
 
-        # Long-range Coulomb (predicted charges only — spins not used).
         coulomb_explicit_forces = None
         coulomb_explicit_virial = None
         if self.coulomb_module is not None:
             assert latent_charges is not None, "CoulombModule requires a LatentChargeHead"
             coulomb_energy, coulomb_explicit_forces, coulomb_explicit_virial = self.coulomb_module(
                 latent_charges, batch
             )
-            raw_energy = raw_energy + coulomb_energy
+            interaction_energy += coulomb_energy
 
-        # Store final energy in `out` (interaction-units for ChargeConditioned, normalized otherwise).
-        if is_charge_conditioned:
-            out[self.energy_name] = raw_energy
-        else:
-            out[self.energy_name] = energy_head.normalize(raw_energy, batch, online=False)
+        out[self.energy_name] = interaction_energy
 
         forces, stress, rotational_grad = compute_gradient_forces_and_stress(
-            energy=raw_energy,
+            energy=interaction_energy,
             positions=batch.node_features["positions"],
             displacement=batch.system_features["stress_displacement"],
             cell=batch.system_features["cell"],
@@ -225,11 +227,11 @@ def forward(self, batch: AtomGraphs) -> dict[str, torch.Tensor]:
 
         # Add explicit/spatial Coulomb force/stress corrections (see CoulombModule docstring).
         if self.coulomb_module is not None:
-            assert coulomb_explicit_forces is not None
-            assert coulomb_explicit_virial is not None
+            assert coulomb_explicit_forces is not None, "Explicit/spatial forces are not computed"
+            assert coulomb_explicit_virial is not None, "Explicit/spatial virial is not computed"
             forces = forces + coulomb_explicit_forces
             if self.has_stress:
-                assert stress is not None
+                assert stress is not None, "has_stress is True but stress is None"
                 cell_3d = batch.system_features["cell"].view(-1, 3, 3)
                 volume = torch.linalg.det(cell_3d).abs()
                 coulomb_stress_3x3 = -coulomb_explicit_virial / volume.view(-1, 1, 1)
@@ -250,26 +252,40 @@ def forward(self, batch: AtomGraphs) -> dict[str, torch.Tensor]:
 
         return out
 
-    def predict(self, batch: AtomGraphs, split: bool = False) -> dict[str, torch.Tensor]:
-        """Predict energy, forces, and stress."""
+    def predict(
+        self,
+        batch: AtomGraphs,
+        split: bool = False,
+        fp64_energy: bool = True,
+    ) -> dict[str, torch.Tensor]:
+        """Predict energy, forces, and stress.
+
+        Args:
+            batch: Input batch.
+            split: If True, split predictions per graph.
+            fp64_energy: If True (default), return absolute energy in fp64;
+                required to preserve kJ/mol resolution since reference
+                energies can be as high as ~1e4-1e5 eV. If False, returns
+                energy in the input dtype.
+        """
         preds = self(batch)
 
         out = {}
-        energy_head = self.heads[self.energy_name]
-        if isinstance(energy_head, ChargeConditionedEnergyHead):
-            # preds[energy_name] is interaction energy in physical units; add reference in fp64.
-            out[self.energy_name] = energy_head.absolute_energy(preds[self.energy_name], batch)
-        else:
-            energy_head = cast(ForcefieldHead, energy_head)
-            out[self.energy_name] = energy_head.denormalize(preds[self.energy_name], batch)
+        energy_head = cast(EnergyHead, self.heads[self.energy_name])
+        out[self.energy_name] = energy_head.absolute_energy(
+            preds[self.energy_name], batch, fp64=fp64_energy
+        )
         out[self.grad_forces_name] = preds[self.grad_forces_name]
         if self.has_stress:
+            assert self.grad_stress_name is not None, (
+                "grad_stress_name must be set if has_stress is True"
+            )
             out[self.grad_stress_name] = preds[self.grad_stress_name]
         out[self.grad_rotation_name] = preds[self.grad_rotation_name]
         for name in self.extra_properties:
             head = self.heads[name]
             if isinstance(head, ForcefieldHead):
-                out[name] = head.denormalize(preds[name], batch)
+                out[name] = preds[name]
             elif isinstance(head, ConfidenceHead):
                 out[name] = torch.softmax(preds[name], dim=-1)
             else:
@@ -287,7 +303,6 @@ def loss(self, batch: AtomGraphs) -> base.ModelOutput:
 
         energy_pred = out[self.energy_name]
         raw_grad_forces_pred = out[self.grad_forces_name]
-        grad_forces_pred = self.grad_forces_normalizer(raw_grad_forces_pred, online=False)
 
         # metrics
         metrics: dict = {}
@@ -309,7 +324,7 @@ def loss(self, batch: AtomGraphs) -> base.ModelOutput:
 
         # Conservative forces
         loss_out = forces_loss_function(
-            pred=grad_forces_pred,
+            raw_pred=raw_grad_forces_pred,
             raw_target=batch.node_targets[self.forces_name],
             raw_gold_target=batch.node_targets[self.forces_name],
             name=self.forces_name,
@@ -326,10 +341,13 @@ def loss(self, batch: AtomGraphs) -> base.ModelOutput:
 
         # Conservative stress (optional)
         if self.has_stress and self.grad_stress_name in out:
+            assert self.stress_name is not None, "stress_name must be set if has_stress is True"
+            assert self.grad_stress_name is not None, (
+                "grad_stress_name must be set if has_stress is True"
+            )
             raw_grad_stress_pred = out[self.grad_stress_name]
-            grad_stress_pred = self.grad_stress_normalizer(raw_grad_stress_pred, online=False)
             loss_out = stress_loss_function(
-                pred=grad_stress_pred,
+                raw_pred=raw_grad_stress_pred,
                 raw_target=batch.system_targets[self.stress_name],
                 raw_gold_target=batch.system_targets[self.stress_name],
                 name=self.stress_name,
@@ -342,23 +360,15 @@ def loss(self, batch: AtomGraphs) -> base.ModelOutput:
             metrics.update({f"{self.grad_prefix}-{k}": v for k, v in loss_out.log.items()})
 
         # Direct forces / stress predictions
-        for grad_name, grad_pred in [
-            (self.grad_forces_name, raw_grad_forces_pred),
-        ] + (
-            [(self.grad_stress_name, out[self.grad_stress_name])]
-            if self.has_stress and self.grad_stress_name in out
-            else []
+        for grad_name in [self.grad_forces_name] + (
+            [self.grad_stress_name] if self.has_stress and self.grad_stress_name in out else []
         ):
+            assert grad_name is not None
             direct_name = grad_name.replace(self.grad_prefix + "_", "")
             if direct_name in self.extra_properties:
                 direct_head = cast(ForcefieldHead, self.heads[direct_name])
                 direct_pred = out[direct_name]
-                if self.distill_direct_heads:
-                    loss_out = direct_head.loss(
-                        direct_pred, batch, alternative_target=grad_pred.detach()
-                    )
-                else:
-                    loss_out = direct_head.loss(direct_pred, batch)
+                loss_out = direct_head.loss(direct_pred, batch)
                 loss = self.loss_weights[direct_name] * loss_out.loss
                 total_loss += loss
                 metrics.update(loss_out.log)