Add OrbMol-v2 with learnable electrostatics

FINETUNING_GUIDE.md

@@ -151,7 +151,7 @@ Number of atoms: 3
 ```bash
 python finetune.py \
   --data_path /path/to/your/dataset.db \
-  --base_model orb_v3_conservative_omol \
+  --base_model orbmol_v2 \
   --energy_loss_weight 0.1 \
   --forces_loss_weight 1.0 \
   --stress_loss_weight 0.0 \
@@ -164,7 +164,7 @@ python finetune.py \
 ```bash
 python finetune.py \
   --data_path /path/to/your/dataset.db \
-  --base_model orb_v3_conservative_omol \
+  --base_model orbmol_v2 \
   --custom_reference_energies /path/to/reference_energies.json \
   --energy_loss_weight 0.1 \
   --forces_loss_weight 1.0
@@ -175,7 +175,7 @@ python finetune.py \
 ```bash
 python finetune.py \
   --data_path /path/to/your/dataset.db \
-  --base_model orb_v3_conservative_omol \
+  --base_model orbmol_v2 \
   --custom_reference_energies /path/to/reference_energies.json \
   --trainable_reference_energies \
   --energy_loss_weight 0.1 \
@@ -187,7 +187,7 @@ python finetune.py \
 ```bash
 python finetune.py \
   --data_path /path/to/your/dataset.db \
-  --base_model orb_v3_conservative_omol \
+  --base_model orbmol_v2 \
   --trainable_reference_energies \
   --energy_loss_weight 0.1 \
   --forces_loss_weight 1.0
@@ -258,7 +258,7 @@ Lines starting with `#` are treated as comments and ignored.
 
 The script automatically handles the differences between conservative and direct models:
 
-- **Conservative models** (e.g., `orb_v3_conservative_omol`):
+- **Conservative models** (e.g., `orbmol_v2`):
   - Use `grad_forces` and `grad_stress` as **loss-weight keys**
   - Compute forces via automatic differentiation
 
@@ -276,7 +276,7 @@ If you prefer to write your own finetuning script, you can use the clean API dir
 from orb_models.forcefield import pretrained
 
 # Load model with custom configuration
-model, atoms_adapter = pretrained.orb_v3_conservative_omol(
+model, atoms_adapter = pretrained.orbmol_v2(
     device='cuda',
     precision='float32-high',
     train=True,
@@ -319,7 +319,7 @@ import torch
 from orb_models.forcefield import pretrained
 
 # Load model architecture (set train=False for inference)
-model, atoms_adapter = pretrained.orb_v3_conservative_omol(train=False)
+model, atoms_adapter = pretrained.orbmol_v2(train=False)
 
 # Load your finetuned checkpoint
 model.load_state_dict(torch.load('path/to/finetuned_checkpoint.pt'))
@@ -331,7 +331,7 @@ You can also specify loss weights when loading for further finetuning:
 
 ```python
 # Load for continued finetuning with different loss weights
-model, atoms_adapter = pretrained.orb_v3_conservative_omol(
+model, atoms_adapter = pretrained.orbmol_v2(
     train=True,
     loss_weights={'energy': 0.5, 'grad_forces': 20.0}
 )
@@ -358,7 +358,7 @@ Finetuning on ORCA wB97M-V data with different reference scheme:
 ```bash
 python finetune.py \
   --data_path my_dataset.db \
-  --base_model orb_v3_conservative_omol \
+  --base_model orbmol_v2 \
   --custom_reference_energies my_refs.json \
   --energy_loss_weight 1.0 \
   --forces_loss_weight 10.0 \
@@ -370,7 +370,7 @@ python finetune.py \
 from orb_models.forcefield import pretrained
 import torch
 
-model, atoms_adapter = pretrained.orb_v3_conservative_omol(train=False)
+model, atoms_adapter = pretrained.orbmol_v2(train=False)
 model.load_state_dict(torch.load('checkpoints/my_finetuned_model.pt'))
 # Reference energies from my_refs.json are now loaded!
 ```
@@ -384,7 +384,7 @@ from orb_models.forcefield import pretrained
 from orb_models.dataset.ase_sqlite_dataset import AseSqliteDataset
 
 # Load model with configuration
-model, atoms_adapter = pretrained.orb_v3_conservative_omol(
+model, atoms_adapter = pretrained.orbmol_v2(
     device='cuda',
     train=True,
     train_reference_energies=False,  # Fixed reference energies

MODELS.md

@@ -7,11 +7,27 @@ 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.
 
+### OrbMol-v2 (learnable electrostatics)
+
+* `orbmol-v2`
+
+OrbMol-v2 extends the OrbMol architecture with **learnable per-atom electrostatics**: a `LatentChargeHead` predicts per-atom latent charges constrained to sum to the system total charge, 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 per-atom charges. Similar to orbmol-v1, system-level total charge and spin are required.
+
+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
+model, atoms_adapter = orbmol_v2(device="cuda")
+# atoms.info["charge"] and atoms.info["spin"] (multiplicity, = 2S+1) must be set.
+```
+
+> **Caution:** While the model does predict per-atom charge values as a latent feature in the charge head, the model has not seen any per-atom charge values during training; these are emergent from optimisation against energies and forces alone. They should therefore be treated with caution: while in at least some cases they appear to correspond to the correct physical values, the reliability and generality of this correspondence is unclear and is the subject of ongoing investigations.
+
 ### [V3 Models](https://arxiv.org/abs/2504.06231)
 
 V3 models use the following naming convention: ```orb-v3-X-Y-Z``` where:

README.md

@@ -21,12 +21,20 @@ Alternatively, you can use Docker to run orb-models; [see instructions below](#d
 
 ### Updates
 
+**May 2026**: Release of OrbMol-v2 — adds a `CoulombModule` for long-range electrostatics on top of the OrbMol architecture, using direct Coulomb summation for non-periodic systems and Particle Mesh Ewald (via `nvalchemiops`) for periodic. Trained on OMol25 and OPoly26 (ωB97M-V/def2-TZVPD); load with `pretrained.orbmol_v2(device="cuda")`. See [MODELS.md](MODELS.md) for the full architecture description.
+
+* **Long-range electrostatics and learnable charges.** GSCDB138 Normalized Error Ratio drops from **6.05 → 1.62** (3.7× lower, comparable to a good DFT functional).
+* **Full-model compilation.** `model.compile(...)` now wraps the full regressor for all models, giving ~1.7× speedup at 10k atoms on a single 80 GB GPU.
+
+`model.predict(...)["energy"]` now returns **fp64** by default to preserve kJ/mol resolution against OMol-scale references (~1e4–1e5 eV). Pass `fp64_energy=False` to opt out.
+
 **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)
 
+
 **August 2025**: Release of the [OrbMol potentials](https://www.orbitalindustries.com/posts/orbmol-extending-orb-to-molecular-systems):
 
 * 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.
@@ -189,7 +197,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_v2(
   device=device,
   precision="float32-high",   # or "float32-highest" / "float64
 )

examples/NaClWaterMD.py

@@ -60,8 +60,8 @@ def run_md_simulation(
     # Set the calculator
     # Note: If you encounter compilation errors (e.g., Triton issues on clusters),
     # you can disable compilation by adding compile=False:
-    # orbff, atoms_adapter = pretrained.orb_v3_conservative_omol(device=device, compile=False)
-    orbff, atoms_adapter = pretrained.orb_v3_conservative_omol(device=device)
+    # orbff, atoms_adapter = pretrained.orbmol_v2(device=device, compile=False)
+    orbff, atoms_adapter = pretrained.orbmol_v2(device=device)
     atoms.calc = ORBCalculator(orbff, atoms_adapter=atoms_adapter, device=device)
 
     # Set the initial velocities

finetune.py

@@ -29,6 +29,9 @@
 from orb_models.common.training.util import get_optim, init_device
 from orb_models.common.utils import seed_everything
 from orb_models.forcefield import pretrained
+from orb_models.forcefield.models.conservative_regressor import (
+    ConservativeForcefieldRegressor,
+)
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 
@@ -439,26 +442,30 @@ def run(args):
     # GPUs and does not appear to hurt training
     precision = "float32-high"
 
-    # Prepare loss weights if specified
-    loss_weights = {}
-    is_conservative_model = "conservative" in args.base_model
+    # Instantiate model with configuration
+    base_model = args.base_model
+    model, atoms_adapter = getattr(pretrained, base_model)(
+        device=device,
+        precision=precision,
+        train=True,
+        train_reference_energies=args.trainable_reference_energies,
+    )
+
+    # Detect conservative vs direct from the instantiated model type.
+    is_conservative_model = isinstance(model, ConservativeForcefieldRegressor)
 
+    # Map CLI loss-weight flags onto the keys the instantiated model expects.
+    loss_weights: dict[str, float] = {}
     if args.energy_loss_weight is not None:
         loss_weights["energy"] = args.energy_loss_weight
 
     if args.forces_loss_weight is not None:
-        # Key depends on model type
-        if is_conservative_model:
-            loss_weights["grad_forces"] = args.forces_loss_weight
-        else:  # direct model
-            loss_weights["forces"] = args.forces_loss_weight
+        key = "grad_forces" if is_conservative_model else "forces"
+        loss_weights[key] = args.forces_loss_weight
 
     if args.stress_loss_weight is not None:
-        # Key depends on model type
-        if is_conservative_model:
-            loss_weights["grad_stress"] = args.stress_loss_weight
-        else:  # direct model
-            loss_weights["stress"] = args.stress_loss_weight
+        key = "grad_stress" if is_conservative_model else "stress"
+        loss_weights[key] = args.stress_loss_weight
 
     if args.equigrad_loss_weight is not None:
         if not is_conservative_model:
@@ -471,16 +478,7 @@ def run(args):
         for key, val in loss_weights.items():
             logging.info(f"  {key}: {val}")
         logging.info("=" * 60)
-
-    # Instantiate model with configuration
-    base_model = args.base_model
-    model, atoms_adapter = getattr(pretrained, base_model)(
-        device=device,
-        precision=precision,
-        train=True,
-        train_reference_energies=args.trainable_reference_energies,
-        loss_weights=loss_weights if loss_weights else None,
-    )
+        model.loss_weights.update(loss_weights)
 
     # Handle custom reference energies if provided
     if args.custom_reference_energies:
@@ -678,12 +676,15 @@ def main():
         type=str,
         help="Base model to finetune.",
         choices=[
+            "orbmol_v2",
+            "orb_v3_conservative_omol",
+            "orb_v3_direct_omol",
+            "orbmol_v1_conservative",
+            "orbmol_v1_direct",
             "orb_v3_conservative_inf_omat",
             "orb_v3_conservative_20_omat",
             "orb_v3_direct_inf_omat",
             "orb_v3_direct_20_omat",
-            "orb_v3_conservative_omol",
-            "orb_v3_direct_omol",
             "orb_v2",
         ],
     )

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/common/models/segment_ops.py

@@ -44,7 +44,7 @@ def aggregate_nodes(
     if reduction == "sum":
         return scatter_sum(tensor, segments, dim=0, dim_size=count)
     elif reduction == "mean":
-        return scatter_mean(tensor, segments, dim=0, dim_size=count)
+        return scatter_mean(tensor, segments, n_node, dim=0, dim_size=count)
     elif reduction == "max":
         return segment_max(tensor, segments, num_segments=count)
     else:
@@ -61,11 +61,6 @@ def segment_max(data: torch.Tensor, segment_ids: torch.Tensor, num_segments: int
     return scatter_max(data, segment_ids, dim=0, dim_size=num_segments)
 
 
-def segment_mean(data: torch.Tensor, segment_ids: torch.Tensor, num_segments: int):
-    """Computes index based mean over segments of a tensor."""
-    return scatter_mean(data, segment_ids, dim=0, dim_size=num_segments)
-
-
 def segment_softmax(
     data: torch.Tensor,
     segment_ids: torch.Tensor,
@@ -222,6 +217,7 @@ def scatter_std(
 def scatter_mean(
     src: torch.Tensor,
     index: torch.Tensor,
+    count: torch.Tensor,
     dim: int = -1,
     out: torch.Tensor | None = None,
     dim_size: int | None = None,
@@ -231,6 +227,7 @@ def scatter_mean(
     Args:
         src (torch.Tensor): The source tensor.
         index (torch.Tensor): The indices of elements to scatter.
+        count (torch.Tensor): Pre-computed group sizes (e.g. n_node).
         dim (int, optional): The dimension along which to index. Defaults to -1.
         out (Optional[torch.Tensor], optional): The output tensor. Defaults to None.
         dim_size (Optional[int], optional): Size of the output tensor. Defaults to None.
@@ -241,20 +238,13 @@ def scatter_mean(
     out = scatter_sum(src, index, dim, out, dim_size)
     dim_size = out.size(dim)
 
-    index_dim = dim
-    if index_dim < 0:
-        index_dim = index_dim + src.dim()
-    if index.dim() <= index_dim:
-        index_dim = index.dim() - 1
-
-    ones = torch.ones(index.size(), dtype=src.dtype, device=src.device)
-    count = scatter_sum(ones, index, index_dim, None, dim_size)
-    count[count < 1] = 1
-    count = _broadcast(count, out, dim)
+    divisor = count.to(dtype=out.dtype)
+    divisor = divisor.clamp(min=1)
+    divisor = _broadcast(divisor, out, dim)
     if out.is_floating_point():
-        out.true_divide_(count)
+        out.true_divide_(divisor)
     else:
-        out.div_(count, rounding_mode="floor")
+        out.div_(divisor, rounding_mode="floor")
     return out