From 9a9c96f1a91c7144869ce847e6a25b388769cd1a Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Thu, 8 Oct 2026 14:11:28 -0400 Subject: [PATCH 1/8] Add EEM partial charges as a per-atom 3D property EEMcharges reproduces RDKit's CalcEEMcharges for every conformer. It is the first property with one output row per atom: a new Property3DExtent says whether a property's rows are conformers or atoms, and per-atom outputs are laid out like the coordinate positions, indexed by the batch's atomStarts. Device3DPropertyResult now carries atom_starts (shared with Device3DResult input) and dense() pads per-atom properties to (n_mols, max_confs, max_atoms) with an atom_mask; the padding helper is shared with Device3DResult.dense(). One block per conformer builds RDKit's bordered EEM system and solves it by LU with partial pivoting, in shared memory when it fits and per-block global scratch otherwise. As in RDKit (which hands its row-major matrix to Eigen as column-major), charges sum to the negated formal charge. SINGLE refines the float32 solution once against a float64 residual: float32 elimination alone drifted up to ~3e-4 on 160-340 atom conformers, refinement brings SINGLE within 8e-7 of FULL. FULL matches RDKit to 1e-12. Atoms the parameter set does not cover produce NaN for their conformer: elements beyond bromine and bond orders above 3, where RDKit reads undefined memory, and, unlike RDKit, zero-parameter atoms within its tables (Si, single-bonded P, Na+ counter-ions, ...), which RDKit silently solves to charges of -5 to -168. Molecules that fail to kekulize and conformers with coincident atoms (where RDKit returns all-zero charges) are NaN as well. --- benchmarks/descriptors3d_bench.py | 18 +- nvmolkit/descriptors3d.cpp | 9 +- nvmolkit/descriptors3d.py | 72 ++++-- nvmolkit/tests/test_descriptors3d.py | 118 ++++++++- nvmolkit/types.py | 85 ++++--- src/descriptors3d.cu | 19 +- src/descriptors3d.h | 43 +++- src/descriptors3d_eem.cuh | 348 +++++++++++++++++++++++++++ src/descriptors3d_mol.cpp | 100 +++++++- src/descriptors3d_mol.h | 7 +- tests/test_descriptors3d.cpp | 131 ++++++++++ 11 files changed, 879 insertions(+), 71 deletions(-) create mode 100644 src/descriptors3d_eem.cuh diff --git a/benchmarks/descriptors3d_bench.py b/benchmarks/descriptors3d_bench.py index 77fc6afb..cb224989 100644 --- a/benchmarks/descriptors3d_bench.py +++ b/benchmarks/descriptors3d_bench.py @@ -37,7 +37,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 +48,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,), @@ -100,6 +102,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 rdMolDescriptors.CalcEEMcharges(mol, confId=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 +170,12 @@ 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() + } def _validate( @@ -199,6 +208,11 @@ 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: + # Atoms RDKit has no parameters for are NaN in both. + eem_atol = EEM_TOLERANCES[precision] if tolerance is None else tolerance + np.testing.assert_allclose(actual, expected[prop], 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 diff --git a/nvmolkit/descriptors3d.cpp b/nvmolkit/descriptors3d.cpp index ba6e0480..77e3d7c7 100644 --- a/nvmolkit/descriptors3d.cpp +++ b/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 bp::object calc3DPropertiesAs(const std::vector& mols, const std::vector& properties, @@ -70,6 +71,7 @@ bp::object calc3DPropertiesAs(const std::vector& 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(width), width); @@ -78,11 +80,12 @@ bp::object calc3DPropertiesAs(const std::vector& 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, diff --git a/nvmolkit/descriptors3d.py b/nvmolkit/descriptors3d.py index 3c7f0347..de3834f8 100644 --- a/nvmolkit/descriptors3d.py +++ b/nvmolkit/descriptors3d.py @@ -17,7 +17,15 @@ 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_index, + _resolve_cuda_stream, + _scatter_atom_rows, +) class Property3D(Enum): @@ -47,6 +55,17 @@ 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. As in RDKit, a molecule's charges + sum to the negative of its formal charge. Conformers containing an + element beyond bromine or a bond order above 3, for which RDKit reads + undefined parameters, conformers with coincident atoms (or non-finite + coordinates), for which RDKit returns all-zero charges, and molecules + that cannot be kekulized, which RDKit rejects, produce NaN. """ PMI1 = "PMI1" @@ -74,6 +93,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,12 +196,15 @@ 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]): @@ -191,6 +218,11 @@ class Device3DPropertyResult(Mapping[str, AsyncGpuResult]): :attr:`~nvmolkit.types.PrecisionMode.FULL`. Keys may be given as names or :class:`Property3D` members. + 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. + Row ``i`` of every property 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. @@ -200,6 +232,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,6 +243,7 @@ def __init__( properties: dict[str, AsyncGpuResult], mol_indices: AsyncGpuResult, conf_indices: AsyncGpuResult, + atom_starts: AsyncGpuResult, gpu_id: int, n_mols: int, ) -> None: @@ -216,6 +251,7 @@ def __init__( 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 +272,30 @@ 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. """ + index = _atom_row_index(self.atom_starts, self.mol_indices, self.conf_indices, self.n_mols) 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, 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 +421,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 +501,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 +527,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), ) diff --git a/nvmolkit/tests/test_descriptors3d.py b/nvmolkit/tests/test_descriptors3d.py index 17319dfa..ea1cec60 100644 --- a/nvmolkit/tests/test_descriptors3d.py +++ b/nvmolkit/tests/test_descriptors3d.py @@ -12,6 +12,7 @@ from rdkit.Geometry import Point3D from nvmolkit.descriptors3d import ( + PER_ATOM_PROPERTIES, Calc3DProperties, DclvOptions, Device3DPropertyResult, @@ -44,7 +45,9 @@ Property3D.DCLV_PACKING_DENSITY, ) # Shape and moment scalars, compared against RDKit by _rdkit_property; the DCLV scalars have their own tests. -SCALAR_PROPERTIES = tuple(prop for prop in Property3D if prop not in VECTOR_PROPERTIES + DCLV_PROPERTIES) +SCALAR_PROPERTIES = tuple( + prop for prop in Property3D if prop not in VECTOR_PROPERTIES + DCLV_PROPERTIES and prop not in PER_ATOM_PROPERTIES +) PAIRWISE_PROPERTIES = (Property3D.RDF, Property3D.MORSE, Property3D.AUTOCORR3D) RDKIT_PAIRWISE = { Property3D.RDF: rdMolDescriptors.CalcRDF, @@ -958,3 +961,116 @@ def test_property_and_coordinate_contract_errors_are_clear(): coordinates = _device_result_from_molecules([mol]) with pytest.raises(ValueError, match=r"coordinates\.n_mols"): Calc3DProperties([mol, mol], ["PMI1"], coordinates=coordinates) + + +# float32 output rounding of charges up to ~20; the float32 solve is refined against a float64 residual. +EEM_SINGLE_ATOL = 2e-6 +EEM_SMILES = ("CC(=O)[O-]", "C[NH3+]", "c1ccncc1O", "O=c1cc[nH]cc1", "CS(=O)(=O)NCl", "FC(F)(F)c1ccc(Br)cc1") + + +def _rdkit_eem_rows(mols): + """RDKit EEM charges of every conformer, concatenated in molecule, then conformer order.""" + return np.concatenate( + [ + np.asarray(rdMolDescriptors.CalcEEMcharges(mol, confId=conf.GetId()), dtype=np.float64) + for mol in mols + for conf in mol.GetConformers() + ] + ) + + +def _assert_eem_matches_rdkit(actual, expected, precision): + if precision == PrecisionMode.FULL: + assert actual.dtype == np.float64 + np.testing.assert_allclose(actual, expected, rtol=1e-10, atol=1e-10, equal_nan=True) + else: + assert actual.dtype == np.float32 + np.testing.assert_allclose(actual, expected, rtol=0, atol=EEM_SINGLE_ATOL, equal_nan=True) + + +@pytest.mark.parametrize("precision", PRECISIONS) +def test_eem_charges_are_per_atom_and_match_rdkit(precision): + mols = [_embed(smiles, 2, 7 + idx) for idx, smiles in enumerate(EEM_SMILES)] + result = Calc3DProperties(mols, [Property3D.NPR1, Property3D.EEM_CHARGES], precision=precision) + + total_atoms = sum(mol.GetNumAtoms() * mol.GetNumConformers() for mol in mols) + expected_starts = np.cumsum([0] + [mol.GetNumAtoms() for mol in mols for _ in mol.GetConformers()]) + assert result[Property3D.NPR1].torch().shape == (result.n_conformers,) + assert result[Property3D.EEM_CHARGES].torch().shape == (total_atoms,) + np.testing.assert_array_equal(result.atom_starts.numpy(), expected_starts) + _assert_eem_matches_rdkit(result[Property3D.EEM_CHARGES].numpy(), _rdkit_eem_rows(mols), precision) + + +def test_eem_charges_sum_to_negated_formal_charge_like_rdkit(): + mols = [_embed(smiles, 1, 3) for smiles in ("CC(=O)[O-]", "C[NH3+]", "CCO")] + result = Calc3DProperties(mols, Property3D.EEM_CHARGES, precision=PrecisionMode.FULL) + starts = result.atom_starts.torch().tolist() + charges = result[Property3D.EEM_CHARGES].numpy() + sums = [charges[starts[row] : starts[row + 1]].sum() for row in range(len(mols))] + np.testing.assert_allclose(sums, [1.0, -1.0, 0.0], atol=1e-9) + + +def test_eem_dense_pads_conformers_and_atoms(): + mols = [_embed("CCO", 3, 1), Chem.AddHs(Chem.MolFromSmiles("CC")), _embed("c1ccncc1O", 1, 2)] + result = Calc3DProperties(mols, [Property3D.PBF, Property3D.EEM_CHARGES], precision=PrecisionMode.FULL) + dense = result.dense() + + max_atoms = max(mol.GetNumAtoms() for mol in mols) + charges = dense.values[Property3D.EEM_CHARGES.value] + assert charges.shape == (3, 3, max_atoms) + assert dense.values[Property3D.PBF.value].shape == (3, 3) + assert dense.atom_mask.shape == (3, 3, max_atoms) + assert int(dense.atom_mask.sum()) == result[Property3D.EEM_CHARGES].torch().numel() + torch.testing.assert_close(charges[dense.atom_mask], result[Property3D.EEM_CHARGES].torch()) + assert torch.isnan(charges[~dense.atom_mask]).all() + assert dense.conf_mask.tolist() == [[True, True, True], [False, False, False], [True, False, False]] + + +@pytest.mark.parametrize("precision", PRECISIONS) +def test_eem_device_coordinates_share_atom_starts(precision): + mols = [Chem.AddHs(Chem.MolFromSmiles(smiles)) for smiles in ("CC(=O)[O-]", "c1ccncc1O")] + params = EmbedParameters() + params.randomSeed = 0xC0FFEE + coordinates = EmbedMolecules(mols, params, confsPerMolecule=2, output=CoordinateOutput.DEVICE) + + result = Calc3DProperties(mols, Property3D.EEM_CHARGES, coordinates=coordinates, precision=precision) + + assert result.atom_starts is coordinates.atom_starts + values = coordinates.values.numpy() + atom_starts = coordinates.atom_starts.torch().tolist() + expected = np.full(len(values), np.nan) + for row, mol_idx in enumerate(coordinates.mol_indices.torch().tolist()): + mol = Chem.Mol(mols[mol_idx]) + conf = Chem.Conformer(mol.GetNumAtoms()) + for atom_idx, (x, y, z) in enumerate(values[atom_starts[row] : atom_starts[row + 1]]): + conf.SetAtomPosition(atom_idx, Point3D(float(x), float(y), float(z))) + conf_id = mol.AddConformer(conf, assignId=True) + expected[atom_starts[row] : atom_starts[row + 1]] = rdMolDescriptors.CalcEEMcharges(mol, confId=conf_id) + _assert_eem_matches_rdkit(result[Property3D.EEM_CHARGES].numpy(), expected, precision) + + +@pytest.mark.parametrize("precision", PRECISIONS) +def test_eem_atoms_without_parameters_give_nan_for_their_conformers(precision): + # RDKit has no EEM parameters past bromine and reads beyond its tables for iodine. + iodide = _embed("CCI", 2, 4) + ethanol = _embed("CCO", 1, 4) + result = Calc3DProperties([iodide, ethanol], Property3D.EEM_CHARGES, precision=precision) + charges = result[Property3D.EEM_CHARGES].numpy() + iodide_rows = 2 * iodide.GetNumAtoms() + assert np.isnan(charges[:iodide_rows]).all() + _assert_eem_matches_rdkit(charges[iodide_rows:], _rdkit_eem_rows([ethanol]), precision) + + +@pytest.mark.parametrize("precision", PRECISIONS) +def test_eem_atoms_with_zero_parameters_give_nan_unlike_rdkit(precision): + # Deliberately differs from RDKit: its parameter tables hold zeros for elements and bond orders they do not + # cover (here Si and single-bonded P), which RDKit silently solves with, returning charges of -10 or beyond. + # nvMolKit reports NaN for those conformers instead. + uncovered = [_embed(smiles, 1, 5) for smiles in ("C[Si](C)(C)C", "CP(C)C")] + ethanol = _embed("CCO", 1, 5) + result = Calc3DProperties([*uncovered, ethanol], Property3D.EEM_CHARGES, precision=precision) + charges = result[Property3D.EEM_CHARGES].numpy() + uncovered_rows = sum(mol.GetNumAtoms() for mol in uncovered) + assert np.isnan(charges[:uncovered_rows]).all() + assert np.abs(_rdkit_eem_rows(uncovered)).max() > 5 # RDKit's zero-parameter charges + _assert_eem_matches_rdkit(charges[uncovered_rows:], _rdkit_eem_rows([ethanol]), precision) diff --git a/nvmolkit/types.py b/nvmolkit/types.py index b0883e72..143ed072 100644 --- a/nvmolkit/types.py +++ b/nvmolkit/types.py @@ -522,45 +522,56 @@ def dense(self, pad_value: float = float("nan")) -> "Dense3DResult": Both masks are ``True`` where the data is real, ``False`` where padded. Reading the index tensors synchronizes implicitly. """ - values = self.values.torch() - atom_starts = self.atom_starts.torch().to(torch.int64) - mol_indices = self.mol_indices.torch().to(torch.int64) - conf_indices = self.conf_indices.torch().to(torch.int64) - - device = values.device - dtype = values.dtype - n_conformers = mol_indices.numel() - - if n_conformers == 0: - dense_vals = torch.full((self.n_mols, 0, 0, 3), pad_value, dtype=dtype, device=device) - conf_mask = torch.zeros((self.n_mols, 0), dtype=torch.bool, device=device) - atom_mask = torch.zeros((self.n_mols, 0, 0), dtype=torch.bool, device=device) - return Dense3DResult(values=dense_vals, conf_mask=conf_mask, atom_mask=atom_mask) - - sizes = atom_starts[1:] - atom_starts[:-1] - confs_per_mol = torch.bincount(mol_indices, minlength=self.n_mols) - max_confs = int(confs_per_mol.max().item()) - max_atoms = int(sizes.max().item()) - - dense_vals = torch.full( - (self.n_mols, max_confs, max_atoms, 3), - pad_value, - dtype=dtype, - device=device, + index = _atom_row_index(self.atom_starts, self.mol_indices, self.conf_indices, self.n_mols) + return Dense3DResult( + values=_scatter_atom_rows(self.values.torch(), index, pad_value), + conf_mask=index.conf_mask, + atom_mask=index.atom_mask, ) - conf_mask = torch.zeros((self.n_mols, max_confs), dtype=torch.bool, device=device) - atom_mask = torch.zeros((self.n_mols, max_confs, max_atoms), dtype=torch.bool, device=device) - conf_mask[mol_indices, conf_indices] = True - mol_idx_per_atom = mol_indices.repeat_interleave(sizes) - conf_idx_per_atom = conf_indices.repeat_interleave(sizes) - total_atoms = values.shape[0] - atom_within_conf = torch.arange(total_atoms, device=device, dtype=torch.int64) - atom_starts[ - :-1 - ].repeat_interleave(sizes) +class _AtomRowIndex(NamedTuple): + """Dense ``(molecule, conformer, atom)`` position of every per-atom row of a CSR conformer batch.""" + + rows: "torch.Tensor" + mol: "torch.Tensor" + conf: "torch.Tensor" + atom: "torch.Tensor" + conf_mask: "torch.Tensor" + atom_mask: "torch.Tensor" + - dense_vals[mol_idx_per_atom, conf_idx_per_atom, atom_within_conf, :] = values - atom_mask[mol_idx_per_atom, conf_idx_per_atom, atom_within_conf] = True +def _atom_row_index( + atom_starts: AsyncGpuResult, mol_indices: AsyncGpuResult, conf_indices: AsyncGpuResult, n_mols: int +) -> _AtomRowIndex: + """Map each conformer's ``atom_starts`` range to padded ``(n_mols, max_confs, max_atoms)`` slots. - return Dense3DResult(values=dense_vals, conf_mask=conf_mask, atom_mask=atom_mask) + Rows outside every conformer's range are not indexed. Reading the index tensors synchronizes implicitly. + """ + starts = atom_starts.torch().to(torch.int64) + mols = mol_indices.torch().to(torch.int64) + confs = conf_indices.torch().to(torch.int64) + device = starts.device + sizes = starts[1:] - starts[:-1] + max_confs = int(torch.bincount(mols, minlength=n_mols).max().item()) if mols.numel() else 0 + max_atoms = int(sizes.max().item()) if sizes.numel() else 0 + + conf_mask = torch.zeros((n_mols, max_confs), dtype=torch.bool, device=device) + conf_mask[mols, confs] = True + first_slot = (torch.cumsum(sizes, 0) - sizes).repeat_interleave(sizes) + atom = torch.arange(int(sizes.sum().item()), device=device, dtype=torch.int64) - first_slot + rows = starts[:-1].repeat_interleave(sizes) + atom + mol = mols.repeat_interleave(sizes) + conf = confs.repeat_interleave(sizes) + atom_mask = torch.zeros((n_mols, max_confs, max_atoms), dtype=torch.bool, device=device) + atom_mask[mol, conf, atom] = True + return _AtomRowIndex(rows=rows, mol=mol, conf=conf, atom=atom, conf_mask=conf_mask, atom_mask=atom_mask) + + +def _scatter_atom_rows(values: "torch.Tensor", index: _AtomRowIndex, pad_value: float) -> "torch.Tensor": + """Pad per-atom rows ``values`` (any trailing shape) to ``(n_mols, max_confs, max_atoms, *values.shape[1:])``.""" + dense = torch.full( + (*index.atom_mask.shape, *values.shape[1:]), pad_value, dtype=values.dtype, device=values.device + ) + dense[index.mol, index.conf, index.atom] = values[index.rows] + return dense diff --git a/src/descriptors3d.cu b/src/descriptors3d.cu index 2c18477c..e9a7eae2 100644 --- a/src/descriptors3d.cu +++ b/src/descriptors3d.cu @@ -8,6 +8,7 @@ #include "src/descriptors3d.h" #include "src/descriptors3d_dclv.cuh" +#include "src/descriptors3d_eem.cuh" #include "src/descriptors3d_getaway.cuh" #include "src/descriptors3d_kernel.cuh" #include "src/descriptors3d_moments.cuh" @@ -71,6 +72,8 @@ std::string_view property3DName(const Property3D property) { return "DCLVCompactness"; case Property3D::DCLVPackingDensity: return "DCLVPackingDensity"; + case Property3D::EEMcharges: + return "EEMcharges"; } throw std::invalid_argument("Unknown Property3D value " + std::to_string(static_cast(property))); } @@ -99,6 +102,7 @@ using descriptors3d_detail::kGroupsPerWarp; using descriptors3d_detail::kWarpSize; using descriptors3d_detail::kWarpsPerBlock; using descriptors3d_detail::launchDclvProperties; +using descriptors3d_detail::launchEemCharges; using descriptors3d_detail::launchGetawayProperties; using descriptors3d_detail::launchPairwiseProperties; using descriptors3d_detail::launchProjectionProperties; @@ -160,6 +164,7 @@ constexpr SharedStageSet directStages(const Property3D property) { case Property3D::DCLVPolarVolume: case Property3D::DCLVCompactness: case Property3D::DCLVPackingDensity: + case Property3D::EEMcharges: return 0; } return 0; @@ -328,12 +333,15 @@ Property3DResults calc3DPropertiesGpu(const DeviceCoordView& coord Real* usrcatOutput = nullptr; Real* getawayOutput = nullptr; DclvOutputs dclvOutputs; + Real* eemOutput = nullptr; for (const Property3D property : properties) { // Bounds the work arrays, which hold one slot per known property. if (std::find(kAllProperty3D.begin(), kAllProperty3D.end(), property) == kAllProperty3D.end()) { throw std::invalid_argument("Unknown Property3D value " + std::to_string(static_cast(property))); } - const size_t outputSize = static_cast(coordinates.numConformers) * property3DWidth(property); + const int64_t outputRows = + property3DExtent(property) == Property3DExtent::Atom ? coordinates.numAtoms : coordinates.numConformers; + const size_t outputSize = static_cast(outputRows) * property3DWidth(property); auto [it, inserted] = results.try_emplace(property, outputSize, stream); if (!inserted) { throw std::invalid_argument("Duplicate 3D property '" + std::string(property3DName(property)) + "'"); @@ -360,6 +368,9 @@ Property3DResults calc3DPropertiesGpu(const DeviceCoordView& coord case Property3DFamily::Dclv: dclvOutputs.values[dclvColumn(property)] = it->second.data(); break; + case Property3DFamily::Eem: + eemOutput = it->second.data(); + break; } } @@ -375,6 +386,7 @@ Property3DResults calc3DPropertiesGpu(const DeviceCoordView& coord } if (coordinates.numConformers == 0) { + launchEemCharges(coordinates, inputs, eemOutput, stream); return results; } if (coordinates.positions == nullptr || coordinates.atomStarts == nullptr || coordinates.molIndices == nullptr || @@ -406,6 +418,10 @@ Property3DResults calc3DPropertiesGpu(const DeviceCoordView& coord throw std::invalid_argument( "Van der Waals radii and polar classes must not be null when a DCLV property is requested"); } + if (eemOutput != nullptr && (inputs.eemElectronegativity == nullptr || inputs.eemHardness == nullptr || + inputs.moleculeFormalCharges == nullptr)) { + throw std::invalid_argument("EEM parameters and formal charges must not be null when EEMcharges is requested"); + } const bool momentWeightsAreUnit = inputs.momentWeights == nullptr; const bool onlySpherocity = hasSpherocity && work.numProperties == 1; @@ -422,6 +438,7 @@ Property3DResults calc3DPropertiesGpu(const DeviceCoordView& coord launchUsrProperties(coordinates, inputs, usrOutput, usrcatOutput, stream); launchGetawayProperties(coordinates, inputs, options.getaway, bondTable.table, getawayOutput, stream); launchDclvProperties(coordinates, inputs, options.dclv, dclvOutputs, stream); + launchEemCharges(coordinates, inputs, eemOutput, stream); return results; } diff --git a/src/descriptors3d.h b/src/descriptors3d.h index af3832cc..f2c21fe9 100644 --- a/src/descriptors3d.h +++ b/src/descriptors3d.h @@ -17,7 +17,7 @@ namespace nvMolKit { -//! Per-conformer 3D properties. Names match the corresponding RDKit descriptor names. +//! Conformer-dependent 3D properties. Names match the corresponding RDKit descriptor names. enum class Property3D : int { PMI1 = 0, PMI2 = 1, @@ -45,9 +45,10 @@ enum class Property3D : int { DCLVPolarVolume = 22, DCLVCompactness = 23, DCLVPackingDensity = 24, + EEMcharges = 25, }; -inline constexpr std::array kAllProperty3D = { +inline constexpr std::array kAllProperty3D = { Property3D::PMI1, Property3D::PMI2, Property3D::PMI3, @@ -73,6 +74,7 @@ inline constexpr std::array kAllProperty3D = { Property3D::DCLVPolarVolume, Property3D::DCLVCompactness, Property3D::DCLVPackingDensity, + Property3D::EEMcharges, }; inline constexpr int kNumWhimProperties = 114; @@ -89,7 +91,18 @@ std::string_view property3DName(Property3D property); //! Parse a canonical property name. @throws std::invalid_argument for unknown names. Property3D property3DFromName(std::string_view name); -//! Number of values emitted per conformer. Scalar properties have width one. +//! What one output row of a property describes. +enum class Property3DExtent : int { + Conformer, //!< One row per conformer, in coordinate-row order. + Atom, //!< One row per atom, laid out like the coordinate positions: conformer i's atoms are rows + //!< `atomStarts[i] .. atomStarts[i + 1]` of DeviceCoordView. +}; + +constexpr Property3DExtent property3DExtent(const Property3D property) { + return property == Property3D::EEMcharges ? Property3DExtent::Atom : Property3DExtent::Conformer; +} + +//! Number of values emitted per output row (see property3DExtent()). Scalar properties have width one. constexpr int property3DWidth(const Property3D property) { switch (property) { case Property3D::WHIM: @@ -120,6 +133,7 @@ enum class Property3DFamily : int { Usr, //!< Distance moments from four reference points: USR and USRCAT. Getaway, //!< Leverage (molecular influence) matrix descriptors: GETAWAY. Dclv, //!< Dot-surface areas and volumes of the atom spheres: RDKit's DoubleCubicLatticeVolume (DCLV*). + Eem, //!< Electronegativity-equalization partial charges: EEMcharges. }; constexpr Property3DFamily property3DFamily(const Property3D property) { @@ -155,6 +169,8 @@ constexpr Property3DFamily property3DFamily(const Property3D property) { case Property3D::DCLVCompactness: case Property3D::DCLVPackingDensity: return Property3DFamily::Dclv; + case Property3D::EEMcharges: + return Property3DFamily::Eem; } return Property3DFamily::Moments; } @@ -197,8 +213,8 @@ struct DclvOptions { bool includeHs = false; }; -//! Per-family options; each family reads only its own member. PBF, RDF, MORSE, AUTOCORR3D, USR and -//! USRCAT have no options. +//! Per-family options; each family reads only its own member. PBF, RDF, MORSE, AUTOCORR3D, USR, USRCAT +//! and EEMcharges have no options. struct Property3DOptions { MomentOptions moments; WhimOptions whim; @@ -248,6 +264,13 @@ struct Property3DDeviceInputs { const uint8_t* dclvPolarClasses = nullptr; //! PBF, GETAWAY: per-conformer RDKit is3D flags (one per coordinate row); null treats every row as 3D. const int8_t* conformerIs3D = nullptr; + //! EEMcharges: per atom, RDKit's EEM electronegativity A and hardness B for the atom's element and highest + //! Kekulé bond order. A non-finite A marks an atom RDKit has no parameters for and gives the conformer NaN, as + //! do coincident atoms. + const double* eemElectronegativity = nullptr; + const double* eemHardness = nullptr; + //! EEMcharges: per molecule, the sum of formal charges. + const double* moleculeFormalCharges = nullptr; //! WHIM, AUTOCORR3D, GETAWAY: largest molecule atom count in the batch (host value). Sizes WHIM's //! per-conformer symmetry-search scratch (rows with more atoms produce NaN) and, with 0 or above 3072 atoms, //! leaves AUTOCORR3D and GETAWAY without the shared bond-distance table. @@ -257,13 +280,18 @@ struct Property3DDeviceInputs { int64_t moleculeAtomPairs = 0; }; -//! One row-major device vector of length numConformers * property3DWidth(property) per property. +//! One row-major device vector per property, of length property3DWidth(property) times the number of output +//! rows: `coordinates.numConformers` for Property3DExtent::Conformer, `coordinates.numAtoms` for +//! Property3DExtent::Atom. //! @p Real is float (PrecisionMode::SINGLE) or double (PrecisionMode::FULL). template using Property3DResults = std::unordered_map>; /** * @brief Calculate the requested 3D properties for every conformer in a coordinate batch. * + * Per-atom properties (Property3DExtent::Atom) have one row per row of @ref DeviceCoordView::positions; + * rows outside every conformer's atom range hold NaN. + * * Each requested family runs as one kernel launch and returns @p Real (float or double), computing in * @p Real except for steps where float32 measurably loses accuracy, which are always float64: coordinate * centering (descriptors3d_detail::centeredPosition), WHIM's PCA (descriptors3d_detail::WhimReal), @@ -271,7 +299,8 @@ template using Property3DResults = std::unordered_map +#include +#include +#include + +#include "src/descriptors3d.h" +#include "src/descriptors3d_kernel.cuh" +#include "src/utils/cuda_error_check.h" +#include "src/utils/device_vector.h" + +namespace nvMolKit::descriptors3d_detail { + +//! RDKit's EEM kappa (Code/GraphMol/Descriptors/EEM.cpp). +constexpr double kEemKappa = 0.5125; + +constexpr int kEemBlockSize = 128; +//! Dynamic shared memory for one conformer's workspace; larger workspaces use global scratch. Leaves room for +//! the solver's static shared memory below the 48 KiB launch limit. +constexpr int kEemSharedBytes = 47 * 1024; + +__host__ __device__ constexpr int64_t alignTo8(const int64_t bytes) { + return (bytes + 7) / 8 * 8; +} + +/** + * @brief Per-conformer workspace of an @p numAtoms -atom EEM solve: the augmented (n + 1) x (n + 2) system of + * @p Real, then the float64 first solution and the pivot rows used by float32 refinement. + */ +template struct EemWorkspace { + Real* system; + double* firstSolution; + int32_t* pivots; + + __host__ __device__ static constexpr int64_t systemBytes(const int64_t numAtoms) { + return alignTo8((numAtoms + 1) * (numAtoms + 2) * static_cast(sizeof(Real))); + } + __host__ __device__ static constexpr int64_t bytes(const int64_t numAtoms) { + return systemBytes(numAtoms) + alignTo8((numAtoms + 1) * static_cast(sizeof(double) + sizeof(int32_t))); + } + __device__ static EemWorkspace at(unsigned char* base, const int numAtoms) { + EemWorkspace workspace; + workspace.system = reinterpret_cast(base); + workspace.firstSolution = reinterpret_cast(base + systemBytes(numAtoms)); + workspace.pivots = reinterpret_cast(workspace.firstSolution + numAtoms + 1); + return workspace; + } +}; + +/** + * @brief Block-collective LU factorization with partial pivoting of the augmented system @p system (rows of + * `size + 1` entries, the last column the right-hand side). Rows are swapped whole, so on return the + * strictly lower part holds the unit-lower factor's multipliers, the upper part the upper factor, the + * last column the forward-eliminated right-hand side, and `pivots[k]` the row swapped into row k. Ties + * between pivot candidates go to the lowest row. + */ +template __device__ void factorAugmentedSystem(Real* system, int32_t* pivots, const int size) { + __shared__ int pivotRow; + const int stride = size + 1; + const int lane = static_cast(threadIdx.x) % kWarpSize; + for (int k = 0; k < size; ++k) { + if (threadIdx.x < kWarpSize) { + Real bestValue = -1; + int bestRow = k; + for (int row = k + lane; row < size; row += kWarpSize) { + const Real value = fabs(system[row * stride + k]); + if (value > bestValue) { + bestValue = value; + bestRow = row; + } + } + for (int offset = kWarpSize / 2; offset > 0; offset >>= 1) { + const Real otherValue = __shfl_xor_sync(0xffffffffu, bestValue, offset); + const int otherRow = __shfl_xor_sync(0xffffffffu, bestRow, offset); + if (otherValue > bestValue || (otherValue == bestValue && otherRow < bestRow)) { + bestValue = otherValue; + bestRow = otherRow; + } + } + if (lane == 0) { + pivotRow = bestRow; + pivots[k] = bestRow; + } + } + __syncthreads(); + const int pivot = pivotRow; + if (pivot != k) { + for (int col = static_cast(threadIdx.x); col < stride; col += blockDim.x) { + const Real swapped = system[k * stride + col]; + system[k * stride + col] = system[pivot * stride + col]; + system[pivot * stride + col] = swapped; + } + __syncthreads(); + } + const Real inversePivot = Real(1) / system[k * stride + k]; + for (int row = k + 1 + static_cast(threadIdx.x); row < size; row += blockDim.x) { + system[row * stride + k] *= inversePivot; + } + __syncthreads(); + const int trailingCols = stride - k - 1; + const int trailing = (size - k - 1) * trailingCols; + for (int idx = static_cast(threadIdx.x); idx < trailing; idx += blockDim.x) { + const int row = k + 1 + idx / trailingCols; + const int col = k + 1 + idx % trailingCols; + system[row * stride + col] -= system[row * stride + k] * system[k * stride + col]; + } + __syncthreads(); + } +} + +//! Block-collective. Column-oriented forward substitution with the unit-lower factor, in place on the +//! right-hand-side column. +template __device__ void forwardSubstitute(Real* system, const int size) { + const int stride = size + 1; + for (int k = 0; k < size - 1; ++k) { + const Real solved = system[k * stride + size]; + for (int row = k + 1 + static_cast(threadIdx.x); row < size; row += blockDim.x) { + system[row * stride + size] -= system[row * stride + k] * solved; + } + __syncthreads(); + } +} + +//! Block-collective. Column-oriented back substitution with the upper factor, leaving the solution in the +//! right-hand-side column. +template __device__ void backSubstitute(Real* system, const int size) { + const int stride = size + 1; + for (int k = size - 1; k >= 0; --k) { + const Real solution = system[k * stride + size] / system[k * stride + k]; + for (int row = static_cast(threadIdx.x); row < k; row += blockDim.x) { + system[row * stride + size] -= system[row * stride + k] * solution; + } + __syncthreads(); + if (threadIdx.x == 0) { + system[k * stride + size] = solution; + } + } + __syncthreads(); +} + +//! Atoms, parameters and charge of one conformer's EEM system. +struct EemConformer { + const double* positions; + const double* electronegativity; + const double* hardness; + double formalCharge; + int numAtoms; +}; + +/** + * @brief Entry (@p row, @p col) of RDKit's EEM system as it solves it, the right-hand side in column + * `numAtoms + 1`. Sets @p coincident when an atom pair's squared distance is not positive. + * + * RDKit fills the system row-major and hands it to Eigen as column-major, so it solves the transpose of the + * matrix its code describes: row i < n is `sum_j J_ij q_j + chi = -A_i` with `J_ii = B_i` and + * `J_ij = kappa / r_ij`, and the last row is `-sum_j q_j = Q`. The charges therefore sum to the negated + * formal charge, which this reproduces. + * + * FP64 required: the coordinate differences, before the conversion to @p T (see centeredPosition()). + */ +template +__device__ __forceinline__ T +eemSystemEntry(const EemConformer& conformer, const int row, const int col, bool& coincident) { + const int numAtoms = conformer.numAtoms; + if (col == numAtoms + 1) { + return row < numAtoms ? static_cast(-conformer.electronegativity[row]) : static_cast(conformer.formalCharge); + } + if (row == numAtoms) { + return col < numAtoms ? T(-1) : T(0); + } + if (col == numAtoms) { + return T(1); + } + if (row == col) { + return static_cast(conformer.hardness[row]); + } + const double* positions = conformer.positions; + const T dx = static_cast(positions[row * 3 + 0] - positions[col * 3 + 0]); + const T dy = static_cast(positions[row * 3 + 1] - positions[col * 3 + 1]); + const T dz = static_cast(positions[row * 3 + 2] - positions[col * 3 + 2]); + const T squared = dx * dx + dy * dy + dz * dz; + coincident |= !(squared > T(0)); + return static_cast(kEemKappa) / sqrt(squared); +} + +/** + * @brief Block-collective. Replaces the float32 solution in the factored @p workspace system's right-hand-side + * column by the correction of one iterative-refinement step, keeping the solution in + * `workspace.firstSolution`. + * + * FP64 required: the residual. float32 elimination of the bordered system accumulates rounding that grows with + * the atom count, ~3e-4 absolute on 160-340-atom conformers against ~5e-7 from rounding the entries alone + * (measured); one refinement step with the float32 factors and a float64 residual recovers ~5e-10. + */ +__device__ inline void refineEemSolution(const EemConformer& conformer, const EemWorkspace& workspace) { + const int size = conformer.numAtoms + 1; + const int stride = size + 1; + float* system = workspace.system; + for (int row = static_cast(threadIdx.x); row < size; row += blockDim.x) { + workspace.firstSolution[row] = system[row * stride + size]; + } + __syncthreads(); + // Residual b - A x in the original row order, written to each row's right-hand side. + for (int row = static_cast(threadIdx.x); row < size; row += blockDim.x) { + bool ignored = false; + double residual = eemSystemEntry(conformer, row, size, ignored); + for (int col = 0; col < size; ++col) { + residual -= eemSystemEntry(conformer, row, col, ignored) * workspace.firstSolution[col]; + } + system[row * stride + size] = static_cast(residual); + } + __syncthreads(); + if (threadIdx.x == 0) { + for (int k = 0; k < size; ++k) { + const int pivot = workspace.pivots[k]; + if (pivot != k) { + const float swapped = system[k * stride + size]; + system[k * stride + size] = system[pivot * stride + size]; + system[pivot * stride + size] = swapped; + } + } + } + __syncthreads(); + forwardSubstitute(system, size); + backSubstitute(system, size); +} + +/** + * @brief One block per conformer (grid-stride). Builds RDKit's EEM system (see eemSystemEntry()) and writes + * one partial charge per atom to @p output, laid out like the coordinate positions. + */ +template +__global__ void __launch_bounds__(kEemBlockSize) eemChargesKernel(const DeviceCoordView coordinates, + const Property3DDeviceInputs inputs, + const int64_t sharedCapacity, + unsigned char* globalScratch, + const int64_t scratchStride, + Real* output) { + extern __shared__ __align__(16) unsigned char eemSharedBytes[]; + for (int conformerIdx = blockIdx.x; conformerIdx < coordinates.numConformers; conformerIdx += gridDim.x) { + const int64_t atomStart = coordinates.atomStarts[conformerIdx]; + const int64_t atomStop = coordinates.atomStarts[conformerIdx + 1]; + if (atomStart < 0 || atomStop <= atomStart || atomStop > coordinates.numAtoms) { + continue; // No rows to write; the output was initialized to NaN. + } + const int numAtoms = static_cast(atomStop - atomStart); + const int moleculeIdx = coordinates.molIndices[conformerIdx]; + bool valid = moleculeIdx >= 0 && moleculeIdx < coordinates.nMols && numAtoms <= inputs.maxMoleculeAtoms; + int paramStart = 0; + if (valid) { + paramStart = inputs.moleculeAtomStarts[moleculeIdx]; + valid = inputs.moleculeAtomStarts[moleculeIdx + 1] - paramStart == numAtoms; + } + bool missingParameters = false; + for (int atomIdx = static_cast(threadIdx.x); valid && atomIdx < numAtoms; atomIdx += blockDim.x) { + missingParameters |= !isfinite(inputs.eemElectronegativity[paramStart + atomIdx]); + } + valid = valid && !__syncthreads_or(missingParameters); + + bool coincident = false; + EemConformer conformer{}; + EemWorkspace workspace{}; + const int size = numAtoms + 1; + const int stride = size + 1; + if (valid) { + conformer = EemConformer{coordinates.positions + atomStart * 3, + inputs.eemElectronegativity + paramStart, + inputs.eemHardness + paramStart, + inputs.moleculeFormalCharges[moleculeIdx], + numAtoms}; + workspace = EemWorkspace::at(EemWorkspace::bytes(numAtoms) <= sharedCapacity ? + eemSharedBytes : + globalScratch + blockIdx.x * scratchStride, + numAtoms); + for (int64_t idx = threadIdx.x; idx < static_cast(size) * stride; idx += blockDim.x) { + workspace.system[idx] = + eemSystemEntry(conformer, static_cast(idx / stride), static_cast(idx % stride), coincident); + } + } + // Coincident atoms (or non-finite coordinates) put an infinite kappa / r in the system, whose solution is + // meaningless; RDKit's full-pivot LU returns all-zero charges for it. + valid = valid && !__syncthreads_or(coincident); + if (!valid) { + for (int atomIdx = static_cast(threadIdx.x); atomIdx < numAtoms; atomIdx += blockDim.x) { + output[atomStart + atomIdx] = static_cast(nan("")); + } + __syncthreads(); + continue; + } + + factorAugmentedSystem(workspace.system, workspace.pivots, size); + backSubstitute(workspace.system, size); + if constexpr (std::is_same_v) { + refineEemSolution(conformer, workspace); + } + for (int atomIdx = static_cast(threadIdx.x); atomIdx < numAtoms; atomIdx += blockDim.x) { + if constexpr (std::is_same_v) { + output[atomStart + atomIdx] = static_cast( + workspace.firstSolution[atomIdx] + static_cast(workspace.system[atomIdx * stride + size])); + } else { + output[atomStart + atomIdx] = workspace.system[atomIdx * stride + size]; + } + } + // The next conformer reuses the workspace. + __syncthreads(); + } +} + +template +void launchEemCharges(const DeviceCoordView& coordinates, + const Property3DDeviceInputs& inputs, + Real* output, + const cudaStream_t stream) { + if (output == nullptr) { + return; + } + // Bytes of 0xFF are NaN for float and double: rows outside every conformer's atom range stay NaN. + cudaCheckError(cudaMemsetAsync(output, 0xFF, static_cast(coordinates.numAtoms) * sizeof(Real), stream)); + if (coordinates.numConformers == 0) { + return; + } + + int device = 0; + int numSms = 0; + cudaCheckError(cudaGetDevice(&device)); + cudaCheckError(cudaDeviceGetAttribute(&numSms, cudaDevAttrMultiProcessorCount, device)); + const int numBlocks = std::min(coordinates.numConformers, numSms * 16); + const int64_t largestWorkspace = EemWorkspace::bytes(inputs.maxMoleculeAtoms); + const int64_t sharedCapacity = std::min(largestWorkspace, kEemSharedBytes); + const int64_t scratchStride = largestWorkspace > sharedCapacity ? largestWorkspace : 0; + AsyncDeviceVector scratch(static_cast(scratchStride * numBlocks), stream); + eemChargesKernel<<>>(coordinates, + inputs, + sharedCapacity, + scratch.data(), + scratchStride, + output); + cudaCheckError(cudaGetLastError()); +} + +} // namespace nvMolKit::descriptors3d_detail + +#endif // NVMOLKIT_DESCRIPTORS3D_EEM_CUH diff --git a/src/descriptors3d_mol.cpp b/src/descriptors3d_mol.cpp index 3c5fc3b3..4a1c87b4 100644 --- a/src/descriptors3d_mol.cpp +++ b/src/descriptors3d_mol.cpp @@ -6,7 +6,10 @@ #include #include #include +#include #include +#include +#include #include #include @@ -39,6 +42,9 @@ struct DeviceDescriptorInputs { AsyncDeviceVector dclvPolarClasses; AsyncDeviceVector defaultConformerRows; AsyncDeviceVector conformerIs3D; + AsyncDeviceVector eemElectronegativity; + AsyncDeviceVector eemHardness; + AsyncDeviceVector moleculeFormalCharges; AsyncDeviceVector moleculeAtomStarts; }; @@ -54,8 +60,75 @@ struct DescriptorInputNeeds { bool dclvAtomInputs = false; bool defaultConformerRows = false; bool conformerFlags = false; + bool eemParameters = false; }; +//! RDKit's EEM parameters, indexed by atomic number up to bromine, for atoms whose highest Kekulé bond order +//! is 1, 2 and 3 (Code/GraphMol/Descriptors/EEM.cpp, from the NEEMP B3LYP/6-311G NPA set). A 0 entry marks an +//! element and type the set has no parameters for. +constexpr int kNumEemElements = 36; +// clang-format off +constexpr std::array, 3> kEemElectronegativity = {{ + {0.0, 2.5473, 0.0, 0.0, 0.0, 0.0, 2.7221, 2.9750, 3.1503, 2.9976, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 2.6511, 2.7026, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 2.6263}, + {0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 2.7667, 2.8895, 3.0486, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 2.2933, + 2.6471, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0}, + {0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 2.6944, 3.0240, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0}, +}}; +constexpr std::array, 3> kEemHardness = {{ + {0.0, 1.1641, 0.0, 0.0, 0.0, 0.0, 0.6403, 0.9083, 1.0577, 0.9983, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.4897, 1.1537, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.1105}, + {0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.6513, 0.6647, 0.8410, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.5759, + 0.4512, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0}, + {0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.6776, 1.4240, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0}, +}}; +// clang-format on + +/** + * @brief Writes each atom's EEM electronegativity and hardness, and returns the molecule's formal charge, + * as RDKit's EEM types them: on a Kekulé copy, by element and highest bond order (aromatic counted + * as 2). + * + * Atoms the parameter set does not cover get a NaN electronegativity so the kernel reports NaN, as do molecules + * that fail to kekulize (which RDKit rejects). For elements beyond bromine RDKit reads past its tables, and for + * bond orders above 3 it leaves the right-hand side uninitialized. Unlike RDKit, which silently solves with zero + * parameters for uncovered atoms within its tables (Si, single-bonded P, B, ...) and returns charges of -10 or + * beyond, those atoms are NaN too. + */ +double writeEemParameters(const RDKit::ROMol& mol, double* electronegativity, double* hardness) { + const double missing = std::numeric_limits::quiet_NaN(); + RDKit::RWMol kekule(mol); + try { + RDKit::MolOps::Kekulize(kekule, true); + } catch (const RDKit::MolSanitizeException&) { + std::fill(electronegativity, electronegativity + mol.getNumAtoms(), missing); + std::fill(hardness, hardness + mol.getNumAtoms(), missing); + return 0.0; + } + for (const auto* atom : kekule.atoms()) { + unsigned int type = 1; + for (const auto* bond : kekule.atomBonds(atom)) { + double order = bond->getBondTypeAsDouble(); + if (order == 1.5) { + order = 2.0; + } + type = std::max(type, static_cast(order)); + } + const unsigned int element = atom->getAtomicNum(); + const unsigned int idx = atom->getIdx(); + if (element >= kNumEemElements || type > 3 || kEemHardness[type - 1][element] == 0.0) { + electronegativity[idx] = missing; + hardness[idx] = missing; + } else { + electronegativity[idx] = kEemElectronegativity[type - 1][element]; + hardness[idx] = kEemHardness[type - 1][element]; + } + } + return RDKit::MolOps::getFormalCharge(kekule); +} + //! Writes molecule @p mol's bond adjacency into the global CSR arrays: atom i's neighbors go to //! @p neighbors from `neighborStarts[i]`, as atom indices within the molecule. @p neighborStarts points at the //! molecule's first atom; its entries are global offsets beginning at @p neighborOffset. @@ -190,6 +263,9 @@ DeviceDescriptorInputs uploadDescriptorInputs(const std::vector eemElectronegativity(needs.eemParameters ? static_cast(totalAtoms) : 0); + std::vector eemHardness(needs.eemParameters ? static_cast(totalAtoms) : 0); + std::vector moleculeFormalCharges(needs.eemParameters ? numMols : 0); std::vector conformerIs3D; if (needs.conformerFlags) { for (const RDKit::ROMol* mol : mols) { @@ -215,7 +291,7 @@ DeviceDescriptorInputs uploadDescriptorInputs(const std::vector(polarClasses.size(), stream), AsyncDeviceVector(defaultConformerRows.size(), stream), AsyncDeviceVector(conformerIs3D.size(), stream), + AsyncDeviceVector(eemElectronegativity.size(), stream), + AsyncDeviceVector(eemHardness.size(), stream), + AsyncDeviceVector(moleculeFormalCharges.size(), stream), AsyncDeviceVector(atomStarts.size(), stream)}; if (!weights.empty()) { result.momentWeights.copyFromHost(weights); @@ -342,6 +428,13 @@ DeviceDescriptorInputs uploadDescriptorInputs(const std::vector calc3DProperties(const std::vector calc3DProperties(const std::vectorgetNumAtoms(); inputs.maxMoleculeAtoms = std::max(inputs.maxMoleculeAtoms, static_cast(numAtoms)); @@ -419,6 +516,7 @@ Property3DBatchResult calc3DProperties(const std::vector(view, inputs, properties, options, stream); result.molIndices = std::move(uploaded.molIndices); result.confIndices = std::move(uploaded.confIndices); + result.atomStarts = std::move(uploaded.atomStarts); return result; } diff --git a/src/descriptors3d_mol.h b/src/descriptors3d_mol.h index 20bd975f..25e4cae5 100644 --- a/src/descriptors3d_mol.h +++ b/src/descriptors3d_mol.h @@ -20,11 +20,12 @@ namespace nvMolKit { //! Properties for a molecule batch together with the labels of each output row. template struct Property3DBatchResult { Property3DResults properties; - //! Row -> input molecule index and per-molecule conformer position. Populated only when - //! coordinates were taken from the molecules; otherwise rows follow the caller's DeviceCoordView - //! and these are empty. + //! Conformer row -> input molecule index and per-molecule conformer position, and each conformer's range of + //! per-atom rows (`numConformers + 1` offsets). Populated only when coordinates were taken from the + //! molecules; otherwise rows follow the caller's DeviceCoordView and these are empty. AsyncDeviceVector molIndices; AsyncDeviceVector confIndices; + AsyncDeviceVector atomStarts; }; /** diff --git a/tests/test_descriptors3d.cpp b/tests/test_descriptors3d.cpp index 9c958954..4f690ce6 100644 --- a/tests/test_descriptors3d.cpp +++ b/tests/test_descriptors3d.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include #include #include @@ -13,6 +14,7 @@ #include #include #include +#include #include #include @@ -493,6 +495,135 @@ TEST(Descriptors3DPairwise, BondSearchFallbackMatchesBondDistanceTable) { } } +//! Deterministic, well-separated coordinates for @p numAtoms atoms along a helix. +std::vector helixPoints(const int numAtoms, const double phase) { + std::vector points; + for (int atomIdx = 0; atomIdx < numAtoms; ++atomIdx) { + const double angle = 1.9 * atomIdx + phase; + points.push_back({1.4 * std::cos(angle), 1.4 * std::sin(angle), 0.55 * atomIdx + 0.1 * phase}); + } + return points; +} + +//! Per-atom EEM charges of every conformer, compared against RDKit's EEM through the batch's atomStarts. +void expectEemMatchesRdkit(const std::vector& mols, const double tolerance) { + const std::vector constMols(mols.begin(), mols.end()); + auto results = nvMolKit::calc3DProperties(constMols, {Property3D::EEMcharges}, {}, nullptr); + const auto values = toHost(results.properties.at(Property3D::EEMcharges)); + const auto atomStarts = toHost(results.atomStarts); + int row = 0; + for (RDKit::ROMol* mol : mols) { + for (int confIdx = 0; confIdx < static_cast(mol->getNumConformers()); ++confIdx, ++row) { + std::vector expected; + RDKit::Descriptors::EEM(*mol, expected, confIdx); + ASSERT_EQ(atomStarts[row + 1] - atomStarts[row], static_cast(expected.size())); + for (size_t atomIdx = 0; atomIdx < expected.size(); ++atomIdx) { + EXPECT_NEAR(values[atomStarts[row] + atomIdx], expected[atomIdx], tolerance) + << RDKit::MolToSmiles(*mol) << " conformer " << confIdx << ", atom " << atomIdx; + } + } + } + EXPECT_EQ(static_cast(atomStarts.back()), values.size()); +} + +TEST(Descriptors3DEem, MatchesRdkitForChargedAndAromaticMolecules) { + std::vector> owned; + for (const char* smiles : {"CC(=O)[O-]", "C[NH3+]", "c1ccncc1O", "O=c1cc[nH]cc1", "C#N", "CS(=O)(=O)NCl"}) { + std::unique_ptr parsed(RDKit::SmilesToMol(smiles)); + const int numAtoms = static_cast(parsed->getNumAtoms()); + owned.push_back(molWithConformers(smiles, {helixPoints(numAtoms, 0.0), helixPoints(numAtoms, 0.7)})); + } + std::vector mols; + for (const auto& mol : owned) { + mols.push_back(mol.get()); + } + expectEemMatchesRdkit(mols, 1e-9); +} + +TEST(Descriptors3DEem, LargeMoleculesUseGlobalScratchAndMatchRdkit) { + // 120 atoms exceed the shared-memory system in both precisions; the small molecule stays in shared memory. + auto large = + molWithConformers(("OC" + std::string(118, 'C')).c_str(), {helixPoints(120, 0.0), helixPoints(120, 0.3)}); + auto small = molWithConformers("CCO", {helixPoints(3, 0.0)}); + expectEemMatchesRdkit({large.get(), small.get()}, 1e-8); +} + +TEST(Descriptors3DEem, UnparameterizedAtomsGiveNanOnlyForTheirMolecule) { + // RDKit has no iodine parameters (it reads past its tables), and a lone aromatic ring atom cannot be kekulized. + auto iodide = molWithConformers("CCI", {helixPoints(3, 0.0)}); + auto ethanol = molWithConformers("CCO", {helixPoints(3, 0.0)}); + RDKit::SmilesParserParams params; + params.sanitize = false; + std::unique_ptr aromatic(RDKit::SmilesToMol("c1cccc1", params)); + auto conformer = std::make_unique(aromatic->getNumAtoms()); + const auto points = helixPoints(5, 0.0); + for (int atomIdx = 0; atomIdx < 5; ++atomIdx) { + conformer->setAtomPos(atomIdx, RDGeom::Point3D(points[atomIdx][0], points[atomIdx][1], points[atomIdx][2])); + } + aromatic->addConformer(conformer.release(), true); + + const std::vector mols = {iodide.get(), ethanol.get(), aromatic.get()}; + auto results = nvMolKit::calc3DProperties(mols, {Property3D::EEMcharges}, {}, nullptr); + const auto values = toHost(results.properties.at(Property3D::EEMcharges)); + ASSERT_EQ(values.size(), 11u); + std::vector expected; + RDKit::Descriptors::EEM(*ethanol, expected, 0); + for (int atomIdx = 0; atomIdx < 11; ++atomIdx) { + if (atomIdx >= 3 && atomIdx < 6) { + EXPECT_NEAR(values[atomIdx], expected[atomIdx - 3], 1e-9); + } else { + EXPECT_TRUE(isNanBits(values[atomIdx])) << "atom row " << atomIdx; + } + } +} + +TEST(Descriptors3DEem, CoincidentAtomsGiveNan) { + // Two counter-ions at the same position make kappa / r infinite; RDKit returns all-zero charges for it. + auto points = helixPoints(5, 0.0); + points[4] = points[3]; + auto salt = molWithConformers("CC[NH3+].[Br-].[Br-]", {points, helixPoints(5, 0.4)}); + const std::vector mols = {salt.get()}; + auto results = nvMolKit::calc3DProperties(mols, {Property3D::EEMcharges}, {}, nullptr); + const auto values = toHost(results.properties.at(Property3D::EEMcharges)); + ASSERT_EQ(values.size(), 10u); + std::vector expected; + RDKit::Descriptors::EEM(*salt, expected, 1); + for (int atomIdx = 0; atomIdx < 5; ++atomIdx) { + EXPECT_TRUE(isNanBits(values[atomIdx])) << "atom " << atomIdx; + EXPECT_NEAR(values[5 + atomIdx], expected[atomIdx], 1e-9); + } +} + +TEST(Descriptors3DEem, DeviceRowsOutsideConformersAreNan) { + auto mol = molWithConformers("CCO", {helixPoints(3, 0.0), helixPoints(3, 0.5)}); + const std::vector mols = {mol.get()}; + auto uploaded = nvMolKit::uploadConformerCoordinates(mols, nullptr, 1); + // Six coordinate rows: conformer 0 reads rows 0-2, conformer 1's 2-atom range (rows 3-4) disagrees with the + // 3-atom molecule so its rows are NaN, and row 5 belongs to no conformer. + nvMolKit::AsyncDeviceVector positions(18, nullptr); + nvMolKit::AsyncDeviceVector atomStarts(3, nullptr); + std::vector hostPositions = toHost(uploaded.positions); + hostPositions.resize(18, 0.0); + positions.copyFromHost(hostPositions); + atomStarts.copyFromHost(std::vector{0, 3, 5}); + nvMolKit::DeviceCoordView view = nvMolKit::makeDeviceCoordView(uploaded); + view.positions = positions.data(); + view.atomStarts = atomStarts.data(); + view.numAtoms = 6; + + auto results = nvMolKit::calc3DProperties(mols, {Property3D::EEMcharges}, {}, nullptr, &view); + const auto values = toHost(results.properties.at(Property3D::EEMcharges)); + ASSERT_EQ(values.size(), 6u); + std::vector expected; + RDKit::Descriptors::EEM(*mol, expected, 0); + for (int atomIdx = 0; atomIdx < 3; ++atomIdx) { + EXPECT_NEAR(values[atomIdx], expected[atomIdx], 1e-9); + } + for (int row = 3; row < 6; ++row) { + EXPECT_TRUE(isNanBits(values[row])) << "row " << row; + } +} + TEST(Descriptors3DUsr, MatchesRdkitUsrAndUsrcat) { auto mol = molWithConformers("CCCO", { From e25ed9fe0d4554eaf52370dad1efbae9cfed0588 Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Mon, 19 Oct 2026 13:36:37 -0400 Subject: [PATCH 2/8] Validate EEM coordinates and complete per-atom descriptor coverage --- benchmarks/descriptors3d_bench.py | 22 +++++++++++++-- nvmolkit/descriptors3d.py | 7 ++--- nvmolkit/tests/test_descriptors3d.py | 41 ++++++++++++++++++++++++++++ nvmolkit/types.py | 16 ++++++----- src/descriptors3d.h | 3 +- src/descriptors3d_eem.cuh | 24 ++++++++-------- 6 files changed, 88 insertions(+), 25 deletions(-) diff --git a/benchmarks/descriptors3d_bench.py b/benchmarks/descriptors3d_bench.py index cb224989..85afedb8 100644 --- a/benchmarks/descriptors3d_bench.py +++ b/benchmarks/descriptors3d_bench.py @@ -79,6 +79,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: @@ -103,7 +122,7 @@ def _calc_rdkit_property(mol: Chem.Mol, conf_id: int, prop: Property3D) -> float if prop == Property3D.USRCAT: return rdMolDescriptors.GetUSRCAT(mol, confId=conf_id) if prop == Property3D.EEM_CHARGES: - return rdMolDescriptors.CalcEEMcharges(mol, confId=conf_id) + 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. @@ -209,7 +228,6 @@ def _validate( np.testing.assert_allclose(actual, expected[prop], rtol=DCLV_RELATIVE_TOLERANCES[precision], atol=atol) continue if prop == Property3D.EEM_CHARGES: - # Atoms RDKit has no parameters for are NaN in both. eem_atol = EEM_TOLERANCES[precision] if tolerance is None else tolerance np.testing.assert_allclose(actual, expected[prop], rtol=0, atol=eem_atol, equal_nan=True) continue diff --git a/nvmolkit/descriptors3d.py b/nvmolkit/descriptors3d.py index de3834f8..e4e4bde1 100644 --- a/nvmolkit/descriptors3d.py +++ b/nvmolkit/descriptors3d.py @@ -62,8 +62,7 @@ class Property3D(Enum): - ``EEMcharges``: RDKit's ``CalcEEMcharges`` electronegativity-equalization partial charges, one value per atom. As in RDKit, a molecule's charges sum to the negative of its formal charge. Conformers containing an - element beyond bromine or a bond order above 3, for which RDKit reads - undefined parameters, conformers with coincident atoms (or non-finite + element or bond type without nonzero parameters, conformers with coincident atoms (or non-finite coordinates), for which RDKit returns all-zero charges, and molecules that cannot be kekulized, which RDKit rejects, produce NaN. """ @@ -208,7 +207,7 @@ class Dense3DPropertyResult: 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 @@ -223,7 +222,7 @@ class Device3DPropertyResult(Mapping[str, AsyncGpuResult]): molecule's atom order. With ``Device3DResult`` coordinates, rows match its ``values`` rows one-to-one, and rows outside every conformer's range hold NaN. - Row ``i`` of every property belongs to conformer ``conf_indices[i]`` of input molecule + 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 diff --git a/nvmolkit/tests/test_descriptors3d.py b/nvmolkit/tests/test_descriptors3d.py index ea1cec60..37ebb252 100644 --- a/nvmolkit/tests/test_descriptors3d.py +++ b/nvmolkit/tests/test_descriptors3d.py @@ -1074,3 +1074,44 @@ def test_eem_atoms_with_zero_parameters_give_nan_unlike_rdkit(precision): assert np.isnan(charges[:uncovered_rows]).all() assert np.abs(_rdkit_eem_rows(uncovered)).max() > 5 # RDKit's zero-parameter charges _assert_eem_matches_rdkit(charges[uncovered_rows:], _rdkit_eem_rows([ethanol]), precision) + + +@pytest.mark.parametrize("precision", PRECISIONS) +@pytest.mark.parametrize("invalid", [float("nan"), float("inf"), -float("inf")]) +@pytest.mark.parametrize("smiles", ["[Br-]", "CCO"]) +def test_eem_nonfinite_coordinates_invalidate_only_their_conformer(precision, invalid, smiles): + molecule = Chem.MolFromSmiles(smiles) + for _ in range(2): + conformer = Chem.Conformer(molecule.GetNumAtoms()) + for atom_idx in range(molecule.GetNumAtoms()): + conformer.SetAtomPosition(atom_idx, Point3D(1.5 * atom_idx, 0, 0)) + molecule.AddConformer(conformer, assignId=True) + molecule.GetConformer(0).SetAtomPosition(0, Point3D(invalid, 0, 0)) + result = Calc3DProperties([molecule], Property3D.EEM_CHARGES, precision=precision) + charges = result[Property3D.EEM_CHARGES].numpy() + assert np.isnan(charges[: molecule.GetNumAtoms()]).all() + expected = np.asarray(rdMolDescriptors.CalcEEMcharges(molecule, confId=1)) + _assert_eem_matches_rdkit(charges[molecule.GetNumAtoms() :], expected, precision) + + +@pytest.mark.parametrize("precision", PRECISIONS) +def test_eem_global_workspace_and_custom_stream(precision): + molecule = Chem.MolFromSmiles("O" + "C" * 119) + conformer = Chem.Conformer(molecule.GetNumAtoms()) + for atom_idx in range(molecule.GetNumAtoms()): + angle = 1.9 * atom_idx + conformer.SetAtomPosition(atom_idx, Point3D(1.4 * np.cos(angle), 1.4 * np.sin(angle), 0.55 * atom_idx)) + molecule.AddConformer(conformer, assignId=True) + stream = torch.cuda.Stream() + result = Calc3DProperties([molecule], Property3D.EEM_CHARGES, precision=precision, stream=stream) + _assert_eem_matches_rdkit(result[Property3D.EEM_CHARGES].numpy(), _rdkit_eem_rows([molecule]), precision) + + +@pytest.mark.parametrize("precision", PRECISIONS) +def test_eem_empty_batches_and_molecules_without_conformers(precision): + for molecules in ([], [Chem.MolFromSmiles("CCO")]): + result = Calc3DProperties(molecules, Property3D.EEM_CHARGES, precision=precision) + assert result.n_conformers == 0 + assert result[Property3D.EEM_CHARGES].torch().numel() == 0 + dense = result.dense() + assert dense.values[Property3D.EEM_CHARGES.value].shape == (len(molecules), 0, 0) diff --git a/nvmolkit/types.py b/nvmolkit/types.py index 143ed072..6769c6dd 100644 --- a/nvmolkit/types.py +++ b/nvmolkit/types.py @@ -534,8 +534,8 @@ class _AtomRowIndex(NamedTuple): """Dense ``(molecule, conformer, atom)`` position of every per-atom row of a CSR conformer batch.""" rows: "torch.Tensor" - mol: "torch.Tensor" - conf: "torch.Tensor" + molecule: "torch.Tensor" + conformer: "torch.Tensor" atom: "torch.Tensor" conf_mask: "torch.Tensor" atom_mask: "torch.Tensor" @@ -561,11 +561,13 @@ def _atom_row_index( first_slot = (torch.cumsum(sizes, 0) - sizes).repeat_interleave(sizes) atom = torch.arange(int(sizes.sum().item()), device=device, dtype=torch.int64) - first_slot rows = starts[:-1].repeat_interleave(sizes) + atom - mol = mols.repeat_interleave(sizes) - conf = confs.repeat_interleave(sizes) + molecule = mols.repeat_interleave(sizes) + conformer = confs.repeat_interleave(sizes) atom_mask = torch.zeros((n_mols, max_confs, max_atoms), dtype=torch.bool, device=device) - atom_mask[mol, conf, atom] = True - return _AtomRowIndex(rows=rows, mol=mol, conf=conf, atom=atom, conf_mask=conf_mask, atom_mask=atom_mask) + atom_mask[molecule, conformer, atom] = True + return _AtomRowIndex( + rows=rows, molecule=molecule, conformer=conformer, atom=atom, conf_mask=conf_mask, atom_mask=atom_mask + ) def _scatter_atom_rows(values: "torch.Tensor", index: _AtomRowIndex, pad_value: float) -> "torch.Tensor": @@ -573,5 +575,5 @@ def _scatter_atom_rows(values: "torch.Tensor", index: _AtomRowIndex, pad_value: dense = torch.full( (*index.atom_mask.shape, *values.shape[1:]), pad_value, dtype=values.dtype, device=values.device ) - dense[index.mol, index.conf, index.atom] = values[index.rows] + dense[index.molecule, index.conformer, index.atom] = values[index.rows] return dense diff --git a/src/descriptors3d.h b/src/descriptors3d.h index f2c21fe9..e2ae6522 100644 --- a/src/descriptors3d.h +++ b/src/descriptors3d.h @@ -271,7 +271,8 @@ struct Property3DDeviceInputs { const double* eemHardness = nullptr; //! EEMcharges: per molecule, the sum of formal charges. const double* moleculeFormalCharges = nullptr; - //! WHIM, AUTOCORR3D, GETAWAY: largest molecule atom count in the batch (host value). Sizes WHIM's + //! WHIM, AUTOCORR3D, GETAWAY, EEMcharges: largest molecule atom count in the batch (host value). Sizes EEM's + //! linear-system scratch and WHIM's //! per-conformer symmetry-search scratch (rows with more atoms produce NaN) and, with 0 or above 3072 atoms, //! leaves AUTOCORR3D and GETAWAY without the shared bond-distance table. int32_t maxMoleculeAtoms = 0; diff --git a/src/descriptors3d_eem.cuh b/src/descriptors3d_eem.cuh index 1ef95beb..189add48 100644 --- a/src/descriptors3d_eem.cuh +++ b/src/descriptors3d_eem.cuh @@ -43,7 +43,7 @@ template struct EemWorkspace { __host__ __device__ static constexpr int64_t bytes(const int64_t numAtoms) { return systemBytes(numAtoms) + alignTo8((numAtoms + 1) * static_cast(sizeof(double) + sizeof(int32_t))); } - __device__ static EemWorkspace at(unsigned char* base, const int numAtoms) { + __device__ __forceinline__ static EemWorkspace at(unsigned char* base, const int numAtoms) { EemWorkspace workspace; workspace.system = reinterpret_cast(base); workspace.firstSolution = reinterpret_cast(base + systemBytes(numAtoms)); @@ -59,7 +59,8 @@ template struct EemWorkspace { * last column the forward-eliminated right-hand side, and `pivots[k]` the row swapped into row k. Ties * between pivot candidates go to the lowest row. */ -template __device__ void factorAugmentedSystem(Real* system, int32_t* pivots, const int size) { +template +__device__ __forceinline__ void factorAugmentedSystem(Real* system, int32_t* pivots, const int size) { __shared__ int pivotRow; const int stride = size + 1; const int lane = static_cast(threadIdx.x) % kWarpSize; @@ -115,7 +116,7 @@ template __device__ void factorAugmentedSystem(Real* system, int //! Block-collective. Column-oriented forward substitution with the unit-lower factor, in place on the //! right-hand-side column. -template __device__ void forwardSubstitute(Real* system, const int size) { +template __device__ __forceinline__ void forwardSubstitute(Real* system, const int size) { const int stride = size + 1; for (int k = 0; k < size - 1; ++k) { const Real solved = system[k * stride + size]; @@ -128,7 +129,7 @@ template __device__ void forwardSubstitute(Real* system, const i //! Block-collective. Column-oriented back substitution with the upper factor, leaving the solution in the //! right-hand-side column. -template __device__ void backSubstitute(Real* system, const int size) { +template __device__ __forceinline__ void backSubstitute(Real* system, const int size) { const int stride = size + 1; for (int k = size - 1; k >= 0; --k) { const Real solution = system[k * stride + size] / system[k * stride + k]; @@ -193,11 +194,9 @@ eemSystemEntry(const EemConformer& conformer, const int row, const int col, bool * column by the correction of one iterative-refinement step, keeping the solution in * `workspace.firstSolution`. * - * FP64 required: the residual. float32 elimination of the bordered system accumulates rounding that grows with - * the atom count, ~3e-4 absolute on 160-340-atom conformers against ~5e-7 from rounding the entries alone - * (measured); one refinement step with the float32 factors and a float64 residual recovers ~5e-10. + * FP64 residuals correct float32 elimination error using the existing factors. */ -__device__ inline void refineEemSolution(const EemConformer& conformer, const EemWorkspace& workspace) { +__device__ __forceinline__ void refineEemSolution(const EemConformer& conformer, const EemWorkspace& workspace) { const int size = conformer.numAtoms + 1; const int stride = size + 1; float* system = workspace.system; @@ -256,11 +255,14 @@ __global__ void __launch_bounds__(kEemBlockSize) eemChargesKernel(const DeviceCo paramStart = inputs.moleculeAtomStarts[moleculeIdx]; valid = inputs.moleculeAtomStarts[moleculeIdx + 1] - paramStart == numAtoms; } - bool missingParameters = false; + bool invalidAtom = false; for (int atomIdx = static_cast(threadIdx.x); valid && atomIdx < numAtoms; atomIdx += blockDim.x) { - missingParameters |= !isfinite(inputs.eemElectronegativity[paramStart + atomIdx]); + invalidAtom |= !isfinite(inputs.eemElectronegativity[paramStart + atomIdx]); + for (int axis = 0; axis < 3; ++axis) { + invalidAtom |= !isfinite(coordinates.positions[(atomStart + atomIdx) * 3 + axis]); + } } - valid = valid && !__syncthreads_or(missingParameters); + valid = valid && !__syncthreads_or(invalidAtom); bool coincident = false; EemConformer conformer{}; From 14c6e046c3c636dd5687e1f123135d79a4928453 Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Mon, 19 Oct 2026 13:40:49 -0400 Subject: [PATCH 3/8] Reserve EEM property identifier after DCLV descriptors --- tests/test_descriptors3d.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_descriptors3d.cpp b/tests/test_descriptors3d.cpp index 4f690ce6..d02848c0 100644 --- a/tests/test_descriptors3d.cpp +++ b/tests/test_descriptors3d.cpp @@ -22,6 +22,7 @@ #include #include #include +#include #include #include #include @@ -758,7 +759,7 @@ TEST_F(Descriptors3DTest, RejectsUnknownPropertyValue) { Property3D::PMI2, Property3D::PMI3, Property3D::RadiusOfGyration, - static_cast(nvMolKit::kAllProperty3D.size())}; + static_cast(std::numeric_limits::max())}; EXPECT_THROW(nvMolKit::calc3DProperties(mols_, properties, nvMolKit::Property3DOptions{}, nullptr), std::invalid_argument); EXPECT_THROW( From 57eb883a5252bb7619f32849fc5ed56447da52ca Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Fri, 23 Oct 2026 09:23:22 -0400 Subject: [PATCH 4/8] Compute correct EEM charges for charged molecules RDKit before 2026.9.2 solves a transposed EEM system (rdkit/rdkit#9688), so its charges for charged molecules sum to the negated formal charge. Solve the published equations instead, so charges sum to the formal charge. Tests compare against RDKit only for neutral molecules on RDKit < 2026.9.2, and add golden-value cases for charged molecules that hold on every version. --- nvmolkit/descriptors3d.py | 11 ++- nvmolkit/tests/test_descriptors3d.py | 117 ++++++++++++++++++++++----- src/descriptors3d_eem.cuh | 18 ++--- tests/CMakeLists.txt | 2 +- tests/test_descriptors3d.cpp | 69 +++++++++++++++- 5 files changed, 184 insertions(+), 33 deletions(-) diff --git a/nvmolkit/descriptors3d.py b/nvmolkit/descriptors3d.py index e4e4bde1..5f0a87f3 100644 --- a/nvmolkit/descriptors3d.py +++ b/nvmolkit/descriptors3d.py @@ -60,8 +60,15 @@ class Property3D(Enum): :class:`Device3DPropertyResult`): - ``EEMcharges``: RDKit's ``CalcEEMcharges`` electronegativity-equalization - partial charges, one value per atom. As in RDKit, a molecule's charges - sum to the negative of its formal charge. Conformers containing an + 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 `_. + 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, and molecules that cannot be kekulized, which RDKit rejects, produce NaN. diff --git a/nvmolkit/tests/test_descriptors3d.py b/nvmolkit/tests/test_descriptors3d.py index 37ebb252..bfc36a3f 100644 --- a/nvmolkit/tests/test_descriptors3d.py +++ b/nvmolkit/tests/test_descriptors3d.py @@ -5,6 +5,7 @@ import numpy as np import pytest +import rdkit import torch from rdkit import Chem from rdkit.Chem import rdDistGeom, rdMolDescriptors @@ -968,18 +969,37 @@ def test_property_and_coordinate_contract_errors_are_clear(): EEM_SMILES = ("CC(=O)[O-]", "C[NH3+]", "c1ccncc1O", "O=c1cc[nH]cc1", "CS(=O)(=O)NCl", "FC(F)(F)c1ccc(Br)cc1") +# RDKit before 2026.9.2 solved the transposed EEM system (rdkit/rdkit#9688), so its charges for charged molecules +# sum to the negated formal charge. nvMolKit gives the correct values; only neutral molecules are compared against +# older RDKit, and the golden tests below cover the charged ones. +RDKIT_EEM_CHARGES_CORRECT = tuple(int(part) for part in rdkit.__version__.split(".")[:3]) >= (2026, 9, 2) +requires_rdkit_eem_fix = pytest.mark.skipif( + not RDKIT_EEM_CHARGES_CORRECT, reason="RDKit before 2026.9.2 returns wrong EEM charges for charged molecules" +) + + +def _rdkit_eem_reference_is_correct(mol): + return RDKIT_EEM_CHARGES_CORRECT or Chem.GetFormalCharge(mol) == 0 + + def _rdkit_eem_rows(mols): - """RDKit EEM charges of every conformer, concatenated in molecule, then conformer order.""" - return np.concatenate( - [ - np.asarray(rdMolDescriptors.CalcEEMcharges(mol, confId=conf.GetId()), dtype=np.float64) - for mol in mols - for conf in mol.GetConformers() - ] - ) + """RDKit EEM charges of every conformer, concatenated in molecule, then conformer order. + Returns the charges and a mask of the rows where RDKit's value is correct for the installed version. + """ + rows, comparable = [], [] + for mol in mols: + for conf in mol.GetConformers(): + charges = np.asarray(rdMolDescriptors.CalcEEMcharges(mol, confId=conf.GetId()), dtype=np.float64) + rows.append(charges) + comparable.append(np.full(len(charges), _rdkit_eem_reference_is_correct(mol))) + return np.concatenate(rows), np.concatenate(comparable) -def _assert_eem_matches_rdkit(actual, expected, precision): + +def _assert_eem_matches_rdkit(actual, reference, precision): + expected, comparable = reference + actual = actual[comparable] + expected = expected[comparable] if precision == PrecisionMode.FULL: assert actual.dtype == np.float64 np.testing.assert_allclose(actual, expected, rtol=1e-10, atol=1e-10, equal_nan=True) @@ -988,6 +1008,12 @@ def _assert_eem_matches_rdkit(actual, expected, precision): np.testing.assert_allclose(actual, expected, rtol=0, atol=EEM_SINGLE_ATOL, equal_nan=True) +def _atom_sums(result, n_mols): + starts = result.atom_starts.torch().tolist() + charges = result[Property3D.EEM_CHARGES].numpy() + return [charges[starts[row] : starts[row + 1]].sum() for row in range(n_mols)] + + @pytest.mark.parametrize("precision", PRECISIONS) def test_eem_charges_are_per_atom_and_match_rdkit(precision): mols = [_embed(smiles, 2, 7 + idx) for idx, smiles in enumerate(EEM_SMILES)] @@ -1001,13 +1027,61 @@ def test_eem_charges_are_per_atom_and_match_rdkit(precision): _assert_eem_matches_rdkit(result[Property3D.EEM_CHARGES].numpy(), _rdkit_eem_rows(mols), precision) -def test_eem_charges_sum_to_negated_formal_charge_like_rdkit(): - mols = [_embed(smiles, 1, 3) for smiles in ("CC(=O)[O-]", "C[NH3+]", "CCO")] - result = Calc3DProperties(mols, Property3D.EEM_CHARGES, precision=PrecisionMode.FULL) - starts = result.atom_starts.torch().tolist() - charges = result[Property3D.EEM_CHARGES].numpy() - sums = [charges[starts[row] : starts[row + 1]].sum() for row in range(len(mols))] - np.testing.assert_allclose(sums, [1.0, -1.0, 0.0], atol=1e-9) +@pytest.mark.parametrize("precision", PRECISIONS) +def test_eem_charges_sum_to_formal_charge(precision): + smiles = ("CC(=O)[O-]", "C[NH3+]", "CCO", "[NH4+]", "[O-]C(=O)CC[NH3+]") + mols = [_embed(smi, 1, 3) for smi in smiles] + result = Calc3DProperties(mols, Property3D.EEM_CHARGES, precision=precision) + expected = [Chem.GetFormalCharge(mol) for mol in mols] + assert expected == [-1, 1, 0, 1, 0] + np.testing.assert_allclose( + _atom_sums(result, len(mols)), expected, atol=1e-5 if precision == PrecisionMode.SINGLE else 1e-9 + ) + + +def _eem_golden_molecule(smiles, coordinates): + mol = Chem.AddHs(Chem.MolFromSmiles(smiles)) + conf = Chem.Conformer(mol.GetNumAtoms()) + for idx, xyz in enumerate(coordinates): + conf.SetAtomPosition(idx, Point3D(*xyz)) + mol.AddConformer(conf) + return mol + + +_TETRAHEDRAL_NH = 1.01 / 3**0.5 +# Charges from solving the published EEM equations (Mortier 1986) with RDKit's parameters, as rdkit/rdkit#9688 does: +# the geometries and values of its reproduction. Independent of the installed RDKit, so these hold for any version. +EEM_GOLDEN = ( + ("[OH-]", [(0, 0, 0), (0.97, 0, 0)], [-1.0632, 0.0632]), + ( + "[NH4+]", + [ + (0, 0, 0), + (_TETRAHEDRAL_NH, _TETRAHEDRAL_NH, _TETRAHEDRAL_NH), + (_TETRAHEDRAL_NH, -_TETRAHEDRAL_NH, -_TETRAHEDRAL_NH), + (-_TETRAHEDRAL_NH, _TETRAHEDRAL_NH, -_TETRAHEDRAL_NH), + (-_TETRAHEDRAL_NH, -_TETRAHEDRAL_NH, _TETRAHEDRAL_NH), + ], + [-0.9845, 0.4961, 0.4961, 0.4961, 0.4961], + ), +) + + +@pytest.mark.parametrize("precision", PRECISIONS) +def test_eem_charged_molecules_match_golden_values(precision): + mols = [_eem_golden_molecule(smiles, coordinates) for smiles, coordinates, _ in EEM_GOLDEN] + result = Calc3DProperties(mols, Property3D.EEM_CHARGES, precision=precision) + expected = np.concatenate([charges for _, _, charges in EEM_GOLDEN]) + # The reference values are rounded to four decimals. + np.testing.assert_allclose(result[Property3D.EEM_CHARGES].numpy(), expected, rtol=0, atol=1e-4) + + +@requires_rdkit_eem_fix +@pytest.mark.parametrize("precision", PRECISIONS) +def test_eem_golden_molecules_match_fixed_rdkit(precision): + mols = [_eem_golden_molecule(smiles, coordinates) for smiles, coordinates, _ in EEM_GOLDEN] + result = Calc3DProperties(mols, Property3D.EEM_CHARGES, precision=precision) + _assert_eem_matches_rdkit(result[Property3D.EEM_CHARGES].numpy(), _rdkit_eem_rows(mols), precision) def test_eem_dense_pads_conformers_and_atoms(): @@ -1039,6 +1113,7 @@ def test_eem_device_coordinates_share_atom_starts(precision): values = coordinates.values.numpy() atom_starts = coordinates.atom_starts.torch().tolist() expected = np.full(len(values), np.nan) + comparable = np.zeros(len(values), dtype=bool) for row, mol_idx in enumerate(coordinates.mol_indices.torch().tolist()): mol = Chem.Mol(mols[mol_idx]) conf = Chem.Conformer(mol.GetNumAtoms()) @@ -1046,7 +1121,8 @@ def test_eem_device_coordinates_share_atom_starts(precision): conf.SetAtomPosition(atom_idx, Point3D(float(x), float(y), float(z))) conf_id = mol.AddConformer(conf, assignId=True) expected[atom_starts[row] : atom_starts[row + 1]] = rdMolDescriptors.CalcEEMcharges(mol, confId=conf_id) - _assert_eem_matches_rdkit(result[Property3D.EEM_CHARGES].numpy(), expected, precision) + comparable[atom_starts[row] : atom_starts[row + 1]] = _rdkit_eem_reference_is_correct(mol) + _assert_eem_matches_rdkit(result[Property3D.EEM_CHARGES].numpy(), (expected, comparable), precision) @pytest.mark.parametrize("precision", PRECISIONS) @@ -1072,7 +1148,7 @@ def test_eem_atoms_with_zero_parameters_give_nan_unlike_rdkit(precision): charges = result[Property3D.EEM_CHARGES].numpy() uncovered_rows = sum(mol.GetNumAtoms() for mol in uncovered) assert np.isnan(charges[:uncovered_rows]).all() - assert np.abs(_rdkit_eem_rows(uncovered)).max() > 5 # RDKit's zero-parameter charges + assert np.abs(_rdkit_eem_rows(uncovered)[0]).max() > 5 # RDKit's zero-parameter charges _assert_eem_matches_rdkit(charges[uncovered_rows:], _rdkit_eem_rows([ethanol]), precision) @@ -1090,8 +1166,11 @@ def test_eem_nonfinite_coordinates_invalidate_only_their_conformer(precision, in result = Calc3DProperties([molecule], Property3D.EEM_CHARGES, precision=precision) charges = result[Property3D.EEM_CHARGES].numpy() assert np.isnan(charges[: molecule.GetNumAtoms()]).all() + valid = charges[molecule.GetNumAtoms() :] + np.testing.assert_allclose(valid.sum(), Chem.GetFormalCharge(molecule), atol=1e-5) expected = np.asarray(rdMolDescriptors.CalcEEMcharges(molecule, confId=1)) - _assert_eem_matches_rdkit(charges[molecule.GetNumAtoms() :], expected, precision) + comparable = np.full(len(expected), _rdkit_eem_reference_is_correct(molecule)) + _assert_eem_matches_rdkit(valid, (expected, comparable), precision) @pytest.mark.parametrize("precision", PRECISIONS) diff --git a/src/descriptors3d_eem.cuh b/src/descriptors3d_eem.cuh index 189add48..d7747e27 100644 --- a/src/descriptors3d_eem.cuh +++ b/src/descriptors3d_eem.cuh @@ -154,13 +154,13 @@ struct EemConformer { }; /** - * @brief Entry (@p row, @p col) of RDKit's EEM system as it solves it, the right-hand side in column - * `numAtoms + 1`. Sets @p coincident when an atom pair's squared distance is not positive. + * @brief Entry (@p row, @p col) of the EEM system, the right-hand side in column `numAtoms + 1`. Sets + * @p coincident when an atom pair's squared distance is not positive. * - * RDKit fills the system row-major and hands it to Eigen as column-major, so it solves the transpose of the - * matrix its code describes: row i < n is `sum_j J_ij q_j + chi = -A_i` with `J_ii = B_i` and - * `J_ij = kappa / r_ij`, and the last row is `-sum_j q_j = Q`. The charges therefore sum to the negated - * formal charge, which this reproduces. + * Row i < n is `sum_j J_ij q_j - chi = -A_i` with `J_ii = B_i` and `J_ij = kappa / r_ij`, and the last row is + * `sum_j q_j = Q`, so the charges sum to the formal charge. RDKit before 2026.9.2 solved the transpose of this + * system and returned charges summing to the negated formal charge for charged molecules (rdkit/rdkit#9688); + * this follows the corrected equations. * * FP64 required: the coordinate differences, before the conversion to @p T (see centeredPosition()). */ @@ -172,10 +172,10 @@ eemSystemEntry(const EemConformer& conformer, const int row, const int col, bool return row < numAtoms ? static_cast(-conformer.electronegativity[row]) : static_cast(conformer.formalCharge); } if (row == numAtoms) { - return col < numAtoms ? T(-1) : T(0); + return col < numAtoms ? T(1) : T(0); } if (col == numAtoms) { - return T(1); + return T(-1); } if (row == col) { return static_cast(conformer.hardness[row]); @@ -230,7 +230,7 @@ __device__ __forceinline__ void refineEemSolution(const EemConformer& conformer, } /** - * @brief One block per conformer (grid-stride). Builds RDKit's EEM system (see eemSystemEntry()) and writes + * @brief One block per conformer (grid-stride). Builds the EEM system (see eemSystemEntry()) and writes * one partial charge per atom to @p output, laid out like the coordinate positions. */ template diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 91269b6c..e718288c 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -420,7 +420,7 @@ target_link_libraries( add_executable(test_descriptors3d test_descriptors3d.cpp) target_link_libraries( test_descriptors3d PRIVATE ${RDKit_LIBS} descriptors3d_mol - conformer_coord_upload device) + conformer_coord_upload device nvmolkit_versions) add_executable(test_descriptors3d_whim_arithmetic test_descriptors3d_whim_arithmetic.cu) diff --git a/tests/test_descriptors3d.cpp b/tests/test_descriptors3d.cpp index d02848c0..dff14ed8 100644 --- a/tests/test_descriptors3d.cpp +++ b/tests/test_descriptors3d.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include #include #include @@ -32,6 +33,7 @@ #include "src/conformer/conformer_coord_upload.h" #include "src/descriptors3d_mol.h" #include "src/utils/device.h" +#include "versions.h" using nvMolKit::Property3D; @@ -506,7 +508,19 @@ std::vector helixPoints(const int numAtoms, const double phase) { return points; } -//! Per-atom EEM charges of every conformer, compared against RDKit's EEM through the batch's atomStarts. +//! RDKit before 2026.9.2 solves a transposed EEM system (rdkit/rdkit#9688): its charges for charged molecules sum to +//! the negated formal charge. nvMolKit gives the correct values, so older RDKit is only a reference for neutral +//! molecules; the golden-value test covers the charged ones. +constexpr bool kRdkitEemChargesCorrect = + RDKIT_VERSION_MAJOR > 2026 || (RDKIT_VERSION_MAJOR == 2026 && + (RDKIT_VERSION_MINOR > 9 || (RDKIT_VERSION_MINOR == 9 && RDKIT_VERSION_PATCH >= 2))); + +bool rdkitEemIsCorrectFor(const RDKit::ROMol& mol) { + return kRdkitEemChargesCorrect || RDKit::MolOps::getFormalCharge(mol) == 0; +} + +//! Per-atom EEM charges of every conformer, compared against RDKit's EEM through the batch's atomStarts. Molecules +//! for which the installed RDKit's EEM is wrong are skipped. void expectEemMatchesRdkit(const std::vector& mols, const double tolerance) { const std::vector constMols(mols.begin(), mols.end()); auto results = nvMolKit::calc3DProperties(constMols, {Property3D::EEMcharges}, {}, nullptr); @@ -515,6 +529,9 @@ void expectEemMatchesRdkit(const std::vector& mols, const double int row = 0; for (RDKit::ROMol* mol : mols) { for (int confIdx = 0; confIdx < static_cast(mol->getNumConformers()); ++confIdx, ++row) { + if (!rdkitEemIsCorrectFor(*mol)) { + continue; + } std::vector expected; RDKit::Descriptors::EEM(*mol, expected, confIdx); ASSERT_EQ(atomStarts[row + 1] - atomStarts[row], static_cast(expected.size())); @@ -528,6 +545,7 @@ void expectEemMatchesRdkit(const std::vector& mols, const double } TEST(Descriptors3DEem, MatchesRdkitForChargedAndAromaticMolecules) { + // Charged molecules are compared only against an RDKit with the corrected EEM; see kRdkitEemChargesCorrect. std::vector> owned; for (const char* smiles : {"CC(=O)[O-]", "C[NH3+]", "c1ccncc1O", "O=c1cc[nH]cc1", "C#N", "CS(=O)(=O)NCl"}) { std::unique_ptr parsed(RDKit::SmilesToMol(smiles)); @@ -541,6 +559,48 @@ TEST(Descriptors3DEem, MatchesRdkitForChargedAndAromaticMolecules) { expectEemMatchesRdkit(mols, 1e-9); } +TEST(Descriptors3DEem, ChargedMoleculesMatchGoldenValues) { + // Charges from the published EEM equations with RDKit's parameters, the reproduction in rdkit/rdkit#9688. They hold + // for any installed RDKit and sum to the formal charge. + constexpr double a = 1.01 / 1.7320508075688772; // Tetrahedral N-H at 1.01 A. + RDKit::SmilesParserParams params; + params.removeHs = false; + std::unique_ptr hydroxide(RDKit::SmilesToMol("[O-][H]", params)); + std::unique_ptr ammonium(RDKit::SmilesToMol("[N+]([H])([H])([H])[H]", params)); + const std::vector> coordinates = { + {{0, 0, 0}, {0.97, 0, 0}}, + {{0, 0, 0}, {a, a, a}, {a, -a, -a}, {-a, a, -a}, {-a, -a, a}}, + }; + const std::vector> expected = {{-1.0632, 0.0632}, {-0.9845, 0.4961, 0.4961, 0.4961, 0.4961}}; + std::vector mols = {hydroxide.get(), ammonium.get()}; + for (size_t molIdx = 0; molIdx < mols.size(); ++molIdx) { + auto* conf = new RDKit::Conformer(mols[molIdx]->getNumAtoms()); + for (size_t atomIdx = 0; atomIdx < coordinates[molIdx].size(); ++atomIdx) { + const Point& p = coordinates[molIdx][atomIdx]; + conf->setAtomPos(atomIdx, RDGeom::Point3D(p[0], p[1], p[2])); + } + mols[molIdx]->addConformer(conf, true); + } + + const std::vector constMols(mols.begin(), mols.end()); + auto results = nvMolKit::calc3DProperties(constMols, {Property3D::EEMcharges}, {}, nullptr); + const auto values = toHost(results.properties.at(Property3D::EEMcharges)); + ASSERT_EQ(values.size(), 7u); + size_t row = 0; + for (size_t molIdx = 0; molIdx < mols.size(); ++molIdx) { + double sum = 0.0; + for (const double charge : expected[molIdx]) { + // The reference values are rounded to 4 decimals. + EXPECT_NEAR(values[row], charge, 1e-4) << RDKit::MolToSmiles(*mols[molIdx]) << " row " << row; + sum += values[row++]; + } + EXPECT_NEAR(sum, RDKit::MolOps::getFormalCharge(*mols[molIdx]), 1e-9); + } + if (kRdkitEemChargesCorrect) { + expectEemMatchesRdkit({hydroxide.get(), ammonium.get()}, 1e-9); + } +} + TEST(Descriptors3DEem, LargeMoleculesUseGlobalScratchAndMatchRdkit) { // 120 atoms exceed the shared-memory system in both precisions; the small molecule stays in shared memory. auto large = @@ -589,10 +649,15 @@ TEST(Descriptors3DEem, CoincidentAtomsGiveNan) { ASSERT_EQ(values.size(), 10u); std::vector expected; RDKit::Descriptors::EEM(*salt, expected, 1); + double sum = 0.0; for (int atomIdx = 0; atomIdx < 5; ++atomIdx) { EXPECT_TRUE(isNanBits(values[atomIdx])) << "atom " << atomIdx; - EXPECT_NEAR(values[5 + atomIdx], expected[atomIdx], 1e-9); + sum += values[5 + atomIdx]; + if (rdkitEemIsCorrectFor(*salt)) { + EXPECT_NEAR(values[5 + atomIdx], expected[atomIdx], 1e-9); + } } + EXPECT_NEAR(sum, RDKit::MolOps::getFormalCharge(*salt), 1e-9); } TEST(Descriptors3DEem, DeviceRowsOutsideConformersAreNan) { From ee6846f1afabaf0e5eb49fe6c8c1218f0d5cc36a Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Fri, 23 Oct 2026 09:24:07 -0400 Subject: [PATCH 5/8] Apply clang-format to EEM tests --- tests/test_descriptors3d.cpp | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/tests/test_descriptors3d.cpp b/tests/test_descriptors3d.cpp index dff14ed8..7b8eabf1 100644 --- a/tests/test_descriptors3d.cpp +++ b/tests/test_descriptors3d.cpp @@ -511,9 +511,10 @@ std::vector helixPoints(const int numAtoms, const double phase) { //! RDKit before 2026.9.2 solves a transposed EEM system (rdkit/rdkit#9688): its charges for charged molecules sum to //! the negated formal charge. nvMolKit gives the correct values, so older RDKit is only a reference for neutral //! molecules; the golden-value test covers the charged ones. -constexpr bool kRdkitEemChargesCorrect = - RDKIT_VERSION_MAJOR > 2026 || (RDKIT_VERSION_MAJOR == 2026 && - (RDKIT_VERSION_MINOR > 9 || (RDKIT_VERSION_MINOR == 9 && RDKIT_VERSION_PATCH >= 2))); +constexpr bool kRdkitEemChargesCorrect = RDKIT_VERSION_MAJOR > 2026 || + (RDKIT_VERSION_MAJOR == 2026 && + (RDKIT_VERSION_MINOR > 9 || + (RDKIT_VERSION_MINOR == 9 && RDKIT_VERSION_PATCH >= 2))); bool rdkitEemIsCorrectFor(const RDKit::ROMol& mol) { return kRdkitEemChargesCorrect || RDKit::MolOps::getFormalCharge(mol) == 0; @@ -565,14 +566,17 @@ TEST(Descriptors3DEem, ChargedMoleculesMatchGoldenValues) { constexpr double a = 1.01 / 1.7320508075688772; // Tetrahedral N-H at 1.01 A. RDKit::SmilesParserParams params; params.removeHs = false; - std::unique_ptr hydroxide(RDKit::SmilesToMol("[O-][H]", params)); - std::unique_ptr ammonium(RDKit::SmilesToMol("[N+]([H])([H])([H])[H]", params)); + std::unique_ptr hydroxide(RDKit::SmilesToMol("[O-][H]", params)); + std::unique_ptr ammonium(RDKit::SmilesToMol("[N+]([H])([H])([H])[H]", params)); const std::vector> coordinates = { {{0, 0, 0}, {0.97, 0, 0}}, {{0, 0, 0}, {a, a, a}, {a, -a, -a}, {-a, a, -a}, {-a, -a, a}}, }; - const std::vector> expected = {{-1.0632, 0.0632}, {-0.9845, 0.4961, 0.4961, 0.4961, 0.4961}}; - std::vector mols = {hydroxide.get(), ammonium.get()}; + const std::vector> expected = { + {-1.0632, 0.0632}, + {-0.9845, 0.4961, 0.4961, 0.4961, 0.4961} + }; + std::vector mols = {hydroxide.get(), ammonium.get()}; for (size_t molIdx = 0; molIdx < mols.size(); ++molIdx) { auto* conf = new RDKit::Conformer(mols[molIdx]->getNumAtoms()); for (size_t atomIdx = 0; atomIdx < coordinates[molIdx].size(); ++atomIdx) { From 3c389611c31ced6826a25869a2eb0410e09bbab3 Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Fri, 23 Oct 2026 10:32:33 -0400 Subject: [PATCH 6/8] Address review: bound EEM molecule size, size scratch by solved conformers, harden dense() - Conformers above 46,000 atoms give NaN instead of overflowing 32-bit indexes. - Size global scratch from the largest conformer actually solved, not the largest molecule in the batch. - Compare EEM charges in the benchmark only where the installed RDKit is correct. - dense() leaves invalid atom ranges NaN and builds per-atom gather indexes only when a per-atom property is requested. --- benchmarks/descriptors3d_bench.py | 23 +++++++++- nvmolkit/descriptors3d.py | 14 ++++-- nvmolkit/tests/test_descriptors3d.py | 52 +++++++++++++++++++++ nvmolkit/types.py | 69 +++++++++++++++++++++------- src/descriptors3d.h | 7 ++- src/descriptors3d_eem.cuh | 42 ++++++++++++++++- tests/test_descriptors3d.cpp | 47 +++++++++++++++++++ 7 files changed, 228 insertions(+), 26 deletions(-) diff --git a/benchmarks/descriptors3d_bench.py b/benchmarks/descriptors3d_bench.py index 85afedb8..4954bd0b 100644 --- a/benchmarks/descriptors3d_bench.py +++ b/benchmarks/descriptors3d_bench.py @@ -18,6 +18,7 @@ from pathlib import Path import numpy as np +import rdkit import torch from bench_utils import ( Deadline, @@ -197,6 +198,23 @@ def _calc_rdkit(mols: list[Chem.Mol], properties: tuple[Property3D, ...]) -> dic } +# 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 rdkit.__version__.split(".")[: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( mols: list[Chem.Mol], properties: tuple[Property3D, ...], @@ -229,7 +247,10 @@ def _validate( continue if prop == Property3D.EEM_CHARGES: eem_atol = EEM_TOLERANCES[precision] if tolerance is None else tolerance - np.testing.assert_allclose(actual, expected[prop], rtol=0, atol=eem_atol, equal_nan=True) + 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) diff --git a/nvmolkit/descriptors3d.py b/nvmolkit/descriptors3d.py index 5f0a87f3..1bc2d650 100644 --- a/nvmolkit/descriptors3d.py +++ b/nvmolkit/descriptors3d.py @@ -22,6 +22,7 @@ Device3DResult, HardwareOptions, PrecisionMode, + _atom_row_gather, _atom_row_index, _resolve_cuda_stream, _scatter_atom_rows, @@ -70,8 +71,9 @@ class Property3D(Enum): 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, and molecules - that cannot be kekulized, which RDKit rejects, produce NaN. + 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" @@ -287,14 +289,18 @@ def dense(self, pad_value: float = float("nan")) -> Dense3DPropertyResult: 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. """ - index = _atom_row_index(self.atom_starts, self.mol_indices, self.conf_indices, self.n_mols) + 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) 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, pad_value) + values[name] = _scatter_atom_rows(source, index, gather, pad_value) continue dense_values = torch.full( (*index.conf_mask.shape, *source.shape[1:]), pad_value, dtype=source.dtype, device=source.device diff --git a/nvmolkit/tests/test_descriptors3d.py b/nvmolkit/tests/test_descriptors3d.py index bfc36a3f..4452fb04 100644 --- a/nvmolkit/tests/test_descriptors3d.py +++ b/nvmolkit/tests/test_descriptors3d.py @@ -12,6 +12,7 @@ from rdkit.Chem.rdDistGeom import EmbedParameters from rdkit.Geometry import Point3D +import nvmolkit.descriptors3d as descriptors3d_module from nvmolkit.descriptors3d import ( PER_ATOM_PROPERTIES, Calc3DProperties, @@ -1194,3 +1195,54 @@ def test_eem_empty_batches_and_molecules_without_conformers(precision): assert result[Property3D.EEM_CHARGES].torch().numel() == 0 dense = result.dense() assert dense.values[Property3D.EEM_CHARGES.value].shape == (len(molecules), 0, 0) + + +def _with_atom_starts(coordinates, atom_starts): + starts = torch.tensor(atom_starts, dtype=torch.int32, device="cuda") + return Device3DResult( + values=coordinates.values, + atom_starts=AsyncGpuResult(starts, coordinates.gpu_id), + mol_indices=coordinates.mol_indices, + conf_indices=coordinates.conf_indices, + gpu_id=coordinates.gpu_id, + n_mols=coordinates.n_mols, + ) + + +@pytest.mark.parametrize( + ("atom_starts", "valid_conformer", "first_row"), + [([4, 8, 12], 0, 4), ([-4, 0, 4], 1, 0)], + ids=["past_the_rows", "negative_start"], +) +def test_dense_charges_leave_invalid_atom_ranges_nan(atom_starts, valid_conformer, first_row): + square = [(0, 0, 0), (1.5, 0, 0), (1.5, 1.5, 0), (0, 1.5, 1.0)] + shifted = [(x + 0.3 * idx, y, z) for idx, (x, y, z) in enumerate(square)] + mol = _mol_with_conformers("CCCC", [square, shifted]) + coordinates = _with_atom_starts(_device_result_from_molecules([mol]), atom_starts) + + result = Calc3DProperties([mol], Property3D.EEM_CHARGES, coordinates=coordinates, precision=PrecisionMode.FULL) + dense = result.dense() + + charges = result[Property3D.EEM_CHARGES].numpy() + invalid_conformer = 1 - valid_conformer + assert dense.atom_mask.tolist() == [[[valid_conformer == 0] * 4, [valid_conformer == 1] * 4]] + np.testing.assert_array_equal( + dense.values[Property3D.EEM_CHARGES.value][0, valid_conformer].cpu().numpy(), + charges[first_row : first_row + 4], + ) + assert torch.isnan(dense.values[Property3D.EEM_CHARGES.value][0, invalid_conformer]).all() + assert not np.isnan(charges[first_row : first_row + 4]).any() + + +def test_dense_scalar_properties_do_not_build_atom_gather_indexes(monkeypatch): + def fail(*args, **kwargs): + raise AssertionError("per-atom gather indexes built for scalar properties") + + monkeypatch.setattr(descriptors3d_module, "_atom_row_gather", fail) + mols = [_embed("CCO", 2, 1), _embed("CCCO", 1, 2)] + result = Calc3DProperties(mols, [Property3D.PMI1, Property3D.PBF], precision=PrecisionMode.FULL) + dense = result.dense() + assert dense.values[Property3D.PMI1.value].shape == (2, 2) + assert dense.conf_mask.tolist() == [[True, True], [True, False]] + assert dense.atom_mask.shape == (2, 2, max(mol.GetNumAtoms() for mol in mols)) + assert int(dense.atom_mask.sum()) == sum(mol.GetNumAtoms() * mol.GetNumConformers() for mol in mols) diff --git a/nvmolkit/types.py b/nvmolkit/types.py index 6769c6dd..65ea6387 100644 --- a/nvmolkit/types.py +++ b/nvmolkit/types.py @@ -522,58 +522,93 @@ def dense(self, pad_value: float = float("nan")) -> "Dense3DResult": Both masks are ``True`` where the data is real, ``False`` where padded. Reading the index tensors synchronizes implicitly. """ - index = _atom_row_index(self.atom_starts, self.mol_indices, self.conf_indices, self.n_mols) + values = self.values.torch() + index = _atom_row_index(self.atom_starts, self.mol_indices, self.conf_indices, self.n_mols, values.shape[0]) return Dense3DResult( - values=_scatter_atom_rows(self.values.torch(), index, pad_value), + values=_scatter_atom_rows(values, index, _atom_row_gather(index), pad_value), conf_mask=index.conf_mask, atom_mask=index.atom_mask, ) class _AtomRowIndex(NamedTuple): - """Dense ``(molecule, conformer, atom)`` position of every per-atom row of a CSR conformer batch.""" + """Padded layout of a CSR conformer batch: masks plus each conformer's valid atom-row range.""" + + conf_mask: "torch.Tensor" + atom_mask: "torch.Tensor" + starts: "torch.Tensor" + sizes: "torch.Tensor" + molecules: "torch.Tensor" + conformers: "torch.Tensor" + + +class _AtomRowGather(NamedTuple): + """Dense ``(molecule, conformer, atom)`` position of every per-atom row of the valid conformers.""" rows: "torch.Tensor" molecule: "torch.Tensor" conformer: "torch.Tensor" atom: "torch.Tensor" - conf_mask: "torch.Tensor" - atom_mask: "torch.Tensor" def _atom_row_index( - atom_starts: AsyncGpuResult, mol_indices: AsyncGpuResult, conf_indices: AsyncGpuResult, n_mols: int + atom_starts: AsyncGpuResult, + mol_indices: AsyncGpuResult, + conf_indices: AsyncGpuResult, + n_mols: int, + n_rows: Optional[int] = None, ) -> _AtomRowIndex: """Map each conformer's ``atom_starts`` range to padded ``(n_mols, max_confs, max_atoms)`` slots. - Rows outside every conformer's range are not indexed. Reading the index tensors synchronizes implicitly. + Conformers whose range is empty, starts below zero or, when ``n_rows`` (the number of per-atom rows) is given, + ends past it, have no atom slots. Reading the index tensors synchronizes implicitly. """ starts = atom_starts.torch().to(torch.int64) mols = mol_indices.torch().to(torch.int64) confs = conf_indices.torch().to(torch.int64) device = starts.device sizes = starts[1:] - starts[:-1] + valid = (starts[:-1] >= 0) & (sizes > 0) + if n_rows is not None: + valid &= starts[1:] <= n_rows + sizes = torch.where(valid, sizes, torch.zeros_like(sizes)) max_confs = int(torch.bincount(mols, minlength=n_mols).max().item()) if mols.numel() else 0 max_atoms = int(sizes.max().item()) if sizes.numel() else 0 conf_mask = torch.zeros((n_mols, max_confs), dtype=torch.bool, device=device) conf_mask[mols, confs] = True - first_slot = (torch.cumsum(sizes, 0) - sizes).repeat_interleave(sizes) - atom = torch.arange(int(sizes.sum().item()), device=device, dtype=torch.int64) - first_slot - rows = starts[:-1].repeat_interleave(sizes) + atom - molecule = mols.repeat_interleave(sizes) - conformer = confs.repeat_interleave(sizes) - atom_mask = torch.zeros((n_mols, max_confs, max_atoms), dtype=torch.bool, device=device) - atom_mask[molecule, conformer, atom] = True + conformer_sizes = torch.zeros((n_mols, max_confs), dtype=torch.int64, device=device) + conformer_sizes[mols, confs] = sizes + atom_mask = torch.arange(max_atoms, device=device) < conformer_sizes.unsqueeze(-1) return _AtomRowIndex( - rows=rows, molecule=molecule, conformer=conformer, atom=atom, conf_mask=conf_mask, atom_mask=atom_mask + conf_mask=conf_mask, + atom_mask=atom_mask, + starts=starts[:-1], + sizes=sizes, + molecules=mols, + conformers=confs, + ) + + +def _atom_row_gather(index: _AtomRowIndex) -> _AtomRowGather: + """Per-atom-row positions of ``index``'s valid conformers; needed only to move per-atom values.""" + device = index.sizes.device + first_slot = (torch.cumsum(index.sizes, 0) - index.sizes).repeat_interleave(index.sizes) + atom = torch.arange(int(index.sizes.sum().item()), device=device, dtype=torch.int64) - first_slot + return _AtomRowGather( + rows=index.starts.repeat_interleave(index.sizes) + atom, + molecule=index.molecules.repeat_interleave(index.sizes), + conformer=index.conformers.repeat_interleave(index.sizes), + atom=atom, ) -def _scatter_atom_rows(values: "torch.Tensor", index: _AtomRowIndex, pad_value: float) -> "torch.Tensor": +def _scatter_atom_rows( + values: "torch.Tensor", index: _AtomRowIndex, gather: _AtomRowGather, pad_value: float +) -> "torch.Tensor": """Pad per-atom rows ``values`` (any trailing shape) to ``(n_mols, max_confs, max_atoms, *values.shape[1:])``.""" dense = torch.full( (*index.atom_mask.shape, *values.shape[1:]), pad_value, dtype=values.dtype, device=values.device ) - dense[index.molecule, index.conformer, index.atom] = values[index.rows] + dense[gather.molecule, gather.conformer, gather.atom] = values[gather.rows] return dense diff --git a/src/descriptors3d.h b/src/descriptors3d.h index e2ae6522..8c684a2a 100644 --- a/src/descriptors3d.h +++ b/src/descriptors3d.h @@ -77,6 +77,9 @@ inline constexpr std::array kAllProperty3D = { Property3D::EEMcharges, }; +//! Largest molecule EEMcharges solves. The (n + 1) x (n + 2) system is indexed with 32-bit integers, which bounds +//! n below 46,339; conformers of larger molecules produce NaN. +inline constexpr int kMaxEemAtoms = 46000; inline constexpr int kNumWhimProperties = 114; inline constexpr int kNumRdfProperties = 210; //!< 7 atom-property channels x 30 radii. inline constexpr int kNumMorseProperties = 224; //!< 7 atom-property channels x 32 scattering values. @@ -271,8 +274,8 @@ struct Property3DDeviceInputs { const double* eemHardness = nullptr; //! EEMcharges: per molecule, the sum of formal charges. const double* moleculeFormalCharges = nullptr; - //! WHIM, AUTOCORR3D, GETAWAY, EEMcharges: largest molecule atom count in the batch (host value). Sizes EEM's - //! linear-system scratch and WHIM's + //! WHIM, AUTOCORR3D, GETAWAY, EEMcharges: largest molecule atom count in the batch (host value). Bounds EEM's + //! conformer size and WHIM's //! per-conformer symmetry-search scratch (rows with more atoms produce NaN) and, with 0 or above 3072 atoms, //! leaves AUTOCORR3D and GETAWAY without the shared bond-distance table. int32_t maxMoleculeAtoms = 0; diff --git a/src/descriptors3d_eem.cuh b/src/descriptors3d_eem.cuh index d7747e27..c0d172ca 100644 --- a/src/descriptors3d_eem.cuh +++ b/src/descriptors3d_eem.cuh @@ -229,6 +229,29 @@ __device__ __forceinline__ void refineEemSolution(const EemConformer& conformer, backSubstitute(system, size); } +//! Largest conformer the EEM kernel solves; larger ones produce NaN. The system's 32-bit indexing bounds it. +__host__ __device__ inline int eemAtomLimit(const Property3DDeviceInputs& inputs) { + return inputs.maxMoleculeAtoms < kMaxEemAtoms ? inputs.maxMoleculeAtoms : kMaxEemAtoms; +} + +/** + * @brief Writes the largest atom count among conformers with a valid atom range that eemChargesKernel() will solve + * to @p maxAtoms (zero-initialized), so that global scratch is sized by those and not by unused molecules. + */ +__global__ void eemMaxSolvedAtomsKernel(const DeviceCoordView coordinates, const int atomLimit, int* maxAtoms) { + const int stride = static_cast(gridDim.x * blockDim.x); + for (int conformerIdx = static_cast(blockIdx.x * blockDim.x + threadIdx.x); + conformerIdx < coordinates.numConformers; + conformerIdx += stride) { + const int64_t atomStart = coordinates.atomStarts[conformerIdx]; + const int64_t atomStop = coordinates.atomStarts[conformerIdx + 1]; + const bool validRange = atomStart >= 0 && atomStop > atomStart && atomStop <= coordinates.numAtoms; + if (validRange && atomStop - atomStart <= atomLimit) { + atomicMax(maxAtoms, static_cast(atomStop - atomStart)); + } + } +} + /** * @brief One block per conformer (grid-stride). Builds the EEM system (see eemSystemEntry()) and writes * one partial charge per atom to @p output, laid out like the coordinate positions. @@ -249,7 +272,7 @@ __global__ void __launch_bounds__(kEemBlockSize) eemChargesKernel(const DeviceCo } const int numAtoms = static_cast(atomStop - atomStart); const int moleculeIdx = coordinates.molIndices[conformerIdx]; - bool valid = moleculeIdx >= 0 && moleculeIdx < coordinates.nMols && numAtoms <= inputs.maxMoleculeAtoms; + bool valid = moleculeIdx >= 0 && moleculeIdx < coordinates.nMols && numAtoms <= eemAtomLimit(inputs); int paramStart = 0; if (valid) { paramStart = inputs.moleculeAtomStarts[moleculeIdx]; @@ -331,8 +354,23 @@ void launchEemCharges(const DeviceCoordView& coordinates, int numSms = 0; cudaCheckError(cudaGetDevice(&device)); cudaCheckError(cudaDeviceGetAttribute(&numSms, cudaDevAttrMultiProcessorCount, device)); + AsyncDeviceVector maxSolvedAtoms(1, stream); + cudaCheckError(cudaMemsetAsync(maxSolvedAtoms.data(), 0, sizeof(int), stream)); + constexpr int kMaxAtomsBlock = 256; + eemMaxSolvedAtomsKernel<<>>(coordinates, eemAtomLimit(inputs), maxSolvedAtoms.data()); + cudaCheckError(cudaGetLastError()); + int largestSolved = 0; + cudaCheckError(cudaMemcpyAsync(&largestSolved, maxSolvedAtoms.data(), sizeof(int), cudaMemcpyDeviceToHost, stream)); + cudaCheckError(cudaStreamSynchronize(stream)); + if (largestSolved == 0) { + return; // Every row stays NaN. + } + const int numBlocks = std::min(coordinates.numConformers, numSms * 16); - const int64_t largestWorkspace = EemWorkspace::bytes(inputs.maxMoleculeAtoms); + const int64_t largestWorkspace = EemWorkspace::bytes(largestSolved); const int64_t sharedCapacity = std::min(largestWorkspace, kEemSharedBytes); const int64_t scratchStride = largestWorkspace > sharedCapacity ? largestWorkspace : 0; AsyncDeviceVector scratch(static_cast(scratchStride * numBlocks), stream); diff --git a/tests/test_descriptors3d.cpp b/tests/test_descriptors3d.cpp index 7b8eabf1..72afe764 100644 --- a/tests/test_descriptors3d.cpp +++ b/tests/test_descriptors3d.cpp @@ -613,6 +613,53 @@ TEST(Descriptors3DEem, LargeMoleculesUseGlobalScratchAndMatchRdkit) { expectEemMatchesRdkit({large.get(), small.get()}, 1e-8); } +//! An acyclic carbon chain of @p numAtoms atoms with the given conformers, built without SMILES parsing. +std::unique_ptr carbonChain(const int numAtoms, const std::vector>& confs) { + auto mol = std::make_unique(); + for (int atomIdx = 0; atomIdx < numAtoms; ++atomIdx) { + mol->addAtom(new RDKit::Atom(6), false, true); + if (atomIdx > 0) { + mol->addBond(atomIdx - 1, atomIdx, RDKit::Bond::SINGLE); + } + } + RDKit::MolOps::sanitizeMol(*mol); + for (const auto& points : confs) { + auto conf = std::make_unique(numAtoms); + for (int atomIdx = 0; atomIdx < numAtoms; ++atomIdx) { + conf->setAtomPos(atomIdx, RDGeom::Point3D(points[atomIdx][0], points[atomIdx][1], points[atomIdx][2])); + } + mol->addConformer(conf.release(), true); + } + return mol; +} + +TEST(Descriptors3DEem, MoleculesWithoutConformersDoNotInflateScratch) { + // A 20,000-atom molecule would need gigabytes of global scratch per block if it sized the workspace; it has no + // conformers, so only the small molecules' conformers (which fit in shared memory) matter. + auto unused = carbonChain(20000, {}); + auto small = molWithConformers("CCO", {helixPoints(3, 0.0), helixPoints(3, 0.4)}); + expectEemMatchesRdkit({unused.get(), small.get()}, 1e-9); +} + +TEST(Descriptors3DEem, MoleculesAboveTheAtomLimitGiveNan) { + // The solver's 32-bit indexing bounds the molecule size; larger conformers are NaN, not an out-of-range read. + const int numAtoms = nvMolKit::kMaxEemAtoms + 1; + auto huge = carbonChain(numAtoms, {helixPoints(numAtoms, 0.0)}); + auto ethanol = molWithConformers("CCO", {helixPoints(3, 0.0)}); + const std::vector mols = {huge.get(), ethanol.get()}; + auto results = nvMolKit::calc3DProperties(mols, {Property3D::EEMcharges}, {}, nullptr); + const auto values = toHost(results.properties.at(Property3D::EEMcharges)); + ASSERT_EQ(values.size(), static_cast(numAtoms) + 3); + for (int atomIdx = 0; atomIdx < numAtoms; ++atomIdx) { + ASSERT_TRUE(isNanBits(values[atomIdx])) << "atom " << atomIdx; + } + std::vector expected; + RDKit::Descriptors::EEM(*ethanol, expected, 0); + for (int atomIdx = 0; atomIdx < 3; ++atomIdx) { + EXPECT_NEAR(values[numAtoms + atomIdx], expected[atomIdx], 1e-9); + } +} + TEST(Descriptors3DEem, UnparameterizedAtomsGiveNanOnlyForTheirMolecule) { // RDKit has no iodine parameters (it reads past its tables), and a lone aromatic ring atom cannot be kekulized. auto iodide = molWithConformers("CCI", {helixPoints(3, 0.0)}); From 0d341ff105853c16afd406645906f227871840e0 Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Fri, 23 Oct 2026 10:33:21 -0400 Subject: [PATCH 7/8] Apply clang-format --- src/descriptors3d.cu | 2 +- src/descriptors3d.h | 4 ++-- src/descriptors3d_eem.cuh | 4 ++-- src/descriptors3d_mol.cpp | 2 +- tests/test_descriptors3d.cpp | 10 +++++----- 5 files changed, 11 insertions(+), 11 deletions(-) diff --git a/src/descriptors3d.cu b/src/descriptors3d.cu index e9a7eae2..545125f9 100644 --- a/src/descriptors3d.cu +++ b/src/descriptors3d.cu @@ -333,7 +333,7 @@ Property3DResults calc3DPropertiesGpu(const DeviceCoordView& coord Real* usrcatOutput = nullptr; Real* getawayOutput = nullptr; DclvOutputs dclvOutputs; - Real* eemOutput = nullptr; + Real* eemOutput = nullptr; for (const Property3D property : properties) { // Bounds the work arrays, which hold one slot per known property. if (std::find(kAllProperty3D.begin(), kAllProperty3D.end(), property) == kAllProperty3D.end()) { diff --git a/src/descriptors3d.h b/src/descriptors3d.h index 8c684a2a..432da6b6 100644 --- a/src/descriptors3d.h +++ b/src/descriptors3d.h @@ -45,7 +45,7 @@ enum class Property3D : int { DCLVPolarVolume = 22, DCLVCompactness = 23, DCLVPackingDensity = 24, - EEMcharges = 25, + EEMcharges = 25, }; inline constexpr std::array kAllProperty3D = { @@ -79,7 +79,7 @@ inline constexpr std::array kAllProperty3D = { //! Largest molecule EEMcharges solves. The (n + 1) x (n + 2) system is indexed with 32-bit integers, which bounds //! n below 46,339; conformers of larger molecules produce NaN. -inline constexpr int kMaxEemAtoms = 46000; +inline constexpr int kMaxEemAtoms = 46000; inline constexpr int kNumWhimProperties = 114; inline constexpr int kNumRdfProperties = 210; //!< 7 atom-property channels x 30 radii. inline constexpr int kNumMorseProperties = 224; //!< 7 atom-property channels x 32 scattering values. diff --git a/src/descriptors3d_eem.cuh b/src/descriptors3d_eem.cuh index c0d172ca..aadd8c0e 100644 --- a/src/descriptors3d_eem.cuh +++ b/src/descriptors3d_eem.cuh @@ -243,8 +243,8 @@ __global__ void eemMaxSolvedAtomsKernel(const DeviceCoordView coordinates, const for (int conformerIdx = static_cast(blockIdx.x * blockDim.x + threadIdx.x); conformerIdx < coordinates.numConformers; conformerIdx += stride) { - const int64_t atomStart = coordinates.atomStarts[conformerIdx]; - const int64_t atomStop = coordinates.atomStarts[conformerIdx + 1]; + const int64_t atomStart = coordinates.atomStarts[conformerIdx]; + const int64_t atomStop = coordinates.atomStarts[conformerIdx + 1]; const bool validRange = atomStart >= 0 && atomStop > atomStart && atomStop <= coordinates.numAtoms; if (validRange && atomStop - atomStart <= atomLimit) { atomicMax(maxAtoms, static_cast(atomStop - atomStart)); diff --git a/src/descriptors3d_mol.cpp b/src/descriptors3d_mol.cpp index 4a1c87b4..721d15a8 100644 --- a/src/descriptors3d_mol.cpp +++ b/src/descriptors3d_mol.cpp @@ -5,8 +5,8 @@ #include #include -#include #include +#include #include #include #include diff --git a/tests/test_descriptors3d.cpp b/tests/test_descriptors3d.cpp index 72afe764..14aa1ed0 100644 --- a/tests/test_descriptors3d.cpp +++ b/tests/test_descriptors3d.cpp @@ -11,8 +11,8 @@ #include #include #include -#include #include +#include #include #include #include @@ -643,10 +643,10 @@ TEST(Descriptors3DEem, MoleculesWithoutConformersDoNotInflateScratch) { TEST(Descriptors3DEem, MoleculesAboveTheAtomLimitGiveNan) { // The solver's 32-bit indexing bounds the molecule size; larger conformers are NaN, not an out-of-range read. - const int numAtoms = nvMolKit::kMaxEemAtoms + 1; - auto huge = carbonChain(numAtoms, {helixPoints(numAtoms, 0.0)}); - auto ethanol = molWithConformers("CCO", {helixPoints(3, 0.0)}); - const std::vector mols = {huge.get(), ethanol.get()}; + const int numAtoms = nvMolKit::kMaxEemAtoms + 1; + auto huge = carbonChain(numAtoms, {helixPoints(numAtoms, 0.0)}); + auto ethanol = molWithConformers("CCO", {helixPoints(3, 0.0)}); + const std::vector mols = {huge.get(), ethanol.get()}; auto results = nvMolKit::calc3DProperties(mols, {Property3D::EEMcharges}, {}, nullptr); const auto values = toHost(results.properties.at(Property3D::EEMcharges)); ASSERT_EQ(values.size(), static_cast(numAtoms) + 3); From c9d5ec73eb1b03f0c6096c7613bb639ebe47410b Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Fri, 23 Oct 2026 15:43:16 -0400 Subject: [PATCH 8/8] Bound EEM scratch allocation and handle RDKit development versions --- benchmarks/descriptors3d_bench.py | 3 +- benchmarks/tests/test_descriptors3d_bench.py | 28 ++++++++++++ nvmolkit/tests/test_descriptors3d.py | 30 ++++++++++++- src/descriptors3d_eem.cuh | 13 +++++- tests/CMakeLists.txt | 6 +++ tests/test_descriptors3d_eem.cu | 45 ++++++++++++++++++++ 6 files changed, 122 insertions(+), 3 deletions(-) create mode 100644 tests/test_descriptors3d_eem.cu diff --git a/benchmarks/descriptors3d_bench.py b/benchmarks/descriptors3d_bench.py index 4954bd0b..e192d87e 100644 --- a/benchmarks/descriptors3d_bench.py +++ b/benchmarks/descriptors3d_bench.py @@ -15,6 +15,7 @@ import argparse import math +import re from pathlib import Path import numpy as np @@ -200,7 +201,7 @@ def _calc_rdkit(mols: list[Chem.Mol], properties: tuple[Property3D, ...]) -> dic # 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 rdkit.__version__.split(".")[:3]) >= (2026, 9, 2) +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: diff --git a/benchmarks/tests/test_descriptors3d_bench.py b/benchmarks/tests/test_descriptors3d_bench.py index de30e311..9dfc6739 100644 --- a/benchmarks/tests/test_descriptors3d_bench.py +++ b/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( diff --git a/nvmolkit/tests/test_descriptors3d.py b/nvmolkit/tests/test_descriptors3d.py index 4452fb04..a4c387f2 100644 --- a/nvmolkit/tests/test_descriptors3d.py +++ b/nvmolkit/tests/test_descriptors3d.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 import gc +import re import numpy as np import pytest @@ -973,7 +974,7 @@ def test_property_and_coordinate_contract_errors_are_clear(): # RDKit before 2026.9.2 solved the transposed EEM system (rdkit/rdkit#9688), so its charges for charged molecules # sum to the negated formal charge. nvMolKit gives the correct values; only neutral molecules are compared against # older RDKit, and the golden tests below cover the charged ones. -RDKIT_EEM_CHARGES_CORRECT = tuple(int(part) for part in rdkit.__version__.split(".")[:3]) >= (2026, 9, 2) +RDKIT_EEM_CHARGES_CORRECT = tuple(int(part) for part in re.findall(r"\d+", rdkit.__version__)[:3]) >= (2026, 9, 2) requires_rdkit_eem_fix = pytest.mark.skipif( not RDKIT_EEM_CHARGES_CORRECT, reason="RDKit before 2026.9.2 returns wrong EEM charges for charged molecules" ) @@ -1187,6 +1188,33 @@ def test_eem_global_workspace_and_custom_stream(precision): _assert_eem_matches_rdkit(result[Property3D.EEM_CHARGES].numpy(), _rdkit_eem_rows([molecule]), precision) +@pytest.mark.parametrize("precision", PRECISIONS) +def test_eem_budget_limited_grid_preserves_mixed_conformer_order(precision): + # 256 atoms need ~262/521 KiB per block in single/full precision. The 64 MiB budget caps + # the grid at 250/125 blocks on GPUs with enough SMs; 515 conformers require scratch reuse. + large = Chem.MolFromSmiles("O" + "C" * 255) + conformer = Chem.Conformer(large.GetNumAtoms()) + for atom_idx in range(large.GetNumAtoms()): + angle = 1.9 * atom_idx + conformer.SetAtomPosition(atom_idx, Point3D(1.4 * np.cos(angle), 1.4 * np.sin(angle), 0.55 * atom_idx)) + large.AddConformer(conformer, assignId=True) + small = _embed("CCO", 2, 42) + molecules = [large, *([small] * 128), large, *([small] * 128), large] + expected = np.concatenate( + [ + rdMolDescriptors.CalcEEMcharges(mol, confId=conf.GetId()) + for mol in molecules + for conf in mol.GetConformers() + ] + ) + stream = torch.cuda.Stream() + for _ in range(2): + result = Calc3DProperties(molecules, Property3D.EEM_CHARGES, precision=precision, stream=stream) + _assert_eem_matches_rdkit( + result[Property3D.EEM_CHARGES].numpy(), (expected, np.ones(len(expected), dtype=bool)), precision + ) + + @pytest.mark.parametrize("precision", PRECISIONS) def test_eem_empty_batches_and_molecules_without_conformers(precision): for molecules in ([], [Chem.MolFromSmiles("CCO")]): diff --git a/src/descriptors3d_eem.cuh b/src/descriptors3d_eem.cuh index aadd8c0e..d9517c5b 100644 --- a/src/descriptors3d_eem.cuh +++ b/src/descriptors3d_eem.cuh @@ -24,6 +24,17 @@ constexpr int kEemBlockSize = 128; //! the solver's static shared memory below the 48 KiB launch limit. constexpr int kEemSharedBytes = 47 * 1024; +//! Global scratch budget across blocks. A single larger workspace is allowed so its conformer can still run. +constexpr int64_t kEemScratchBudgetBytes = 64 * 1024 * 1024; + +inline int eemBlockCount(const int numConformers, const int numSms, const int64_t scratchStride) { + const int blockLimit = std::min(numConformers, numSms * 16); + if (scratchStride == 0) { + return blockLimit; + } + return static_cast(std::min(blockLimit, std::max(1, kEemScratchBudgetBytes / scratchStride))); +} + __host__ __device__ constexpr int64_t alignTo8(const int64_t bytes) { return (bytes + 7) / 8 * 8; } @@ -369,10 +380,10 @@ void launchEemCharges(const DeviceCoordView& coordinates, return; // Every row stays NaN. } - const int numBlocks = std::min(coordinates.numConformers, numSms * 16); const int64_t largestWorkspace = EemWorkspace::bytes(largestSolved); const int64_t sharedCapacity = std::min(largestWorkspace, kEemSharedBytes); const int64_t scratchStride = largestWorkspace > sharedCapacity ? largestWorkspace : 0; + const int numBlocks = eemBlockCount(coordinates.numConformers, numSms, scratchStride); AsyncDeviceVector scratch(static_cast(scratchStride * numBlocks), stream); eemChargesKernel<<>>(coordinates, inputs, diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index e718288c..e0d5eda1 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -431,6 +431,11 @@ add_executable(test_descriptors3d_topology test_descriptors3d_topology.cu) target_link_libraries(test_descriptors3d_topology PRIVATE descriptors3d cuda_error_check) +add_executable(test_descriptors3d_eem test_descriptors3d_eem.cu) +target_link_libraries( + test_descriptors3d_eem PRIVATE conformer_types device_vector cuda_error_check + CUDA::cudart) + add_executable(test_tfd_cpu test_tfd_cpu.cpp) target_link_libraries(test_tfd_cpu PRIVATE ${RDKit_LIBS} tfd_cpu test_utils nvmolkit_versions) @@ -452,6 +457,7 @@ set(TEST_LIST test_descriptors3d test_descriptors3d_whim_arithmetic test_descriptors3d_topology + test_descriptors3d_eem test_boolean_tree test_bitbirch test_butina diff --git a/tests/test_descriptors3d_eem.cu b/tests/test_descriptors3d_eem.cu new file mode 100644 index 00000000..7ef0a9c7 --- /dev/null +++ b/tests/test_descriptors3d_eem.cu @@ -0,0 +1,45 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include + +#include + +#include "src/descriptors3d_eem.cuh" + +using namespace nvMolKit::descriptors3d_detail; + +TEST(Descriptors3DEemScratch, SharedMemoryNeedsNoBudgetLimit) { + EXPECT_EQ(eemBlockCount(0, 20, 0), 0); + EXPECT_EQ(eemBlockCount(1, 20, 0), 1); + EXPECT_EQ(eemBlockCount(1000, 20, 0), 320); +} + +TEST(Descriptors3DEemScratch, RespectsConformerAndOccupancyLimits) { + EXPECT_EQ(eemBlockCount(0, 20, 1024), 0); + EXPECT_EQ(eemBlockCount(1, 20, 1024), 1); + EXPECT_EQ(eemBlockCount(100, 20, 1024), 100); + EXPECT_EQ(eemBlockCount(1000, 20, 1024), 320); +} + +TEST(Descriptors3DEemScratch, BudgetBoundaryAndSingleOversizedWorkspace) { + EXPECT_EQ(eemBlockCount(1000, 20, kEemScratchBudgetBytes / 2), 2); + EXPECT_EQ(eemBlockCount(1000, 20, kEemScratchBudgetBytes / 2 + 1), 1); + EXPECT_EQ(eemBlockCount(1000, 20, kEemScratchBudgetBytes), 1); + EXPECT_EQ(eemBlockCount(1000, 20, kEemScratchBudgetBytes + 1), 1); + EXPECT_EQ(eemBlockCount(1000, 20, EemWorkspace::bytes(nvMolKit::kMaxEemAtoms)), 1); +} + +TEST(Descriptors3DEemScratch, BothPrecisionsStayWithinBudget) { + for (const int numAtoms : {120, 256, 512, 1024}) { + for (const int64_t stride : {EemWorkspace::bytes(numAtoms), EemWorkspace::bytes(numAtoms)}) { + const int blocks = eemBlockCount(10000, 100, stride); + EXPECT_GE(blocks, 1); + EXPECT_LE(blocks, 1600); + EXPECT_LE(blocks * stride, kEemScratchBudgetBytes); + if (blocks < 1600) { + EXPECT_GT((blocks + 1) * stride, kEemScratchBudgetBytes); + } + } + } +}