Advanced Usage

This guide covers advanced features of Auto3D: multi-GPU workflows, performance optimization, custom neural network potentials, and fine-tuning parameters.

Multi-GPU Processing

Auto3D supports multi-GPU processing for large datasets, automatically distributing molecules across GPUs.

CLI Usage

# Use multiple GPUs (recommended for large datasets)
auto3d run large_dataset.smi --k=1 --gpu --gpu-idx="0,1,2,3"

# Use specific GPUs on shared systems
auto3d run large_dataset.smi --k=1 --gpu --gpu-idx="2,3"

# Use CUDA_VISIBLE_DEVICES environment variable
CUDA_VISIBLE_DEVICES=0,1,2,3 auto3d run large_dataset.smi --k=1 --gpu

Python API

from Auto3D import Auto3DOptions, main

if __name__ == "__main__":
    config = Auto3DOptions(
        path="large_dataset.smi",
        k=1,
        use_gpu=True,
        gpu_idx=[0, 1, 2, 3],  # Use GPUs 0, 1, 2, and 3
    )
    output = main(config)

Performance Optimization

Quick Settings Reference

# Fast screening (fastest, good enough for ranking)
auto3d run input.smi --k=1 --engine=ANI2xt --gpu

# Balanced (default)
auto3d run input.smi --k=1 --engine=AIMNET --gpu

# High accuracy (for final production runs)
AUTO3D_USE_ENSEMBLE=1 auto3d run input.smi --k=1 --engine=AIMNET --gpu

Configuration Presets

Auto3D provides presets for common use cases:

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

# Use a preset
auto3d run input.smi -c quick.yaml

TF32 Acceleration

Enable TensorFloat-32 for ~1.5x faster computation on Ampere+ GPUs (RTX 30xx, A100, H100):

# Create config with TF32 enabled
cat > performance.yaml << EOF
allow_tf32: true
use_gpu: true
EOF

auto3d run input.smi --k=1 -c performance.yaml

Python API:

config = Auto3DOptions(
    path="input.smi",
    k=1,
    use_gpu=True,
    allow_tf32=True,  # ~1.5x faster matrix operations
)

Note

TF32 reduces precision slightly (19 mantissa bits vs 23 for FP32). This is acceptable for conformer generation but may affect very tight energy comparisons.

torch.compile() Optimization

Enable PyTorch 2.0 compilation for ANI models (~1.25x speedup):

# Enable via environment variable
AUTO3D_COMPILE_MODEL=1 auto3d run input.smi --k=1 --engine=ANI2x --gpu

# Combine multiple optimizations
AUTO3D_COMPILE_MODEL=1 auto3d run input.smi --k=1 --engine=ANI2xt --gpu

Python API:

import os
os.environ["AUTO3D_COMPILE_MODEL"] = "1"

# Or via create_model for custom workflows
from Auto3D import create_model
model = create_model("ANI2xt", device=device, compile_model=True)

Note

torch.compile() works best with ANI2x/ANI2xt. AIMNET uses optimized JIT compilation internally.

Batch Size Tuning

Adjust batch size based on GPU memory:

# Create config with tuned batch size
cat > gpu_tuning.yaml << EOF
batchsize_atoms: 2048   # For 24GB+ GPUs
# batchsize_atoms: 512  # For 8GB GPUs
use_gpu: true
EOF

auto3d run input.smi --k=1 -c gpu_tuning.yaml

Recommended batch sizes:

GPU	Memory	batchsize_atoms
RTX 3070, RTX 4070	8 GB	512
RTX 3080, V100	16 GB	1024 (default)
RTX 3090, A5000, RTX 4090	24 GB	1536
A100, H100	40-80 GB	2048

Python API:

config = Auto3DOptions(
    path="input.smi",
    k=1,
    batchsize_atoms=2048,  # Adjust for your GPU
)

Large Dataset Memory Management

For very large datasets (100k+ molecules), configure memory allocation:

# Create config for large-scale processing
cat > large_scale.yaml << EOF
memory: 64       # Assign 64GB RAM
capacity: 50     # Molecules per GB
use_gpu: true
gpu_idx: [0, 1, 2, 3]
EOF

auto3d run huge_dataset.smi --k=1 -c large_scale.yaml

Python API:

config = Auto3DOptions(
    path="huge_dataset.smi",
    k=1,
    memory=64,        # Assign 64GB RAM
    capacity=50,      # Molecules per GB
    use_gpu=True,
    gpu_idx=[0, 1, 2, 3],
)

Environment Variables

Control Auto3D behavior via environment variables:

# Enable torch.compile for ANI models
export AUTO3D_COMPILE_MODEL=1

# Use AIMNET ensemble (slower, highest accuracy)
export AUTO3D_USE_ENSEMBLE=1

# Set OpenEye license path
export OE_LICENSE=/path/to/oe_license.txt

# Run with environment settings
auto3d run input.smi --k=1 --gpu

Variable	Default	Description
`AUTO3D_COMPILE_MODEL`	`0`	Enable torch.compile() for ANI models
`AUTO3D_USE_ENSEMBLE`	`0`	Use AIMNET 8-model ensemble
`OE_LICENSE`	(none)	OpenEye license for Omega isomer engine

Optimization Parameters

Fine-tune geometry optimization via CLI configuration files or Python API.

Quick vs Accurate Settings

# Quick screening (fast, slightly looser convergence)
cat > quick.yaml << EOF
optimizing_engine: ANI2xt
convergence_threshold: 0.02
patience: 100
opt_steps: 1000
max_confs: 50
EOF

# Accurate (tight convergence for production)
cat > accurate.yaml << EOF
optimizing_engine: AIMNET
convergence_threshold: 0.003
patience: 500
opt_steps: 5000
EOF

# Run with settings
auto3d run input.smi --k=1 -c quick.yaml --gpu
auto3d run input.smi --k=1 -c accurate.yaml --gpu

Python API:

# Quick screening
config = Auto3DOptions(
    path="input.smi",
    k=1,
    optimizing_engine="ANI2xt",
    convergence_threshold=0.02,
    patience=100,
    opt_steps=1000,
    max_confs=50,
)

# Accurate
config = Auto3DOptions(
    path="input.smi",
    k=1,
    optimizing_engine="AIMNET",
    convergence_threshold=0.003,
    patience=500,
    opt_steps=5000,
)

Parameter Reference

Parameter	Default	Description
`opt_steps`	2000	Maximum optimization steps per structure
`convergence_threshold`	0.01	Force convergence threshold in eV/A
`patience`	250	Steps before dropping oscillating structures
`batchsize_atoms`	1024	Atoms per optimization batch per GB memory
`max_confs`	(auto)	Maximum initial conformers per molecule

Custom Neural Network Potentials

Auto3D supports custom PyTorch NNP models for specialized applications.

CLI Usage

# Use custom model by path
auto3d run input.smi --k=1 --engine=/path/to/my_model.pt --gpu

# With configuration file
cat > custom.yaml << EOF
optimizing_engine: /path/to/my_model.pt
use_gpu: true
EOF

auto3d run input.smi --k=1 -c custom.yaml

Python API

from Auto3D import Auto3DOptions, main

if __name__ == "__main__":
    config = Auto3DOptions(
        path="input.smi",
        k=1,
        optimizing_engine="/path/to/my_model.pt",
        use_gpu=True,
    )
    output = main(config)

NNPModel Protocol

Your custom model must implement this interface:

import torch

class MyNNP(torch.nn.Module):
    # Required attributes
    coord_pad = 0       # Padding value for coordinates
    species_pad = -1    # Padding value for species

    def forward(
        self,
        species: torch.Tensor,    # Shape: (batch_size, max_atoms)
        coords: torch.Tensor,     # Shape: (batch_size, max_atoms, 3)
        charges: torch.Tensor,    # Shape: (batch_size,)
    ) -> torch.Tensor:
        """
        Calculate energies for a batch of molecules.

        Args:
            species: Atomic numbers, padded with species_pad
            coords: Atomic coordinates in Angstroms
            charges: Total molecular charges

        Returns:
            Energies tensor of shape (batch_size,) in eV
        """
        energies = self.calculate_energies(species, coords, charges)
        return energies

Example Custom Model Wrapper

import torch
import torch.nn as nn

class CustomNNPWrapper(nn.Module):
    """Wrapper to make an external NNP compatible with Auto3D."""

    coord_pad = 0
    species_pad = -1

    def __init__(self, underlying_model):
        super().__init__()
        self.model = underlying_model

    def forward(self, species, coords, charges):
        batch_size = species.shape[0]
        energies = []

        for i in range(batch_size):
            mask = species[i] != self.species_pad
            valid_species = species[i][mask]
            valid_coords = coords[i][mask]
            charge = charges[i]

            energy = self.model.predict(valid_species, valid_coords, charge)
            energies.append(energy)

        return torch.stack(energies)

# Save for use with Auto3D
model = CustomNNPWrapper(your_model)
torch.save(model, "my_model.pt")

Tautomer Enumeration

Enable tautomer enumeration for drug-like molecules with multiple possible forms.

CLI Usage

# Enable tautomer enumeration
auto3d run input.smi --k=1 --enumerate-tautomer --gpu

# With ANI2xt (recommended for tautomers)
auto3d run input.smi --k=1 --enumerate-tautomer --engine=ANI2xt --gpu

# Advanced configuration
cat > tautomer.yaml << EOF
enumerate_tautomer: true
tauto_engine: rdkit
optimizing_engine: ANI2xt
max_confs: 10
patience: 200
EOF

auto3d run input.smi --k=1 -c tautomer.yaml --gpu

Python API

from Auto3D import Auto3DOptions
from Auto3D.tautomer import get_stable_tautomers

config = Auto3DOptions(
    path="input.smi",
    k=1,
    enumerate_tautomer=True,
    tauto_engine="rdkit",      # or "oechem" with license
    optimizing_engine="ANI2xt", # Recommended for tautomers
    max_confs=10,
    patience=200,
)

# Get stable tautomers with top-3 per input
output = get_stable_tautomers(config, tauto_k=3)

Available Models

Model	Supported Elements	Charges	Speed
`AIMNET`	H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I	Neutral + charged	Fast (default)
`ANI2x`	H, C, N, O, F, S, Cl	Neutral only	Very fast
`ANI2xt`	H, C, N, O, F, S, Cl	Neutral only	Ultra-fast

Single Model vs Ensemble (AIMNET)

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

# Default: single model (fast)
auto3d run input.smi --k=1 --gpu

# Ensemble: highest accuracy (slower)
AUTO3D_USE_ENSEMBLE=1 auto3d run input.smi --k=1 --gpu

The single model is accurate enough for conformer generation and ranking. Use ensemble only when you need the most accurate absolute energies.

Model Factory API (Python)

For custom workflows, create models directly:

import torch
from Auto3D import create_model
from Auto3D.model_factory import ModelFactory

device = torch.device("cuda:0")

# Create model
model = create_model("AIMNET", device=device)

# Use for calculations
energies = model(species, coords, charges)

# List available models
print(ModelFactory.available_models())

# Clear cache to free GPU memory
ModelFactory.clear_cache()

Troubleshooting

GPU Memory Issues

If you encounter CUDA out-of-memory errors:

# 1. Reduce batch size
cat > low_memory.yaml << EOF
batchsize_atoms: 512
EOF
auto3d run input.smi --k=1 -c low_memory.yaml --gpu

# 2. Disable ensemble
AUTO3D_USE_ENSEMBLE=0 auto3d run input.smi --k=1 --gpu

# 3. Use CPU as fallback
auto3d run input.smi --k=1 --no-gpu

Python solutions:

# Reduce batch size
config = Auto3DOptions(path="input.smi", k=1, batchsize_atoms=512)

# Clear model cache
from Auto3D.model_factory import ModelFactory
ModelFactory.clear_cache()

Slow Processing

If optimization is slower than expected:

# 1. Use fastest model
auto3d run input.smi --k=1 --engine=ANI2xt --gpu

# 2. Use quick preset
auto3d config init -p quick -o quick.yaml
auto3d run input.smi --k=1 -c quick.yaml --gpu

# 3. Enable TF32 on Ampere+ GPUs
cat > fast.yaml << EOF
optimizing_engine: ANI2xt
allow_tf32: true
convergence_threshold: 0.02
patience: 150
EOF
auto3d run input.smi --k=1 -c fast.yaml --gpu