From 43746d2c8ad82366eb985b359921aae475758bb9 Mon Sep 17 00:00:00 2001 From: cvmbb Date: Tue, 25 Aug 2026 18:32:05 +0200 Subject: [PATCH] Expose OrbMol-v2 per-atom charges and dipole via ASE LatentChargeHead already computed per-atom charges and discarded them. Return them as "charges" from predict(), plus a derived "dipole", through ORBCalculator so atoms.get_charges() and get_dipole_moment() work. --- MODELS.md | 2 +- orb_models/forcefield/inference/calculator.py | 15 +++++- .../models/conservative_regressor.py | 12 ++++- tests/forcefield/test_calculator.py | 52 +++++++++++++++++++ tests/forcefield/test_conservative.py | 12 +++++ 5 files changed, 90 insertions(+), 3 deletions(-) diff --git a/MODELS.md b/MODELS.md index b1f53ff..a30b6e2 100644 --- a/MODELS.md +++ b/MODELS.md @@ -27,7 +27,7 @@ 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. +> **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. For trial purposes they are returned as `charges` from `predict()`, and as ASE's `charges` and `dipole` properties from `ORBCalculator`. ### [V3 Models](https://arxiv.org/abs/2504.06231) diff --git a/orb_models/forcefield/inference/calculator.py b/orb_models/forcefield/inference/calculator.py index 1a627df..3675135 100644 --- a/orb_models/forcefield/inference/calculator.py +++ b/orb_models/forcefield/inference/calculator.py @@ -62,7 +62,10 @@ def __init__( conditioner, ChargeSpinConditioner ) - self.implemented_properties = model.properties # type: ignore + properties = list(model.properties) # type: ignore + if "charges" in properties: + properties.append("dipole") + self.implemented_properties = properties def check_state(self, atoms, tol=1e-15): """Check if calculation is needed. @@ -122,7 +125,17 @@ def _update_results(self, out: dict[str, torch.Tensor]): # ASE expects: # - stresses to be squeezed to a 1D array of shape (6,) # - forces to never be squeezed i.e. single-atom systems should be (1, 3) + # - charges to be (n_atoms,) if prop == "stress": self.results[prop] = to_numpy(out[out_key].squeeze()) + elif prop == "charges": + self.results[prop] = out[out_key].detach().reshape(-1).cpu().numpy() else: self.results[prop] = to_numpy(out[out_key]) + + if "dipole" in self.implemented_properties and "charges" in self.results: + atoms = self.atoms + # Point-charge dipole sum_i q_i r_i, in e*A (ASE's convention). Only + # defined for non-periodic systems, and origin-dependent unless neutral. + if atoms is not None and not atoms.pbc.any(): + self.results["dipole"] = self.results["charges"] @ atoms.get_positions() diff --git a/orb_models/forcefield/models/conservative_regressor.py b/orb_models/forcefield/models/conservative_regressor.py index 09e2fad..b937650 100644 --- a/orb_models/forcefield/models/conservative_regressor.py +++ b/orb_models/forcefield/models/conservative_regressor.py @@ -99,7 +99,7 @@ def __init__( self.has_stress = has_stress - collisions = {"forces", "stress"} & heads.keys() + collisions = {"forces", "stress", "charges"} & heads.keys() assert not collisions, ( f"Heads {collisions} collide with gradient-based prediction keys in predict()." ) @@ -128,6 +128,8 @@ def properties(self) -> list[str]: if self.has_stress: props.append("stress") props.extend(self.extra_properties) + if "latent_charges" in self.heads: + props.append("charges") return props @property @@ -158,6 +160,7 @@ def forward( "energy" — absolute energy (B,); fp64 when fp64_energy=True "forces" — total forces (N, 3), autograd + explicit "stress" — total stress (B, 6) in Voigt notation + "charges" — per-atom charges (N,) in e, if a latent_charges head exists Components (used by loss and pipeline frameworks): "interaction_energy" — energy without reference (B,) @@ -182,6 +185,8 @@ def forward( latent_charges = None if "latent_charges" in self.heads: latent_charges = self.heads["latent_charges"](node_features, batch) + # (N,) for consumers; the (N, 1) form feeds the energy head / CoulombModule. + out["charges"] = latent_charges.squeeze(-1) latent_spins = None if "latent_spins" in self.heads: @@ -262,6 +267,8 @@ def predict( "energy" — absolute energy (B,) "forces" — total forces (N, 3) "stress" — total stress in Voigt notation (B, 6) + "charges" — per-atom charges (N,) in e, if a latent_charges head exists; + emergent from energy/force fitting alone, see MODELS.md """ # self() not self.forward() to respect torch.compile preds = self( @@ -286,6 +293,9 @@ def predict( else: raise ValueError(f"Expected ForcefieldHead or ConfidenceHead, got {type(head)}.") + if "charges" in preds: + out["charges"] = preds["charges"] + if split: for name, pred in out.items(): out[name] = split_prediction(pred, batch.n_node) diff --git a/tests/forcefield/test_calculator.py b/tests/forcefield/test_calculator.py index 8a18575..a0bab49 100644 --- a/tests/forcefield/test_calculator.py +++ b/tests/forcefield/test_calculator.py @@ -1,3 +1,8 @@ +import numpy as np +import pytest +from ase import Atoms +from ase.build import molecule + from orb_models.forcefield.forcefield_adapter import ForcefieldAtomsAdapter from orb_models.forcefield.inference.calculator import ORBCalculator @@ -77,3 +82,50 @@ def test_direct_stress_enabled(direct_regressor, mptraj_10_systems_db): calc.calculate(atoms) assert "stress" in calc.results assert "forces" in calc.results + + +@pytest.mark.parametrize( + ("atoms", "total_charge"), + [ + (molecule("H2O"), 1), + # Single atom: guards against to_numpy collapsing (1,) to a Python float. + (Atoms("H", positions=[[0.0, 0.0, 0.0]]), 0), + ], +) +def test_charges(conservative_regressor, atoms, total_charge): + """Charges are exposed to ASE as (n_atoms,) and sum to the requested total.""" + atoms.info["charge"] = total_charge + atoms.info["spin"] = 1 + calc = ORBCalculator( + model=conservative_regressor, + atoms_adapter=ForcefieldAtomsAdapter(6.0, 20), + ) + assert "charges" in calc.implemented_properties + atoms.calc = calc + + charges = atoms.get_charges() + assert charges.shape == (len(atoms),) + assert np.isfinite(charges).all() + assert charges.sum() == pytest.approx(total_charge, abs=1e-5) + + +def test_dipole(conservative_regressor, mptraj_10_systems_db): + """Dipole is the point-charge sum, and only available for non-periodic systems.""" + adapter = ForcefieldAtomsAdapter(6.0, 20) + + atoms = molecule("H2O") + atoms.info["charge"] = 0 + atoms.info["spin"] = 1 + atoms.calc = ORBCalculator(model=conservative_regressor, atoms_adapter=adapter) + + dipole = atoms.get_dipole_moment() + assert dipole.shape == (3,) + np.testing.assert_allclose(dipole, atoms.get_charges() @ atoms.get_positions(), rtol=1e-6) + + periodic = mptraj_10_systems_db.get_atoms(1) + periodic.info["charge"] = 0 + periodic.info["spin"] = 1 + calc = ORBCalculator(model=conservative_regressor, atoms_adapter=adapter) + calc.calculate(periodic) + assert "charges" in calc.results + assert "dipole" not in calc.results diff --git a/tests/forcefield/test_conservative.py b/tests/forcefield/test_conservative.py index ae4ee44..9f0b6c8 100644 --- a/tests/forcefield/test_conservative.py +++ b/tests/forcefield/test_conservative.py @@ -5,6 +5,7 @@ from ase import Atom, Atoms from orb_models.common.atoms.batch.graph_batch import AtomGraphs +from orb_models.common.models.segment_ops import aggregate_nodes from orb_models.forcefield.forcefield_adapter import ForcefieldAtomsAdapter from orb_models.forcefield.models.conservative_regressor import ConservativeForcefieldRegressor @@ -18,6 +19,17 @@ def test_regressor_forward(request, conservative_regressor, graph_name): assert "stress" in out +def test_predict_charges(conservative_regressor, batch): + """Per-atom charges are exposed as (N,) and sum to the system total charge.""" + assert "charges" in conservative_regressor.properties + charges = conservative_regressor.predict(batch)["charges"] + assert charges.shape == (batch.n_node.sum(),) + + # The batch fixture carries no total_charge, so charges are centered on zero. + per_system = aggregate_nodes(charges, batch.n_node, reduction="sum") + torch.testing.assert_close(per_system, torch.zeros_like(per_system), atol=1e-6, rtol=0) + + def test_regressor_loss(conservative_regressor, batch): out = conservative_regressor.loss(batch) out.loss.backward()