EEM Partial Charges for 3D descriptors

benchmarks/descriptors3d_bench.py

@@ -15,9 +15,11 @@
 
 import argparse
 import math
+import re
 from pathlib import Path
 
 import numpy as np
+import rdkit
 import torch
 from bench_utils import (
     Deadline,
@@ -37,7 +39,7 @@
 from rdkit import Chem
 from rdkit.Chem import rdMolDescriptors
 
-from nvmolkit.descriptors3d import Calc3DProperties, Property3D
+from nvmolkit.descriptors3d import PER_ATOM_PROPERTIES, Calc3DProperties, Property3D
 from nvmolkit.types import AsyncGpuResult, Device3DResult, HardwareOptions, PrecisionMode
 
 PRECISIONS = {"single": PrecisionMode.SINGLE, "full": PrecisionMode.FULL}
@@ -48,6 +50,8 @@
 USR_SKEW_TOLERANCES = {PrecisionMode.SINGLE: (5e-3, 0.1), PrecisionMode.FULL: (1e-4, 1e-3)}
 # Relative tolerances for the DCLV properties, which sum float32 dot contributions over every atom.
 DCLV_RELATIVE_TOLERANCES = {PrecisionMode.SINGLE: 1e-5, PrecisionMode.FULL: 1e-9}
+# Absolute tolerances for EEM charges.
+EEM_TOLERANCES = {PrecisionMode.SINGLE: 2e-6, PrecisionMode.FULL: 1e-9}
 
 PROPERTY_SETS = {
     "single": (Property3D.RADIUS_OF_GYRATION,),
@@ -77,6 +81,25 @@
 }
 
 
+def _calc_rdkit_eem(mol: Chem.Mol, conf_id: int) -> list[float]:
+    """Reference charges, with NaN for unsupported types or invalid coordinates."""
+    missing = [math.nan] * mol.GetNumAtoms()
+    molecule = Chem.Mol(mol)
+    try:
+        Chem.Kekulize(molecule, clearAromaticFlags=True)
+    except Chem.MolSanitizeException:
+        return missing
+    supported_elements = {1: {1, 6, 7, 8, 9, 16, 17, 35}, 2: {6, 7, 8, 15, 16}, 3: {6, 7}}
+    for atom in molecule.GetAtoms():
+        bond_order = max((int(bond.GetBondTypeAsDouble()) for bond in atom.GetBonds()), default=1)
+        if atom.GetAtomicNum() not in supported_elements.get(bond_order, set()):
+            return missing
+    positions = molecule.GetConformer(conf_id).GetPositions()
+    if not np.isfinite(positions).all() or len(np.unique(positions, axis=0)) != len(positions):
+        return missing
+    return rdMolDescriptors.CalcEEMcharges(molecule, confId=conf_id)
+
+
 def _calc_rdkit_property(mol: Chem.Mol, conf_id: int, prop: Property3D) -> float | list[float]:
     """Calculate one property using the options mirrored by the benchmark."""
     if prop == Property3D.SPHEROCITY_INDEX:
@@ -100,6 +123,8 @@ def _calc_rdkit_property(mol: Chem.Mol, conf_id: int, prop: Property3D) -> float
         return rdMolDescriptors.GetUSR(mol, confId=conf_id)
     if prop == Property3D.USRCAT:
         return rdMolDescriptors.GetUSRCAT(mol, confId=conf_id)
+    if prop == Property3D.EEM_CHARGES:
+        return _calc_rdkit_eem(mol, conf_id)
     if prop == Property3D.GETAWAY:
         if len(Chem.GetMolFrags(mol)) > 1:
             # RDKit's GETAWAY does not finish in reasonable time on multi-fragment molecules.
@@ -166,7 +191,29 @@ def _calc_rdkit(mols: list[Chem.Mol], properties: tuple[Property3D, ...]) -> dic
         for conf in mol.GetConformers():
             for prop, value in _calc_rdkit_conformer(mol, conf.GetId(), properties).items():
                 values[prop].append(value)
-    return {prop: np.asarray(rows, dtype=np.float64) for prop, rows in values.items()}
+    return {
+        prop: np.concatenate([np.asarray(row, dtype=np.float64) for row in rows])
+        if prop in PER_ATOM_PROPERTIES
+        else np.asarray(rows, dtype=np.float64)
+        for prop, rows in values.items()
+    }
+
+
+# RDKit before 2026.9.2 returns wrong EEM charges for charged molecules (rdkit/rdkit#9688); nvMolKit's are correct,
+# so only neutral molecules are compared against older RDKit.
+RDKIT_EEM_CHARGES_CORRECT = tuple(int(part) for part in re.findall(r"\d+", rdkit.__version__)[:3]) >= (2026, 9, 2)
+
+
+def _eem_comparable_rows(mols: list[Chem.Mol]) -> np.ndarray:
+    """Mask of the per-atom EEM rows, in molecule/conformer order, where the installed RDKit is a valid reference."""
+    return np.concatenate(
+        [
+            np.full(mol.GetNumAtoms(), RDKIT_EEM_CHARGES_CORRECT or Chem.GetFormalCharge(mol) == 0)
+            for mol in mols
+            for _ in mol.GetConformers()
+        ]
+        or [np.zeros(0, dtype=bool)]
+    )
 
 
 def _validate(
@@ -199,6 +246,13 @@ def _validate(
             if prop in DCLV_GETTERS:
                 np.testing.assert_allclose(actual, expected[prop], rtol=DCLV_RELATIVE_TOLERANCES[precision], atol=atol)
                 continue
+            if prop == Property3D.EEM_CHARGES:
+                eem_atol = EEM_TOLERANCES[precision] if tolerance is None else tolerance
+                comparable = _eem_comparable_rows(check_mols)
+                np.testing.assert_allclose(
+                    actual[comparable], expected[prop][comparable], rtol=0, atol=eem_atol, equal_nan=True
+                )
+                continue
             if prop not in (Property3D.USR, Property3D.USRCAT):
                 np.testing.assert_allclose(actual, expected[prop], rtol=rtol, atol=atol)
                 continue

benchmarks/tests/test_descriptors3d_bench.py

@@ -4,6 +4,7 @@
 """Regression tests for capped descriptor benchmark reporting."""
 
 import csv
+import importlib
 from contextlib import nullcontext
 
 import descriptors3d_bench as bench
@@ -13,6 +14,33 @@
 from rdkit import Chem
 
 
+@pytest.mark.parametrize(
+    ("version", "correct"),
+    [
+        ("2025.09.6", False),
+        ("2026.09.1", False),
+        ("2026.09.1pre", False),
+        ("2026.09.2", True),
+        ("2026.09.3pre", True),
+        ("2027.03.1", True),
+    ],
+)
+def test_eem_reference_version_accepts_development_suffix(monkeypatch, version, correct):
+    """Development RDKit builds import successfully and select the appropriate reference mask."""
+    try:
+        with monkeypatch.context() as patch:
+            patch.setattr(bench.rdkit, "__version__", version)
+            importlib.reload(bench)
+            assert bench.RDKIT_EEM_CHARGES_CORRECT is correct
+            charged = Chem.MolFromSmiles("C[NH3+]")
+            charged.AddConformer(Chem.Conformer(charged.GetNumAtoms()), assignId=True)
+            neutral = Chem.MolFromSmiles("CCO")
+            neutral.AddConformer(Chem.Conformer(neutral.GetNumAtoms()), assignId=True)
+            assert bench._eem_comparable_rows([charged, neutral]).tolist() == [correct] * 2 + [True] * 3
+    finally:
+        importlib.reload(bench)
+
+
 @pytest.mark.parametrize("completed", [0, 1, 2])
 @pytest.mark.parametrize("device_input", [False, True])
 def test_descriptor_rows_and_csv_handle_zero_partial_and_complete_cpu_work(

nvmolkit/descriptors3d.cpp

@@ -57,7 +57,8 @@ bp::object toOwnedPyArray(nvMolKit::PyArray* array) {
   return bp::object(bp::handle<>(Converter()(array)));
 }
 
-//! Returns (properties dict in request order, molIndices or None, confIndices or None).
+//! Returns (properties dict in request order, molIndices, confIndices and atomStarts, each None when
+//! coordinates came from the caller).
 template <typename Real>
 bp::object calc3DPropertiesAs(const std::vector<const RDKit::ROMol*>&  mols,
                               const std::vector<nvMolKit::Property3D>& properties,
@@ -70,6 +71,7 @@ bp::object calc3DPropertiesAs(const std::vector<const RDKit::ROMol*>&  mols,
   bp::dict output;
   for (const nvMolKit::Property3D property : properties) {
     auto&      values = result.properties.at(property);
+    // Rows are conformers or atoms (property3DExtent); either way the buffer holds rows x width values.
     const int  width  = nvMolKit::property3DWidth(property);
     const auto shape =
       width == 1 ? bp::make_tuple(values.size()) : bp::make_tuple(values.size() / static_cast<size_t>(width), width);
@@ -78,11 +80,12 @@ bp::object calc3DPropertiesAs(const std::vector<const RDKit::ROMol*>&  mols,
   // Row labels exist only when coordinates came from the molecules; otherwise the caller already
   // holds the labels of its own coordinate batch.
   if (view != nullptr) {
-    return bp::make_tuple(output, bp::object(), bp::object());
+    return bp::make_tuple(output, bp::object(), bp::object(), bp::object());
   }
   return bp::make_tuple(output,
                         toOwnedPyArray(nvMolKit::makePyArray(result.molIndices)),
-                        toOwnedPyArray(nvMolKit::makePyArray(result.confIndices)));
+                        toOwnedPyArray(nvMolKit::makePyArray(result.confIndices)),
+                        toOwnedPyArray(nvMolKit::makePyArray(result.atomStarts)));
 }
 
 bp::object calc3DProperties(const bp::list&                mols,

nvmolkit/descriptors3d.py

@@ -17,7 +17,16 @@
 import torch
 
 from nvmolkit import _descriptors3d
-from nvmolkit.types import AsyncGpuResult, Device3DResult, HardwareOptions, PrecisionMode, _resolve_cuda_stream
+from nvmolkit.types import (
+    AsyncGpuResult,
+    Device3DResult,
+    HardwareOptions,
+    PrecisionMode,
+    _atom_row_gather,
+    _atom_row_index,
+    _resolve_cuda_stream,
+    _scatter_atom_rows,
+)
 
 
 class Property3D(Enum):
@@ -47,6 +56,24 @@ class Property3D(Enum):
     - ``USRCAT``: 60 values, ``USR`` for all atoms and then for RDKit's
       hydrophobic, aromatic, acceptor and donor atoms. Same atom-count rule.
     - ``GETAWAY``: 273 values in RDKit's ``CalcGETAWAY`` order.
+
+    Per-atom properties have one row per atom rather than per conformer (see
+    :class:`Device3DPropertyResult`):
+
+    - ``EEMcharges``: RDKit's ``CalcEEMcharges`` electronegativity-equalization
+      partial charges, one value per atom. A molecule's charges sum to its
+      formal charge.
+
+      .. note::
+         Partial charges on charged molecules will not match RDKit < 2026.9.2
+         due to RDKit bug `#9688 <https://github.com/rdkit/rdkit/issues/9688>`_.
+         nvMolKit gives the corrected values regardless of the RDKit version.
+
+      Conformers containing an
+      element or bond type without nonzero parameters, conformers with coincident atoms (or non-finite
+      coordinates), for which RDKit returns all-zero charges, molecules
+      that cannot be kekulized, which RDKit rejects, and conformers of more than
+      46,000 atoms produce NaN.
     """
 
     PMI1 = "PMI1"
@@ -74,6 +101,11 @@ class Property3D(Enum):
     DCLV_POLAR_VOLUME = "DCLVPolarVolume"
     DCLV_COMPACTNESS = "DCLVCompactness"
     DCLV_PACKING_DENSITY = "DCLVPackingDensity"
+    EEM_CHARGES = "EEMcharges"
+
+
+#: Properties with one output row per atom instead of per conformer.
+PER_ATOM_PROPERTIES = frozenset({Property3D.EEM_CHARGES})
 
 
 @dataclass(frozen=True)
@@ -172,16 +204,19 @@ class Dense3DPropertyResult:
         values: One tensor of shape ``(n_mols, max_confs, *property_shape)`` per property, in request
             order, with the dtype of the source result. Scalar properties have no trailing dimensions;
             vector properties have ``property_shape == (width,)`` (widths listed on :class:`Property3D`).
+            Per-atom properties have shape ``(n_mols, max_confs, max_atoms, *property_shape)``.
             Padded slots hold the ``pad_value`` passed to :meth:`Device3DPropertyResult.dense`.
         conf_mask: bool ``(n_mols, max_confs)``; ``True`` where a real conformer exists.
+        atom_mask: bool ``(n_mols, max_confs, max_atoms)``; ``True`` where a real atom exists.
     """
 
     values: dict[str, torch.Tensor]
     conf_mask: torch.Tensor
+    atom_mask: torch.Tensor
 
 
 class Device3DPropertyResult(Mapping[str, AsyncGpuResult]):
-    """Per-conformer 3D properties on the GPU, labeled by molecule and conformer.
+    """Conformer-dependent 3D properties on the GPU, labeled by molecule and conformer.
 
     Behaves as a read-only mapping from property name to an :class:`~nvmolkit.types.AsyncGpuResult`
     whose first dimension is ``n_conformers``, in request order. Scalar properties have shape
@@ -191,7 +226,12 @@ class Device3DPropertyResult(Mapping[str, AsyncGpuResult]):
     :attr:`~nvmolkit.types.PrecisionMode.FULL`. Keys may be given as names or
     :class:`Property3D` members.
 
-    Row ``i`` of every property belongs to conformer ``conf_indices[i]`` of input molecule
+    Per-atom properties (:data:`PER_ATOM_PROPERTIES`) instead have one row per atom, laid out like
+    the coordinates: rows ``atom_starts[i]:atom_starts[i + 1]`` belong to conformer ``i``, in the
+    molecule's atom order. With ``Device3DResult`` coordinates, rows match its ``values`` rows
+    one-to-one, and rows outside every conformer's range hold NaN.
+
+    Conformer row ``i`` belongs to conformer ``conf_indices[i]`` of input molecule
     ``mol_indices[i]``. ``conf_indices`` is the conformer's position within its molecule, not
     its RDKit conformer ID.
     When the properties were calculated from a ``Device3DResult``, rows follow that result and the
@@ -200,6 +240,8 @@ class Device3DPropertyResult(Mapping[str, AsyncGpuResult]):
     Attributes:
         mol_indices: int32 ``(n_conformers,)`` input-molecule index of each row.
         conf_indices: int32 ``(n_conformers,)`` per-molecule conformer position of each row.
+        atom_starts: int32 ``(n_conformers + 1,)`` offsets of each conformer's per-atom rows; shared
+            with ``coordinates.atom_starts`` when calculated from a ``Device3DResult``.
         gpu_id: GPU holding every buffer.
         n_mols: Number of input molecules, including those without conformers.
     """
@@ -209,13 +251,15 @@ def __init__(
         properties: dict[str, AsyncGpuResult],
         mol_indices: AsyncGpuResult,
         conf_indices: AsyncGpuResult,
+        atom_starts: AsyncGpuResult,
         gpu_id: int,
         n_mols: int,
     ) -> None:
         """Create a result from per-property device buffers and their row labels."""
         self._properties = properties
         self.mol_indices = mol_indices
         self.conf_indices = conf_indices
+        self.atom_starts = atom_starts
         self.gpu_id = gpu_id
         self.n_mols = n_mols
 
@@ -236,31 +280,34 @@ def __len__(self) -> int:
 
     @property
     def n_conformers(self) -> int:
-        """Number of rows (conformers) in every property buffer."""
+        """Number of conformers, the row count of every per-conformer property buffer."""
         return self.mol_indices.torch().numel()
 
     def dense(self, pad_value: float = float("nan")) -> Dense3DPropertyResult:
-        """Materialize padded molecule/conformer tensors for every property.
+        """Materialize padded molecule/conformer (and, for per-atom properties, atom) tensors.
 
-        Molecules with fewer than ``max_confs`` conformers (including none) receive ``pad_value``.
-        Reading the index tensors synchronizes implicitly.
+        Molecules with fewer than ``max_confs`` conformers (including none), and conformers with fewer
+        than ``max_atoms`` atoms, receive ``pad_value``. Reading the index tensors synchronizes implicitly.
         """
+        per_atom = [name for name in self._properties if Property3D(name) in PER_ATOM_PROPERTIES]
+        # Per-atom values bound the valid atom ranges; the gather positions are only built when they are needed.
+        n_rows = self._properties[per_atom[0]].torch().shape[0] if per_atom else None
+        index = _atom_row_index(self.atom_starts, self.mol_indices, self.conf_indices, self.n_mols, n_rows)
+        gather = _atom_row_gather(index) if per_atom else None
         mol_indices = self.mol_indices.torch().to(torch.int64)
         conf_indices = self.conf_indices.torch().to(torch.int64)
-        device = mol_indices.device
-        max_confs = int(torch.bincount(mol_indices, minlength=self.n_mols).max().item()) if mol_indices.numel() else 0
-
-        conf_mask = torch.zeros((self.n_mols, max_confs), dtype=torch.bool, device=device)
-        conf_mask[mol_indices, conf_indices] = True
         values = {}
         for name, result in self._properties.items():
             source = result.torch()
+            if Property3D(name) in PER_ATOM_PROPERTIES:
+                values[name] = _scatter_atom_rows(source, index, gather, pad_value)
+                continue
             dense_values = torch.full(
-                (self.n_mols, max_confs, *source.shape[1:]), pad_value, dtype=source.dtype, device=device
+                (*index.conf_mask.shape, *source.shape[1:]), pad_value, dtype=source.dtype, device=source.device
             )
             dense_values[mol_indices, conf_indices] = source
             values[name] = dense_values
-        return Dense3DPropertyResult(values=values, conf_mask=conf_mask)
+        return Dense3DPropertyResult(values=values, conf_mask=index.conf_mask, atom_mask=index.atom_mask)
 
 
 _DCLV_PROPERTIES = frozenset(prop for prop in Property3D if prop.value.startswith("DCLV"))
@@ -386,9 +433,11 @@ def Calc3DProperties(
 
     Returns:
         A :class:`Device3DPropertyResult` mapping each requested property name,
-        in request order, to a device vector with one row per conformer, plus
-        the molecule and conformer labels of every row. Without ``coordinates``,
-        rows follow input-molecule order, then RDKit conformer order.
+        in request order, to a device vector with one row per conformer (one per
+        atom for :data:`PER_ATOM_PROPERTIES`), plus the molecule and conformer
+        labels of every row and each conformer's per-atom row range. Without
+        ``coordinates``, rows follow input-molecule order, then RDKit conformer
+        order.
 
     Any subset of :class:`Property3D` can be requested in one call, mixing
     families and giving members or names. Each requested family runs once for
@@ -464,7 +513,7 @@ def Calc3DProperties(
         else _device_coordinate_interfaces(coordinates, len(normalized_mols), active_stream.device)
     )
 
-    raw_results, raw_mol_indices, raw_conf_indices = _descriptors3d.Calc3DProperties(
+    raw_results, raw_mol_indices, raw_conf_indices, raw_atom_starts = _descriptors3d.Calc3DProperties(
         normalized_mols,
         [prop.value for prop in normalized_properties],
         options.moments.useAtomicMasses,
@@ -490,13 +539,16 @@ def Calc3DProperties(
     if coordinates is None:
         mol_indices = AsyncGpuResult(raw_mol_indices, gpu_id=gpu_id)
         conf_indices = AsyncGpuResult(raw_conf_indices, gpu_id=gpu_id)
+        atom_starts = AsyncGpuResult(raw_atom_starts, gpu_id=gpu_id)
     else:
         mol_indices = coordinates.mol_indices
         conf_indices = coordinates.conf_indices
+        atom_starts = coordinates.atom_starts
     return Device3DPropertyResult(
         properties_by_name,
         mol_indices=mol_indices,
         conf_indices=conf_indices,
+        atom_starts=atom_starts,
         gpu_id=gpu_id,
         n_mols=len(normalized_mols),
     )