Expose OrbMol-v2 per-atom charges and dipole via ASE

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)
 

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()

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)

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

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()