Large-Scale Conformer Generation
This notebook demonstrates best practices for processing large molecular datasets (1000+ molecules) with Auto3D.
[ ]:
import os
import time
import torch
import Auto3D
from Auto3D import Auto3DOptions, main
from pathlib import Path
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)}")
print(f"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
1. Memory and Chunking Configuration
For large datasets, Auto3D automatically splits processing into chunks based on available memory.
[ ]:
# Get system memory info
import psutil
total_ram = psutil.virtual_memory().total / 1e9
available_ram = psutil.virtual_memory().available / 1e9
print(f"Total RAM: {total_ram:.1f} GB")
print(f"Available RAM: {available_ram:.1f} GB")
# Recommended memory allocation
recommended_memory = int(available_ram * 0.8) # Use 80% of available
print(f"\nRecommended 'memory' setting: {recommended_memory} GB")
2. Optimal Configuration for Large Datasets
[ ]:
def create_large_scale_config(input_path, output_k=1, use_gpu=True):
"""Create optimized configuration for large-scale processing."""
# Determine optimal settings based on hardware
if use_gpu and torch.cuda.is_available():
gpu_mem = torch.cuda.get_device_properties(0).total_memory / 1e9
n_gpus = torch.cuda.device_count()
# Batch size based on GPU memory
if gpu_mem < 12:
batchsize = 512
elif gpu_mem < 24:
batchsize = 1024
else:
batchsize = 2048
gpu_idx = list(range(n_gpus)) if n_gpus > 1 else 0
else:
batchsize = 1024
gpu_idx = 0
# Calculate available memory
available_ram = psutil.virtual_memory().available / 1e9
memory = int(available_ram * 0.8)
config = Auto3DOptions(
path=input_path,
k=output_k,
# Performance settings
use_gpu=use_gpu,
gpu_idx=gpu_idx,
optimizing_engine="ANI2xt", # Fastest engine
allow_tf32=True, # Enable on Ampere+ GPUs
# Memory management
memory=memory,
capacity=50, # Molecules per GB
batchsize_atoms=batchsize,
# Optimization settings (balanced)
opt_steps=2000,
convergence_threshold=0.01,
patience=200,
# Limit initial conformers for speed
max_confs=50,
)
return config
# Example usage
# config = create_large_scale_config("large_dataset.smi", output_k=1)
3. Monitoring Progress
For long-running jobs, monitor progress using the verbose output.
[ ]:
# Create a sample large dataset for demonstration
sample_smiles = [
"c1ccccc1",
"CCO",
"CC(=O)O",
"c1ccc2c(c1)cccc2",
"CCCCCC",
] * 10 # 50 molecules for demo
# Write to temp file
demo_file = "demo_large_dataset.smi"
with open(demo_file, "w") as f:
for i, smi in enumerate(sample_smiles):
f.write(f"{smi} mol_{i}\n")
print(f"Created {demo_file} with {len(sample_smiles)} molecules")
[ ]:
# Process with timing
if __name__ == "__main__":
start_time = time.time()
config = Auto3DOptions(
path=demo_file,
k=1,
use_gpu=torch.cuda.is_available(),
optimizing_engine="ANI2xt",
max_confs=20,
patience=100,
)
output = main(config)
elapsed = time.time() - start_time
throughput = len(sample_smiles) / elapsed
print(f"\n" + "="*50)
print(f"Total time: {elapsed:.1f} seconds")
print(f"Throughput: {throughput:.1f} molecules/second")
print(f"Output: {output}")
4. Multi-GPU Processing
For systems with multiple GPUs, distribute work automatically.
[ ]:
if torch.cuda.is_available():
n_gpus = torch.cuda.device_count()
print(f"Available GPUs: {n_gpus}")
for i in range(n_gpus):
props = torch.cuda.get_device_properties(i)
print(f" GPU {i}: {props.name} ({props.total_memory / 1e9:.1f} GB)")
if n_gpus > 1:
print(f"\nUsing all {n_gpus} GPUs: gpu_idx={list(range(n_gpus))}")
else:
print("No GPU available")
5. Processing in Batches
For very large datasets, process in separate batches with checkpointing.
[ ]:
def split_file(input_path, batch_size=10000):
"""Split large SMILES file into batches."""
with open(input_path) as f:
lines = f.readlines()
batches = []
for i in range(0, len(lines), batch_size):
batch_path = f"{input_path}.batch_{i//batch_size}.smi"
with open(batch_path, "w") as f:
f.writelines(lines[i:i+batch_size])
batches.append(batch_path)
return batches
def process_with_checkpoints(input_path, output_dir, batch_size=10000):
"""Process large dataset with checkpointing."""
from rdkit import Chem
output_dir = Path(output_dir)
output_dir.mkdir(exist_ok=True)
batches = split_file(input_path, batch_size)
results = []
for i, batch_path in enumerate(batches):
checkpoint = output_dir / f"batch_{i}_done.flag"
# Skip if already processed
if checkpoint.exists():
print(f"Batch {i} already processed, skipping")
continue
print(f"Processing batch {i+1}/{len(batches)}...")
config = Auto3DOptions(
path=batch_path,
k=1,
job_name=str(output_dir / f"batch_{i}"),
)
output = main(config)
results.append(output)
# Mark as done
checkpoint.touch()
# Clean up batch file
os.remove(batch_path)
return results
# Example:
# results = process_with_checkpoints("huge_dataset.smi", "output/", batch_size=5000)
6. Merging Results
[ ]:
def merge_sdf_files(sdf_files, output_path):
"""Merge multiple SDF files into one."""
from rdkit import Chem
writer = Chem.SDWriter(output_path)
total = 0
for sdf_file in sdf_files:
for mol in Chem.SDMolSupplier(sdf_file):
if mol is not None:
writer.write(mol)
total += 1
writer.close()
print(f"Merged {total} molecules into {output_path}")
return output_path
# Example:
# merge_sdf_files(["batch_0.sdf", "batch_1.sdf"], "merged_output.sdf")
7. Resource Monitoring
[ ]:
def print_resource_usage():
"""Print current resource usage."""
import psutil
# CPU
cpu_percent = psutil.cpu_percent(interval=1)
# Memory
mem = psutil.virtual_memory()
mem_used = mem.used / 1e9
mem_total = mem.total / 1e9
print(f"CPU: {cpu_percent:.1f}%")
print(f"RAM: {mem_used:.1f} / {mem_total:.1f} GB ({mem.percent:.1f}%)")
# GPU
if torch.cuda.is_available():
for i in range(torch.cuda.device_count()):
mem_used = torch.cuda.memory_allocated(i) / 1e9
mem_total = torch.cuda.get_device_properties(i).total_memory / 1e9
print(f"GPU {i}: {mem_used:.2f} / {mem_total:.1f} GB")
print_resource_usage()
8. Cleanup
[ ]:
# Clean up demo file
if os.path.exists(demo_file):
os.remove(demo_file)
print(f"Cleaned up {demo_file}")
# Clear GPU cache if needed
if torch.cuda.is_available():
torch.cuda.empty_cache()
print("GPU cache cleared")
Summary: Best Practices for Large-Scale Processing
Use the right engine: ANI2xt for speed, AIMNET for accuracy
Enable GPU: Significant speedup, especially for large batches
Use multi-GPU: Linear scaling with number of GPUs
Tune batch size: Match to GPU memory
Limit initial conformers:
max_confs=50for speedProcess in batches: Checkpoint for resumability
Monitor resources: Watch memory usage
Clean up: Remove temporary files, clear GPU cache
CLI Commands for Large-Scale Processing
# Basic large-scale processing
auto3d run large_dataset.smi --k=1 --engine=ANI2xt --gpu
# Multi-GPU processing
auto3d run large_dataset.smi --k=1 --gpu --gpu-idx="0,1,2,3"
# Validate input before processing
auto3d validate large_dataset.smi
# Use configuration file for reproducibility
auto3d config init -p quick -o large_scale.yaml
auto3d run large_dataset.smi -c large_scale.yaml
Example large_scale.yaml:
k: 1
optimizing_engine: ANI2xt
use_gpu: true
gpu_idx: [0, 1, 2, 3]
memory: 64
capacity: 50
batchsize_atoms: 2048
max_confs: 50
patience: 200
allow_tf32: true
Batch Processing with Shell Scripts
#!/bin/bash
# process_batches.sh
# Split large file
split -l 10000 huge_dataset.smi batch_
# Process each batch
for batch in batch_*; do
echo "Processing $batch..."
auto3d run "$batch" --k=1 --engine=ANI2xt --gpu
done
# Merge results (using Python/RDKit)
echo "Done. Merge output SDF files as needed."
HPC SLURM Script
#!/bin/bash
#SBATCH --job-name=auto3d_large
#SBATCH --gpus=4
#SBATCH --cpus-per-task=16
#SBATCH --mem=128G
#SBATCH --time=24:00:00
module load cuda
conda activate auto3d
auto3d run large_dataset.smi --k=1 --gpu --gpu-idx="0,1,2,3" --engine=ANI2xt