Performance Tuning Guide

This notebook covers optimization strategies for getting the best performance from Auto3D.

CLI Quick Reference (Recommended)

# Quick screening (fastest)
auto3d run input.smi --k=1 --engine=ANI2xt --gpu

# Standard workflow
auto3d run input.smi --k=1 --engine=AIMNET --gpu

# High accuracy (with ensemble)
AUTO3D_USE_ENSEMBLE=1 auto3d run input.smi --k=1 --engine=AIMNET --gpu

# With torch.compile optimization (ANI models)
AUTO3D_COMPILE_MODEL=1 auto3d run input.smi --k=1 --engine=ANI2x --gpu

# Multi-GPU processing
auto3d run large_dataset.smi --k=1 --gpu --gpu-idx="0,1,2,3"

# Generate preset configurations
auto3d config init -p quick -o quick.yaml      # Fast screening
auto3d config init -p balanced -o balanced.yaml # Default settings
auto3d config init -p thorough -o thorough.yaml # High accuracy

# Run with configuration file
auto3d run input.smi -c quick.yaml

Optimal Settings by Use Case:

Use Case	CLI Command
Quick screening	auto3d run input.smi --k=1 --engine=ANI2xt --gpu
Standard workflow	auto3d run input.smi --k=1 --engine=AIMNET --gpu
High accuracy	AUTO3D_USE_ENSEMBLE=1 auto3d run input.smi --k=1 --gpu
Multi-GPU	auto3d run input.smi --k=1 --gpu --gpu-idx="0,1,2,3"

---

Python API Performance Guide

Below we explore performance optimization using the Python API:

import time
import torch
import Auto3D
from Auto3D import Auto3DOptions, smiles2mols, create_model

print(f"Auto3D: {Auto3D.__version__}")
print(f"PyTorch: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")

1. Choosing the Right Engine

Speed comparison of available engines:

# Test molecules
test_smiles = ["c1ccc2c(c1)cccc2", "CCCCc1ccccc1", "O=C(NC)Nc1ccccc1"]

engines = ["ANI2xt", "ANI2x", "AIMNET"]

for engine in engines:
    start = time.time()
    config = Auto3DOptions(k=1, optimizing_engine=engine, use_gpu=False)
    mols = smiles2mols(test_smiles, config)
    elapsed = time.time() - start
    print(f"{engine}: {elapsed:.2f}s")

Recommendations:

• ANI2xt: Fastest, use for initial screening

• ANI2x: Good balance of speed and accuracy

• AIMNET: Most accurate, supports charged molecules

2. GPU Acceleration

GPU provides significant speedup for large batches:

if torch.cuda.is_available():
    # CPU timing
    start = time.time()
    config = Auto3DOptions(k=1, use_gpu=False)
    mols_cpu = smiles2mols(test_smiles, config)
    cpu_time = time.time() - start

    # GPU timing
    start = time.time()
    config = Auto3DOptions(k=1, use_gpu=True)
    mols_gpu = smiles2mols(test_smiles, config)
    gpu_time = time.time() - start

    print(f"CPU: {cpu_time:.2f}s")
    print(f"GPU: {gpu_time:.2f}s")
    print(f"Speedup: {cpu_time/gpu_time:.1f}x")
else:
    print("No GPU available for comparison")

3. Single Model vs Ensemble (AIMNET)

By default, Auto3D uses a single AIMNet2 model for ~35x faster optimization:

import os

# Single model (default, faster)
os.environ["AUTO3D_USE_ENSEMBLE"] = "0"
start = time.time()
config = Auto3DOptions(k=1, optimizing_engine="AIMNET", use_gpu=False)
mols = smiles2mols(test_smiles[:1], config)
single_time = time.time() - start

print(f"Single model: {single_time:.2f}s")
print("Use ensemble=True for highest accuracy (slower)")

4. TF32 Acceleration (Ampere+ GPUs)

Enable TensorFloat-32 for faster matrix operations on RTX 30xx, A100, H100:

if torch.cuda.is_available():
    # Check GPU architecture
    capability = torch.cuda.get_device_capability(0)
    is_ampere_plus = capability[0] >= 8

    if is_ampere_plus:
        print(f"GPU supports TF32 (compute capability {capability[0]}.{capability[1]})")

        # With TF32
        config = Auto3DOptions(k=1, use_gpu=True, allow_tf32=True)
        start = time.time()
        mols = smiles2mols(test_smiles, config)
        tf32_time = time.time() - start
        print(f"With TF32: {tf32_time:.2f}s")
    else:
        print(f"GPU compute capability {capability[0]}.{capability[1]} - TF32 not available")

5. Batch Size Optimization

Tune batch size for your GPU memory:

# GPU memory recommendations:
# 8 GB:  batchsize_atoms=512
# 16 GB: batchsize_atoms=1024 (default)
# 24 GB: batchsize_atoms=1536
# 40+ GB: batchsize_atoms=2048

if torch.cuda.is_available():
    total_mem = torch.cuda.get_device_properties(0).total_memory / 1e9

    if total_mem < 12:
        recommended = 512
    elif total_mem < 20:
        recommended = 1024
    elif total_mem < 32:
        recommended = 1536
    else:
        recommended = 2048

    print(f"GPU memory: {total_mem:.1f} GB")
    print(f"Recommended batchsize_atoms: {recommended}")

6. Convergence Tuning

Trade off accuracy for speed:

# Fast (screening)
fast_config = Auto3DOptions(
    k=1,
    convergence_threshold=0.02,  # Loose threshold
    patience=100,                # Quick dropout
    opt_steps=1000,              # Fewer max steps
    use_gpu=False,
)

# Default (balanced)
default_config = Auto3DOptions(
    k=1,
    convergence_threshold=0.01,  # Default
    patience=250,
    opt_steps=2000,
    use_gpu=False,
)

# Accurate (production)
accurate_config = Auto3DOptions(
    k=1,
    convergence_threshold=0.003,  # Tight threshold
    patience=500,
    opt_steps=5000,
    use_gpu=False,
)

# Compare
for name, config in [("Fast", fast_config), ("Default", default_config), ("Accurate", accurate_config)]:
    start = time.time()
    mols = smiles2mols(["CCCCc1ccccc1"], config)
    elapsed = time.time() - start
    energy = float(mols[0].GetProp("E_tot"))
    print(f"{name}: {elapsed:.2f}s, E={energy:.6f} Hartree")

7. Reducing Initial Conformers

Limit the number of initial conformers generated:

# For flexible molecules, limit initial conformers
config = Auto3DOptions(
    k=1,
    max_confs=50,  # Limit initial conformers (None = dynamic)
    use_gpu=False,
)

start = time.time()
mols = smiles2mols(["CCCCCCCCCC"], config)  # decane
print(f"Time: {time.time() - start:.2f}s")

8. Multi-GPU Processing

Distribute work across multiple GPUs:

if torch.cuda.is_available():
    n_gpus = torch.cuda.device_count()
    print(f"Available GPUs: {n_gpus}")

    if n_gpus > 1:
        # Use all GPUs
        config = Auto3DOptions(
            k=1,
            use_gpu=True,
            gpu_idx=list(range(n_gpus)),  # [0, 1, 2, ...]
        )
        print(f"Using GPUs: {list(range(n_gpus))}")
    else:
        print("Single GPU mode")

9. Memory Profiling

if torch.cuda.is_available():
    torch.cuda.reset_peak_memory_stats()

    config = Auto3DOptions(k=1, use_gpu=True)
    mols = smiles2mols(["c1ccc2c(c1)cccc2"], config)

    peak_memory = torch.cuda.max_memory_allocated() / 1e9
    print(f"Peak GPU memory: {peak_memory:.2f} GB")

Summary: Optimal Settings by Use Case

Use Case	Engine	GPU	TF32	batchsize	threshold
Quick screening	ANI2xt	Yes	Yes	1024	0.02
Standard workflow	AIMNET	Yes	Yes	1024	0.01
High accuracy	AIMNET + ensemble	Yes	No	512	0.003
Limited GPU memory	ANI2xt	Yes	Yes	256	0.01
CPU only	ANI2xt	No	N/A	1024	0.01

CLI Quick Reference for Performance

# Quick screening (fastest)
auto3d run input.smi --k=1 --engine=ANI2xt --gpu

# Standard workflow
auto3d run input.smi --k=1 --engine=AIMNET --gpu

# High accuracy (with ensemble)
AUTO3D_USE_ENSEMBLE=1 auto3d run input.smi --k=1 --engine=AIMNET --gpu

# With torch.compile optimization (ANI models)
AUTO3D_COMPILE_MODEL=1 auto3d run input.smi --k=1 --engine=ANI2x --gpu

# Multi-GPU processing
auto3d run large_dataset.smi --k=1 --gpu --gpu-idx="0,1,2,3"

# Generate preset configurations
auto3d config init -p quick -o quick.yaml      # Fast screening
auto3d config init -p balanced -o balanced.yaml # Default settings
auto3d config init -p thorough -o thorough.yaml # High accuracy

# Run with configuration file
auto3d run input.smi -c quick.yaml

Example performance.yaml for HPC:

optimizing_engine: ANI2xt
use_gpu: true
gpu_idx: [0, 1, 2, 3]
batchsize_atoms: 2048
allow_tf32: true
convergence_threshold: 0.01
patience: 200