Add GETAWAY 3D descriptors

benchmarks/descriptors3d_bench.py

@@ -84,6 +84,11 @@ 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.GETAWAY:
+        if len(Chem.GetMolFrags(mol)) > 1:
+            # RDKit's GETAWAY does not finish in reasonable time on multi-fragment molecules.
+            return [math.nan] * 273
+        return rdMolDescriptors.CalcGETAWAY(mol, confId=conf_id)
     return RDKIT_CALCULATORS[prop](mol, confId=conf_id, useAtomicMasses=True)
 
 
@@ -151,6 +156,11 @@ def _validate(
     for result in results:
         for prop in properties:
             actual = result[prop.value].numpy()
+            if prop == Property3D.GETAWAY:
+                # Rows RDKit could not compute (multi-fragment molecules) are NaN and not compared.
+                computed = ~np.isnan(expected[prop]).all(axis=1)
+                np.testing.assert_allclose(actual[computed], expected[prop][computed], rtol=0, atol=max(atol, 1.1e-3))
+                continue
             if prop not in (Property3D.USR, Property3D.USRCAT):
                 np.testing.assert_allclose(actual, expected[prop], rtol=rtol, atol=atol)
                 continue
@@ -256,6 +266,11 @@ def run(
                 }
 
                 rdkit_timing = None
+                if not no_rdkit and Property3D.GETAWAY in properties:
+                    skipped = sum(len(Chem.GetMolFrags(mol)) > 1 for mol in mols)
+                    if skipped:
+                        print(f"RDKit GETAWAY skips {skipped} multi-fragment molecules (it would not finish on them)")
+                    row["rdkit_getaway_skipped_mols"] = skipped
                 if not no_rdkit:
                     rdkit_timing = time_it(
                         lambda mols=mols, properties=properties: _calc_rdkit(mols, properties),

docs/api/nvmolkit.rst

@@ -176,6 +176,7 @@ Conformer RMSD
    descriptors3d.Property3DOptions
    descriptors3d.MomentOptions
    descriptors3d.WhimOptions
+   descriptors3d.GetawayOptions
 
 Torsion Fingerprint Deviation (TFD)
 -----------------------------------

nvmolkit/descriptors3d.cpp

@@ -89,6 +89,7 @@ bp::object calc3DProperties(const bp::list&                mols,
                             const bp::list&                propertyNames,
                             const bool                     useAtomicMasses,
                             const double                   whimThreshold,
+                            const unsigned int             getawayPrecision,
                             const bp::object&              coordinates,
                             const nvMolKit::PrecisionMode& precision,
                             const int                      preprocessingThreads,
@@ -100,6 +101,7 @@ bp::object calc3DProperties(const bp::list&                mols,
   nvMolKit::Property3DOptions options;
   options.moments.useAtomicMasses = useAtomicMasses;
   options.whim.threshold          = whimThreshold;
+  options.getaway.precision       = getawayPrecision;
 
   std::vector<nvMolKit::Property3D> properties;
   for (int i = 0; i < bp::len(propertyNames); ++i) {
@@ -131,6 +133,7 @@ BOOST_PYTHON_MODULE(_descriptors3d) {
            bp::arg("properties"),
            bp::arg("useAtomicMasses"),
            bp::arg("whimThreshold"),
+           bp::arg("getawayPrecision"),
            bp::arg("coordinates"),
            bp::arg("precision"),
            bp::arg("preprocessingThreads"),

nvmolkit/descriptors3d.py

@@ -40,6 +40,7 @@ class Property3D(Enum):
       Conformers with fewer than 3 atoms, which RDKit rejects, produce NaN.
     - ``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.
     """
 
     PMI1 = "PMI1"
@@ -59,6 +60,7 @@ class Property3D(Enum):
     AUTOCORR3D = "AUTOCORR3D"
     USR = "USR"
     USRCAT = "USRCAT"
+    GETAWAY = "GETAWAY"
 
 
 @dataclass(frozen=True)
@@ -85,6 +87,18 @@ class WhimOptions:
     threshold: float = 0.001
 
 
+@dataclass(frozen=True)
+class GetawayOptions:
+    """Options for :attr:`Property3D.GETAWAY`.
+
+    Attributes:
+        precision: Significant digits the heavy-atom leverages are rounded to before ``ITH`` and ``ISH``
+            group them, RDKit's default ``2``. Must be between 1 and 6.
+    """
+
+    precision: int = 2
+
+
 @dataclass(frozen=True)
 class Property3DOptions:
     """Per-family options for :func:`Calc3DProperties`; each family reads only its own member.
@@ -106,10 +120,12 @@ class Property3DOptions:
     Attributes:
         moments: Options for the moment-based properties.
         whim: Options for ``WHIM``.
+        getaway: Options for ``GETAWAY``.
     """
 
     moments: MomentOptions = MomentOptions()
     whim: WhimOptions = WhimOptions()
+    getaway: GetawayOptions = GetawayOptions()
 
 
 @dataclass(frozen=True)
@@ -306,16 +322,20 @@ def Calc3DProperties(
             ``values``, or whose atom count differs from their molecule's
             produce NaN rather than an error, so no host synchronization is
             needed. Device coordinate rows are treated as three-dimensional;
-            molecule conformers preserve their RDKit ``is3D`` flag for PBF.
+            molecule conformers keep their RDKit ``is3D`` flag, which PBF and
+            GETAWAY's HIC read. HIC takes its flatness (PBF and ``is3D``) from
+            each molecule's first conformer for all of its conformers; with
+            ``coordinates``, each row uses its own geometry.
         options: Per-family options; defaults to :class:`Property3DOptions`
             (RDKit's defaults).
         precision: ``PrecisionMode.SINGLE`` (default) computes and returns
             float32; ``PrecisionMode.FULL`` computes and returns float64.
             Steps where float32 measurably loses accuracy always compute in
             float64: coordinate centering (conformers far from the origin),
             WHIM's PCA (inverse kurtosis of near-planar conformers),
-            MORSE's first scattering value (a large sum of positive terms)
-            and USR's reference-atom choice (near-equidistant atoms).
+            MORSE's first scattering value (a large sum of positive terms),
+            USR's reference-atom choice (near-equidistant atoms) and GETAWAY's
+            leverages (rounded and clustered for ITH and ISH).
         hardwareOptions: Only ``preprocessingThreads`` applies: the CPU
             threads used to extract coordinates and atom weights from the
             molecules (default ``-1``, all threads).
@@ -379,6 +399,13 @@ def Calc3DProperties(
         options = Property3DOptions()
     elif not isinstance(options, Property3DOptions):
         raise TypeError(f"options must be a Property3DOptions or None, got {type(options).__name__}")
+    getaway_precision = GetawayOptions().precision
+    if Property3D.GETAWAY in normalized_properties:
+        getaway_precision = options.getaway.precision
+        if isinstance(getaway_precision, bool) or not isinstance(getaway_precision, int):
+            raise TypeError(f"GETAWAY precision must be an int, got {type(getaway_precision).__name__}")
+        if not 1 <= getaway_precision <= 6:
+            raise ValueError(f"GETAWAY precision must be between 1 and 6 significant digits, got {getaway_precision}")
     if hardwareOptions is None:
         hardwareOptions = HardwareOptions()
     elif not isinstance(hardwareOptions, HardwareOptions):
@@ -398,6 +425,7 @@ def Calc3DProperties(
         [prop.value for prop in normalized_properties],
         options.moments.useAtomicMasses,
         options.whim.threshold,
+        getaway_precision,
         coordinate_interfaces,
         precision,
         hardwareOptions.preprocessingThreads,

nvmolkit/tests/test_descriptors3d.py

@@ -14,6 +14,7 @@
 from nvmolkit.descriptors3d import (
     Calc3DProperties,
     Device3DPropertyResult,
+    GetawayOptions,
     MomentOptions,
     Property3D,
     Property3DOptions,
@@ -30,6 +31,7 @@
     Property3D.AUTOCORR3D,
     Property3D.USR,
     Property3D.USRCAT,
+    Property3D.GETAWAY,
 )
 SCALAR_PROPERTIES = tuple(prop for prop in Property3D if prop not in VECTOR_PROPERTIES)
 PAIRWISE_PROPERTIES = (Property3D.RDF, Property3D.MORSE, Property3D.AUTOCORR3D)
@@ -547,6 +549,119 @@ def test_autocorr3d_non_finite_coordinates_match_rdkit(precision):
     )
 
 
+def _rdkit_getaway_rows(mols, precision_digits=2):
+    return np.asarray(
+        [
+            rdMolDescriptors.CalcGETAWAY(mol, confId=conf.GetId(), precision=precision_digits)
+            for mol in mols
+            for conf in mol.GetConformers()
+        ]
+    )
+
+
+@pytest.mark.parametrize("precision", PRECISIONS)
+def test_getaway_matches_rdkit(precision):
+    rng = np.random.default_rng(11)
+    steps = rng.normal(size=(60, 3))
+    chain = np.cumsum(1.5 * steps / np.linalg.norm(steps, axis=1, keepdims=True), axis=0)
+    mols = [
+        _embed("CC(=O)Nc1ccc(O)cc1", 3, 97),
+        # Planar: X^T X has a zero singular value, dropped from the pseudo-inverse, and HIC uses D = 2.
+        _embed("c1ccncc1", 2, 101),
+        _mol_with_conformers("C" * 60, [chain]),
+    ]
+    result = Calc3DProperties(mols, Property3D.GETAWAY, precision=precision)
+    assert result[Property3D.GETAWAY].torch().shape == (6, 273)
+    _assert_rounded_matches_rdkit(result[Property3D.GETAWAY].numpy(), _rdkit_getaway_rows(mols), precision)
+
+
+@pytest.fixture(scope="module")
+def getaway_salt():
+    # RDKit's GETAWAY takes seconds on this 4-atom salt, so its reference is computed once.
+    mol = _mol_with_conformers("CCO.N", [[(0.0, 0.0, 0.0), (1.5, 0.1, 0.0), (2.1, 1.4, 0.2), (5.0, 1.0, -1.0)]])
+    return mol, _rdkit_getaway_rows([mol])
+
+
+@pytest.mark.parametrize("precision", PRECISIONS)
+def test_getaway_counts_disconnected_pairs_in_totals_like_rdkit(getaway_salt, precision):
+    mol, expected = getaway_salt
+    result = Calc3DProperties([mol], Property3D.GETAWAY, precision=precision)
+    _assert_rounded_matches_rdkit(result[Property3D.GETAWAY].numpy(), expected, precision)
+
+
+@pytest.mark.parametrize("precision", PRECISIONS)
+def test_getaway_hic_uses_default_conformer_dimension(precision):
+    # RDKit's HIC takes its 2D/3D choice from PBF(mol), i.e. the default (first) conformer: planar here.
+    planar = [(0.0, 0.0, 0.0), (1.5, 0.0, 0.0), (2.1, 1.4, 0.0), (3.6, 1.5, 0.0)]
+    bent = [(0.0, 0.0, 0.0), (1.5, 0.1, 0.3), (2.1, 1.4, -0.4), (3.6, 1.5, 0.6)]
+    mol = _mol_with_conformers("CCCO", [planar, bent])
+    result = Calc3DProperties([mol], Property3D.GETAWAY, precision=precision)
+    _assert_rounded_matches_rdkit(result[Property3D.GETAWAY].numpy(), _rdkit_getaway_rows([mol]), precision)
+
+
+@pytest.mark.parametrize("precision", PRECISIONS)
+def test_getaway_tiny_leverage_matches_rdkit(precision):
+    # The central carbon sits 0.003 A from the centroid: its leverage (~2e-6) prints in scientific notation
+    # in RDKit's ITH/ISH digit-string clustering.
+    corner = 0.89
+    coordinates = [
+        (0.003, 0.0, 0.0),
+        (corner, corner, corner),
+        (corner, -corner, -corner),
+        (-corner, corner, -corner),
+        (-corner, -corner, corner),
+    ]
+    mol = _mol_with_conformers("C(C)(C)(C)C", [coordinates])
+    result = Calc3DProperties([mol], Property3D.GETAWAY, precision=precision)
+    _assert_rounded_matches_rdkit(result[Property3D.GETAWAY].numpy(), _rdkit_getaway_rows([mol]), precision)
+
+
+@pytest.mark.parametrize("precision_digits", [1, 3])
+def test_getaway_precision_option_matches_rdkit(precision_digits):
+    mols = [_embed("CC(=O)Nc1ccc(O)cc1", 2, 103)]
+    result = Calc3DProperties(
+        mols,
+        Property3D.GETAWAY,
+        options=Property3DOptions(getaway=GetawayOptions(precision=precision_digits)),
+        precision=PrecisionMode.FULL,
+    )
+    _assert_rounded_matches_rdkit(
+        result[Property3D.GETAWAY].numpy(), _rdkit_getaway_rows(mols, precision_digits), PrecisionMode.FULL
+    )
+
+
+@pytest.mark.parametrize("precision_digits", [-1, 0, 7])
+def test_getaway_rejects_out_of_range_precision(precision_digits):
+    options = Property3DOptions(getaway=GetawayOptions(precision=precision_digits))
+    with pytest.raises(ValueError, match="GETAWAY precision"):
+        Calc3DProperties([_embed("CCO", 1, 5)], Property3D.GETAWAY, options=options)
+
+
+def test_getaway_rejects_non_integer_precision():
+    options = Property3DOptions(getaway=GetawayOptions(precision=2.5))
+    with pytest.raises(TypeError, match="GETAWAY precision"):
+        Calc3DProperties([_embed("CCO", 1, 5)], Property3D.GETAWAY, options=options)
+
+
+@pytest.mark.parametrize("precision", PRECISIONS)
+def test_getaway_non_finite_coordinates_give_nan_information_indices(precision):
+    mols = [
+        _mol_with_conformers("CCCO", [[(0.0, 0.0, 0.0), (1.5, float("nan"), 0.0), (2.0, 1.4, 0.0), (3.4, 1.5, 0.2)]]),
+        _mol_with_conformers("CCCO", [[(0.0, 0.0, 0.0), (1.5, 0.2, 0.0), (2.0, 1.4, float("inf")), (3.4, 1.5, 0.2)]]),
+        _embed("CCCO", 1, 47),
+    ]
+    values = Calc3DProperties(mols, Property3D.GETAWAY, precision=precision)[Property3D.GETAWAY].numpy()
+    assert np.isnan(values[:2, :2]).all()  # ITH and ISH
+    _assert_rounded_matches_rdkit(values[2:], _rdkit_getaway_rows(mols[2:]), precision)
+
+
+@pytest.mark.parametrize("precision_digits", [0, 2.5])
+def test_getaway_precision_ignored_when_not_requested(precision_digits):
+    options = Property3DOptions(getaway=GetawayOptions(precision=precision_digits))
+    result = Calc3DProperties([_embed("CCO", 1, 5)], Property3D.PMI1, options=options)
+    assert result[Property3D.PMI1].torch().shape == (1,)
+
+
 @pytest.mark.parametrize("precision", PRECISIONS)
 def test_autocorr3d_accepts_batches_without_bonds(precision):
     mols = [

src/descriptors3d.cu

@@ -7,10 +7,12 @@
 #include <string>
 
 #include "src/descriptors3d.h"
+#include "src/descriptors3d_getaway.cuh"
 #include "src/descriptors3d_kernel.cuh"
 #include "src/descriptors3d_moments.cuh"
 #include "src/descriptors3d_pairwise.cuh"
 #include "src/descriptors3d_projection.cuh"
+#include "src/descriptors3d_topology.cuh"
 #include "src/descriptors3d_usr.cuh"
 #include "src/utils/cuda_error_check.h"
 
@@ -52,6 +54,8 @@ std::string_view property3DName(const Property3D property) {
       return "USR";
     case Property3D::USRCAT:
       return "USRCAT";
+    case Property3D::GETAWAY:
+      return "GETAWAY";
   }
   throw std::invalid_argument("Unknown Property3D value " + std::to_string(static_cast<int>(property)));
 }
@@ -77,6 +81,7 @@ using descriptors3d_detail::kGroupSize;
 using descriptors3d_detail::kGroupsPerWarp;
 using descriptors3d_detail::kWarpSize;
 using descriptors3d_detail::kWarpsPerBlock;
+using descriptors3d_detail::launchGetawayProperties;
 using descriptors3d_detail::launchPairwiseProperties;
 using descriptors3d_detail::launchProjectionProperties;
 using descriptors3d_detail::launchUsrProperties;
@@ -129,6 +134,7 @@ constexpr SharedStageSet directStages(const Property3D property) {
     case Property3D::AUTOCORR3D:
     case Property3D::USR:
     case Property3D::USRCAT:
+    case Property3D::GETAWAY:
       return 0;
   }
   return 0;
@@ -293,8 +299,9 @@ Property3DResults<Real> calc3DPropertiesGpu(const DeviceCoordView&         coord
   Real*                   pbfOutput     = nullptr;
   Real*                   whimOutput    = nullptr;
   PairwiseOutputs<Real>   pairwiseOutputs;
-  Real*                   usrOutput    = nullptr;
-  Real*                   usrcatOutput = nullptr;
+  Real*                   usrOutput     = nullptr;
+  Real*                   usrcatOutput  = nullptr;
+  Real*                   getawayOutput = 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()) {
@@ -321,13 +328,19 @@ Property3DResults<Real> calc3DPropertiesGpu(const DeviceCoordView&         coord
       case Property3DFamily::Usr:
         (property == Property3D::USR ? usrOutput : usrcatOutput) = it->second.data();
         break;
+      case Property3DFamily::Getaway:
+        getawayOutput = it->second.data();
+        break;
     }
   }
 
   // Options are validated only for requested families; the rest are ignored.
   if (whimOutput != nullptr && (!std::isfinite(options.whim.threshold) || options.whim.threshold < 0)) {
     throw std::invalid_argument("WHIM threshold must be finite and non-negative");
   }
+  if (getawayOutput != nullptr && (options.getaway.precision < 1 || options.getaway.precision > 6)) {
+    throw std::invalid_argument("GETAWAY precision must be between 1 and 6 significant digits");
+  }
 
   if (coordinates.numConformers == 0) {
     return results;
@@ -338,9 +351,13 @@ Property3DResults<Real> calc3DPropertiesGpu(const DeviceCoordView&         coord
   }
   const bool anyPairwise =
     pairwiseOutputs.rdf != nullptr || pairwiseOutputs.morse != nullptr || pairwiseOutputs.autocorr3D != nullptr;
-  if ((whimOutput != nullptr || anyPairwise) && inputs.atomPropertyWeights == nullptr) {
+  if ((whimOutput != nullptr || anyPairwise || getawayOutput != nullptr) && inputs.atomPropertyWeights == nullptr) {
     throw std::invalid_argument(
-      "Atom-property weights must not be null when WHIM, RDF, MORSE or AUTOCORR3D is requested");
+      "Atom-property weights must not be null when WHIM, RDF, MORSE, AUTOCORR3D or GETAWAY is requested");
+  }
+  if (getawayOutput != nullptr &&
+      (inputs.bondNeighborStarts == nullptr || inputs.bondNeighbors == nullptr || inputs.heavyAtomFlags == nullptr)) {
+    throw std::invalid_argument("Bond adjacency and heavy-atom flags must not be null when GETAWAY is requested");
   }
   if (pairwiseOutputs.autocorr3D != nullptr &&
       (inputs.covalentRadiusWeights == nullptr || inputs.bondNeighborStarts == nullptr ||
@@ -360,8 +377,14 @@ Property3DResults<Real> calc3DPropertiesGpu(const DeviceCoordView&         coord
   const double* momentWeights           = onlySpherocity ? nullptr : inputs.momentWeights;
   launchMomentProperties(coordinates, momentWeights, inputs.moleculeAtomStarts, work, separateSpherocityState, stream);
   launchProjectionProperties(coordinates, inputs, options.whim, pbfOutput, whimOutput, stream);
-  launchPairwiseProperties(coordinates, inputs, pairwiseOutputs, stream);
+  // Bond distances depend only on the molecule, so AUTOCORR3D and GETAWAY share one table for all conformers.
+  descriptors3d_detail::BondDistanceTableStorage bondTable;
+  if (pairwiseOutputs.autocorr3D != nullptr || getawayOutput != nullptr) {
+    descriptors3d_detail::buildBondDistanceTable(coordinates, inputs, bondTable, stream);
+  }
+  launchPairwiseProperties(coordinates, inputs, pairwiseOutputs, bondTable.table, stream);
   launchUsrProperties(coordinates, inputs, usrOutput, usrcatOutput, stream);
+  launchGetawayProperties(coordinates, inputs, options.getaway, bondTable.table, getawayOutput, stream);
   return results;
 }