DFS Substruct backend algorithm

benchmarks/substruct_bench.py

@@ -35,6 +35,9 @@
     # Get all matches instead of just boolean:
     python substruct_bench.py --smiles <smiles_file> --smarts <smarts_file> --mode getSubstructMatches
 
+    # Select the adjacency-anchored DFS backend (GSI is the default):
+    python substruct_bench.py --smiles <smiles_file> --smarts <smarts_file> --algorithm dfs
+
     # Limit to first 10 matches per target/query pair:
     python substruct_bench.py --smiles <smiles_file> --smarts <smarts_file> --mode getSubstructMatches --max_matches 10
 
@@ -51,7 +54,8 @@
     python substruct_bench.py --smiles <smiles_file> --smarts <smarts_file> \
         --rdkit_match_mode raw substructlib --rdkit_threads 1 4 16
 
-    # Run multiple configurations from a dataframe (smarts, batch_size, workers, prep_threads, mode, num_gpus):
+    # Run multiple configurations from a dataframe
+    # (smarts, batch_size, workers, prep_threads, mode, num_gpus, optional algorithm):
     python substruct_bench.py --smiles <smiles_file> --config <config.csv>
 
 """
@@ -67,15 +71,16 @@
 import pandas as pd
 from bench_utils import add_rdkit_max_seconds_arg, load_pickle, load_smarts, load_smiles, time_it_bounded
 from benchmark_timing import time_it as _time_it
+from rdkit import Chem
+from rdkit.Chem import rdSubstructLibrary
+
 from nvmolkit import autotune as nv_autotune
 from nvmolkit.substructure import (
     SubstructSearchConfig,
     countSubstructMatches,
     getSubstructMatches,
     hasSubstructMatch,
 )
-from rdkit import Chem
-from rdkit.Chem import rdSubstructLibrary
 
 OPTUNA_AVAILABLE = nv_autotune.is_available()
 
@@ -325,7 +330,7 @@ def main():
         "--config",
         help=(
             "Path to config dataframe (.csv/.pkl/.pickle/.parquet) with columns: "
-            "smarts, batch_size, workers, prep_threads, mode, num_gpus"
+            "smarts, batch_size, workers, prep_threads, mode, num_gpus, optional algorithm"
         ),
     )
     parser.add_argument("--num_mols", "-n", type=int, default=0, help="Max number of molecules (default: 0 = all)")
@@ -346,6 +351,12 @@ def main():
         default="hasSubstructMatch",
         help="Search mode (default: hasSubstructMatch)",
     )
+    parser.add_argument(
+        "--algorithm",
+        choices=["gsi", "dfs"],
+        default=None,
+        help="nvmolkit matching backend (default: gsi)",
+    )
     parser.add_argument(
         "--max_matches", type=int, default=0, help="Maximum matches per target/query pair, 0 = all (default: 0)"
     )
@@ -406,7 +417,8 @@ def main():
         default=None,
         help=(
             "Path to a previously-saved SubstructSearchConfig JSON. "
-            "Overrides --batch_size/--workers/--prep_threads (and --num_gpus if gpuIds present in the file)."
+            "Overrides --batch_size/--workers/--prep_threads (and --num_gpus if gpuIds present in the file); "
+            "its algorithm is retained unless --algorithm is specified."
         ),
     )
     parser.add_argument(
@@ -509,7 +521,11 @@ def main():
         print(f"  SMARTS file: {args.smarts}")
         print(f"  Mode: {args.mode}")
         if not args.no_nvmolkit:
-            print(f"  nvmolkit config:")
+            print("  nvmolkit config:")
+            if args.autotune_load and args.algorithm is None:
+                print("    algorithm: from loaded config")
+            else:
+                print(f"    algorithm: {args.algorithm or 'gsi'}")
             print(f"    batch_size: {args.batch_size}")
             print(f"    num_gpus: {args.num_gpus}")
             print(f"    workers: {args.workers if args.workers >= 0 else 'auto'}")
@@ -536,6 +552,7 @@ def main():
                 "prep_threads": args.prep_threads,
                 "mode": args.mode,
                 "num_gpus": args.num_gpus,
+                "algorithm": args.algorithm or "gsi",
             }
         ]
 
@@ -545,12 +562,21 @@ def main():
     for config_row in config_rows:
         smarts_path = config_row["smarts"]
         mode = config_row["mode"]
+        row_algorithm = config_row.get("algorithm", args.algorithm or "gsi")
+        algorithm = "gsi" if pd.isna(row_algorithm) else str(row_algorithm).lower()
+        if algorithm not in {"gsi", "dfs"}:
+            print(f"Error: algorithm must be 'gsi' or 'dfs', got {row_algorithm!r}")
+            sys.exit(1)
 
         print("\nRun configuration:")
         print(f"  SMARTS file: {smarts_path}")
         print(f"  Mode: {mode}")
         if not args.no_nvmolkit:
-            print(f"  nvmolkit config:")
+            print("  nvmolkit config:")
+            if args.autotune_load and args.algorithm is None:
+                print("    algorithm: from loaded config")
+            else:
+                print(f"    algorithm: {algorithm}")
             print(f"    batch_size: {config_row['batch_size']}")
             print(f"    num_gpus: {config_row['num_gpus']}")
             print(f"    workers: {config_row['workers'] if config_row['workers'] >= 0 else 'auto'}")
@@ -600,7 +626,8 @@ def main():
                     if not config.gpuIds:
                         config.gpuIds = gpu_ids
                     print(
-                        f"  Loaded: batchSize={config.batchSize}, workerThreads={config.workerThreads}, "
+                        f"  Loaded: algorithm={config.algorithm}, batchSize={config.batchSize}, "
+                        f"workerThreads={config.workerThreads}, "
                         f"preprocessingThreads={config.preprocessingThreads}, gpuIds={list(config.gpuIds)}"
                     )
                 elif args.autotune:
@@ -625,6 +652,7 @@ def main():
                         api=api_for_mode,
                         maxMatches=args.max_matches,
                         gpuIds=gpu_ids,
+                        algorithm=algorithm,
                         n_trials=args.autotune_trials,
                         target_seconds_per_trial=args.autotune_time_budget,
                         calibration_set=explicit_calibration,
@@ -650,6 +678,9 @@ def main():
                     config.gpuIds = gpu_ids
                     if args.max_matches > 0:
                         config.maxMatches = args.max_matches
+                if not args.autotune_load or args.algorithm is not None:
+                    config.algorithm = algorithm
+                algorithm = config.algorithm
 
                 ran_nvmolkit = True
                 torch_module = torch
@@ -774,6 +805,7 @@ def main():
             applied_workers = config_row["workers"]
             applied_prep_threads = config_row["prep_threads"]
             applied_num_gpus = config_row["num_gpus"]
+        applied_algorithm = algorithm if ran_nvmolkit else "N/A"
 
         if args.autotune:
             config_source = "autotuned"
@@ -805,6 +837,7 @@ def main():
                 (
                     name,
                     mode,
+                    applied_algorithm if name == "nvmolkit" else "N/A",
                     smarts_path,
                     input_file,
                     input_type,
@@ -836,7 +869,7 @@ def main():
 
     print("\n\nCSV Results:")
     print(
-        "method,mode,smarts,input_file,input_type,sanitize,num_mols,num_patterns,"
+        "method,mode,algorithm,smarts,input_file,input_type,sanitize,num_mols,num_patterns,"
         "max_matches,batch_size,num_gpus,workers,prep_threads,nvmolkit_config_source,"
         "rdkit_threads,rdkit_match_mode,time_ms,std_ms,"
         "pairs_processed,rdkit_max_seconds,pairs_per_second,vs_rdkit_throughput_ratio"
@@ -845,6 +878,7 @@ def main():
         (
             name,
             mode,
+            algorithm,
             smarts_path,
             input_file,
             input_type,
@@ -871,7 +905,7 @@ def main():
             f"{rdkit_max_seconds:g}" if isinstance(rdkit_max_seconds, float) else str(rdkit_max_seconds)
         )
         print(
-            f"{name},{mode},{smarts_path},{input_file},{input_type},{sanitize},"
+            f"{name},{mode},{algorithm},{smarts_path},{input_file},{input_type},{sanitize},"
             f"{num_mols},{num_patterns},{max_matches},{batch_size},{num_gpus},{workers},{prep_threads},"
             f"{nvmolkit_config_source},{rdkit_threads},{rdkit_match_mode},{avg_ms:.2f},{std_ms:.2f},"
             f"{pairs_done},{rdkit_max_seconds_str},{throughput:.2f},{vs_rdkit_str}"

nvmolkit/autotune/tune_substructure.py

@@ -119,6 +119,7 @@ def tune_substructure(
     maxMatches: int = 0,
     uniquify: bool = False,
     gpuIds: Optional[Iterable[int]] = None,
+    algorithm: str = "gsi",
     calibration_set: Optional[Iterable[int]] = None,
     calibration_fraction: float = 0.1,
     calibration_max_size: int = 50_000,
@@ -143,6 +144,7 @@ def tune_substructure(
             Held constant across trials.
         uniquify: ``uniquify`` flag forwarded to the resulting config.
         gpuIds: GPU device IDs to use. Fixed across the study.
+        algorithm: Matching backend to tune, either ``"gsi"`` or ``"dfs"``.
         calibration_set: Optional explicit indices into ``targets``.
         calibration_fraction: Fraction of the workload to auto-sample.
         calibration_max_size: Cap on the auto-sampled calibration size.
@@ -194,6 +196,7 @@ def _make_config(values: dict[str, Any]) -> SubstructSearchConfig:
             maxMatches=int(maxMatches),
             uniquify=bool(uniquify),
             gpuIds=fixed_gpu_ids if fixed_gpu_ids else None,
+            algorithm=algorithm,
         )
 
     queries_list = list(queries)
@@ -209,6 +212,7 @@ def default_runner(state: CalibrationState) -> int:
                 maxMatches=int(maxMatches),
                 uniquify=bool(uniquify),
                 gpuIds=fixed_gpu_ids if fixed_gpu_ids else None,
+                algorithm=algorithm,
             ),
             state,
         )

nvmolkit/substructure.cpp

@@ -20,6 +20,8 @@
 #include <boost/python/stl_iterator.hpp>
 #include <cstdint>
 #include <memory>
+#include <stdexcept>
+#include <string>
 #include <vector>
 
 #include "nvmolkit/boost_python_utils.h"
@@ -92,6 +94,28 @@ void setGpuIdsPy(nvMolKit::SubstructSearchConfig& config, const object& iterable
   config.gpuIds = listFromIterable<int>(iterable);
 }
 
+std::string getAlgorithmPy(const nvMolKit::SubstructSearchConfig& config) {
+  switch (config.algorithm) {
+    case nvMolKit::SubstructAlgorithm::GSI:
+      return "gsi";
+    case nvMolKit::SubstructAlgorithm::DFS:
+      return "dfs";
+    case nvMolKit::SubstructAlgorithm::VF2:
+      throw std::invalid_argument("VF2 is not supported by the Python substructure bindings");
+  }
+  throw std::invalid_argument("Unknown substructure algorithm");
+}
+
+void setAlgorithmPy(nvMolKit::SubstructSearchConfig& config, const std::string& algorithm) {
+  if (algorithm == "gsi") {
+    config.algorithm = nvMolKit::SubstructAlgorithm::GSI;
+  } else if (algorithm == "dfs") {
+    config.algorithm = nvMolKit::SubstructAlgorithm::DFS;
+  } else {
+    throw std::invalid_argument("algorithm must be 'gsi' or 'dfs'");
+  }
+}
+
 }  // namespace
 
 BOOST_PYTHON_MODULE(_substructure) {
@@ -104,7 +128,8 @@ BOOST_PYTHON_MODULE(_substructure) {
     .def_readwrite("preprocessingThreads", &nvMolKit::SubstructSearchConfig::preprocessingThreads)
     .def_readwrite("maxMatches", &nvMolKit::SubstructSearchConfig::maxMatches)
     .def_readwrite("uniquify", &nvMolKit::SubstructSearchConfig::uniquify)
-    .add_property("gpuIds", &getGpuIdsPy, &setGpuIdsPy);
+    .add_property("gpuIds", &getGpuIdsPy, &setGpuIdsPy)
+    .add_property("algorithm", &getAlgorithmPy, &setAlgorithmPy);
 
   def(
     "getSubstructMatches",
@@ -133,12 +158,7 @@ BOOST_PYTHON_MODULE(_substructure) {
       extractRange.pop();
 
       nvMolKit::SubstructSearchResults results;
-      nvMolKit::getSubstructMatches(targetsVec,
-                                    queriesVec,
-                                    results,
-                                    nvMolKit::SubstructAlgorithm::GSI,
-                                    nullptr,
-                                    config);
+      nvMolKit::getSubstructMatches(targetsVec, queriesVec, results, config.algorithm, nullptr, config);
 
       auto csrPtr        = std::make_unique<SubstructMatchesCSR>();
       csrPtr->numTargets = results.numTargets;
@@ -244,12 +264,7 @@ BOOST_PYTHON_MODULE(_substructure) {
       const int numQueries = static_cast<int>(queriesVec.size());
 
       auto countsPtr = std::make_unique<std::vector<int>>();
-      nvMolKit::countSubstructMatches(targetsVec,
-                                      queriesVec,
-                                      *countsPtr,
-                                      nvMolKit::SubstructAlgorithm::GSI,
-                                      nullptr,
-                                      config);
+      nvMolKit::countSubstructMatches(targetsVec, queriesVec, *countsPtr, config.algorithm, nullptr, config);
 
       nvMolKit::ScopedNvtxRange wrapRange("Python: wrap numpy array", nvMolKit::NvtxColor::kGreen);
       int*                      dataPtr = countsPtr->data();
@@ -304,12 +319,7 @@ BOOST_PYTHON_MODULE(_substructure) {
       extractRange.pop();
 
       auto resultsPtr = std::make_unique<nvMolKit::HasSubstructMatchResults>();
-      nvMolKit::hasSubstructMatch(targetsVec,
-                                  queriesVec,
-                                  *resultsPtr,
-                                  nvMolKit::SubstructAlgorithm::GSI,
-                                  nullptr,
-                                  config);
+      nvMolKit::hasSubstructMatch(targetsVec, queriesVec, *resultsPtr, config.algorithm, nullptr, config);
 
       nvMolKit::ScopedNvtxRange wrapRange("Python: wrap numpy array", nvMolKit::NvtxColor::kGreen);
       const int                 numTargets = resultsPtr->numTargets;

nvmolkit/substructure.py

@@ -49,6 +49,13 @@
 ]
 
 
+def _normalize_algorithm(value: str) -> str:
+    algorithm = str(value).lower()
+    if algorithm not in {"gsi", "dfs"}:
+        raise ValueError("algorithm must be 'gsi' or 'dfs'")
+    return algorithm
+
+
 class SubstructSearchConfig:
     """Configuration for GPU substructure search execution.
 
@@ -70,6 +77,7 @@ def __init__(
         maxMatches: int = 0,
         uniquify: bool = False,
         gpuIds: list[int] | None = None,
+        algorithm: str = "dfs",
     ) -> None:
         native = _NativeSubstructSearchConfig()
         native.batchSize = int(batchSize)
@@ -78,6 +86,7 @@ def __init__(
         native.maxMatches = int(maxMatches)
         native.uniquify = bool(uniquify)
         native.gpuIds = list(gpuIds) if gpuIds is not None else []
+        native.algorithm = _normalize_algorithm(algorithm)
         self._native = native
 
     @property
@@ -134,6 +143,20 @@ def gpuIds(self) -> list[int]:
     def gpuIds(self, value: list[int]) -> None:
         self._native.gpuIds = list(value)
 
+    @property
+    def algorithm(self) -> str:
+        """Internal: matching backend, ``"dfs"`` (default) or ``"gsi"``.
+
+        Not part of the public API -- exposed for benchmarking and testing
+        the two backends against each other, and may be removed once ``gsi``
+        is retired.
+        """
+        return self._native.algorithm
+
+    @algorithm.setter
+    def algorithm(self, value: str) -> None:
+        self._native.algorithm = _normalize_algorithm(value)
+
     def _as_native(self):
         """Internal: return the underlying native config object."""
         return self._native
@@ -147,12 +170,13 @@ def to_dict(self) -> dict[str, Any]:
             "maxMatches": self.maxMatches,
             "uniquify": self.uniquify,
             "gpuIds": list(self.gpuIds),
+            "algorithm": self.algorithm,
         }
 
     @classmethod
     def from_dict(cls, data: dict[str, Any]) -> "SubstructSearchConfig":
         """Create a :class:`SubstructSearchConfig` from a dictionary produced by :meth:`to_dict`."""
-        known = {"batchSize", "workerThreads", "preprocessingThreads", "maxMatches", "uniquify", "gpuIds"}
+        known = {"batchSize", "workerThreads", "preprocessingThreads", "maxMatches", "uniquify", "gpuIds", "algorithm"}
         unknown = set(data) - known
         if unknown:
             raise ValueError(f"Unknown SubstructSearchConfig keys: {sorted(unknown)}")

nvmolkit/tests/test_autotune.py

@@ -163,6 +163,7 @@ def test_substruct_config_to_from_dict_roundtrip():
         maxMatches=8,
         uniquify=True,
         gpuIds=[0],
+        algorithm="dfs",
     )
     encoded = config.to_dict()
     assert encoded == {
@@ -172,6 +173,7 @@ def test_substruct_config_to_from_dict_roundtrip():
         "maxMatches": 8,
         "uniquify": True,
         "gpuIds": [0],
+        "algorithm": "dfs",
     }
     restored = SubstructSearchConfig.from_dict(encoded)
     assert restored.batchSize == 512
@@ -180,6 +182,7 @@ def test_substruct_config_to_from_dict_roundtrip():
     assert restored.maxMatches == 8
     assert restored.uniquify is True
     assert restored.gpuIds == [0]
+    assert restored.algorithm == "dfs"
 
 
 def test_save_load_hardware_options_roundtrip(tmp_path):