tune_batched_forcefield#

nvmolkit.autotune.tune_batched_forcefield(
molecules: list[Mol],
factory: Callable[[list[Mol], HardwareOptions], MMFFBatchedForcefield | UFFBatchedForcefield],
*,
maxIters: int | None = None,
forceTol: float | None = None,
gpuIds: Iterable[int] | None = None,
calibration_set: Iterable[int] | None = None,
calibration_fraction: float = 0.1,
calibration_max_size: int = 2000,
target_seconds_per_trial: float = 10.0,
n_trials: int = 30,
search_space_overrides: dict[str, Any] | None = None,
cpu_budget: int | None = None,
sampler: Any = None,
seed: int | None = None,
verbose: bool = False,
) → TuneResult#

Tune HardwareOptions for a batched-forcefield workflow.

Parameters:

  • molecules – Workload of pre-embedded RDKit molecules.

  • factory – Callable invoked once per trial with the cloned calibration molecule slice and the trial-suggested HardwareOptions. Should return a fully-configured MMFFBatchedForcefield or UFFBatchedForcefield (with any per-element properties or constraints already applied).

  • maxIters – Optional override for the minimize maxIters argument. When None, the default of the returned forcefield class is used.

  • forceTol – Optional override for the minimize forceTol argument. When None, the default of the returned forcefield class is used.

  • gpuIds – GPU device IDs to use. Fixed across the study.

  • calibration_set – Optional explicit indices into molecules.

  • calibration_fraction – Fraction of the workload to auto-sample.

  • calibration_max_size – Cap on the auto-sampled calibration size.

  • target_seconds_per_trial – Target wall-clock budget for one trial.

  • n_trials – Number of Optuna trials to run after warm-up.

  • search_space_overrides – Optional overrides for batchSize / batchesPerGpu ranges.

  • cpu_budget – Optional explicit cap on total CPU threads. The default (None) uses os.cpu_count(). Set this when normalizing tuning runs across machines with different core counts so the search space stays comparable.

  • sampler – Optional Optuna sampler.

  • seed – Seed for the default sampler.

  • verbose – Print warm-up and trial diagnostics.

Returns:

TuneResult with best_config set to a fully-populated HardwareOptions instance.