{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "# Performance Tuning Guide\n\nThis notebook covers optimization strategies for getting the best performance from Auto3D.\n\n## CLI Quick Reference (Recommended)\n\n```bash\n# Quick screening (fastest)\nauto3d run input.smi --k=1 --engine=ANI2xt --gpu\n\n# Standard workflow\nauto3d run input.smi --k=1 --engine=AIMNET --gpu\n\n# High accuracy (with ensemble)\nAUTO3D_USE_ENSEMBLE=1 auto3d run input.smi --k=1 --engine=AIMNET --gpu\n\n# With torch.compile optimization (ANI models)\nAUTO3D_COMPILE_MODEL=1 auto3d run input.smi --k=1 --engine=ANI2x --gpu\n\n# Multi-GPU processing\nauto3d run large_dataset.smi --k=1 --gpu --gpu-idx=\"0,1,2,3\"\n\n# Generate preset configurations\nauto3d config init -p quick -o quick.yaml      # Fast screening\nauto3d config init -p balanced -o balanced.yaml # Default settings  \nauto3d config init -p thorough -o thorough.yaml # High accuracy\n\n# Run with configuration file\nauto3d run input.smi -c quick.yaml\n```\n\n**Optimal Settings by Use Case:**\n\n| Use Case | CLI Command |\n|----------|-------------|\n| Quick screening | `auto3d run input.smi --k=1 --engine=ANI2xt --gpu` |\n| Standard workflow | `auto3d run input.smi --k=1 --engine=AIMNET --gpu` |\n| High accuracy | `AUTO3D_USE_ENSEMBLE=1 auto3d run input.smi --k=1 --gpu` |\n| Multi-GPU | `auto3d run input.smi --k=1 --gpu --gpu-idx=\"0,1,2,3\"` |\n\n---\n\n## Python API Performance Guide\n\nBelow we explore performance optimization using the Python API:"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import torch\n",
    "import Auto3D\n",
    "from Auto3D import Auto3DOptions, smiles2mols, create_model\n",
    "\n",
    "print(f\"Auto3D: {Auto3D.__version__}\")\n",
    "print(f\"PyTorch: {torch.__version__}\")\n",
    "print(f\"CUDA available: {torch.cuda.is_available()}\")\n",
    "if torch.cuda.is_available():\n",
    "    print(f\"GPU: {torch.cuda.get_device_name(0)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Choosing the Right Engine\n",
    "\n",
    "Speed comparison of available engines:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test molecules\n",
    "test_smiles = [\"c1ccc2c(c1)cccc2\", \"CCCCc1ccccc1\", \"O=C(NC)Nc1ccccc1\"]\n",
    "\n",
    "engines = [\"ANI2xt\", \"ANI2x\", \"AIMNET\"]\n",
    "\n",
    "for engine in engines:\n",
    "    start = time.time()\n",
    "    config = Auto3DOptions(k=1, optimizing_engine=engine, use_gpu=False)\n",
    "    mols = smiles2mols(test_smiles, config)\n",
    "    elapsed = time.time() - start\n",
    "    print(f\"{engine}: {elapsed:.2f}s\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Recommendations:**\n",
    "- **ANI2xt**: Fastest, use for initial screening\n",
    "- **ANI2x**: Good balance of speed and accuracy\n",
    "- **AIMNET**: Most accurate, supports charged molecules"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. GPU Acceleration\n",
    "\n",
    "GPU provides significant speedup for large batches:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if torch.cuda.is_available():\n",
    "    # CPU timing\n",
    "    start = time.time()\n",
    "    config = Auto3DOptions(k=1, use_gpu=False)\n",
    "    mols_cpu = smiles2mols(test_smiles, config)\n",
    "    cpu_time = time.time() - start\n",
    "    \n",
    "    # GPU timing\n",
    "    start = time.time()\n",
    "    config = Auto3DOptions(k=1, use_gpu=True)\n",
    "    mols_gpu = smiles2mols(test_smiles, config)\n",
    "    gpu_time = time.time() - start\n",
    "    \n",
    "    print(f\"CPU: {cpu_time:.2f}s\")\n",
    "    print(f\"GPU: {gpu_time:.2f}s\")\n",
    "    print(f\"Speedup: {cpu_time/gpu_time:.1f}x\")\n",
    "else:\n",
    "    print(\"No GPU available for comparison\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Single Model vs Ensemble (AIMNET)\n",
    "\n",
    "By default, Auto3D uses a single AIMNet2 model for ~35x faster optimization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Single model (default, faster)\n",
    "os.environ[\"AUTO3D_USE_ENSEMBLE\"] = \"0\"\n",
    "start = time.time()\n",
    "config = Auto3DOptions(k=1, optimizing_engine=\"AIMNET\", use_gpu=False)\n",
    "mols = smiles2mols(test_smiles[:1], config)\n",
    "single_time = time.time() - start\n",
    "\n",
    "print(f\"Single model: {single_time:.2f}s\")\n",
    "print(\"Use ensemble=True for highest accuracy (slower)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. TF32 Acceleration (Ampere+ GPUs)\n",
    "\n",
    "Enable TensorFloat-32 for faster matrix operations on RTX 30xx, A100, H100:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if torch.cuda.is_available():\n",
    "    # Check GPU architecture\n",
    "    capability = torch.cuda.get_device_capability(0)\n",
    "    is_ampere_plus = capability[0] >= 8\n",
    "    \n",
    "    if is_ampere_plus:\n",
    "        print(f\"GPU supports TF32 (compute capability {capability[0]}.{capability[1]})\")\n",
    "        \n",
    "        # With TF32\n",
    "        config = Auto3DOptions(k=1, use_gpu=True, allow_tf32=True)\n",
    "        start = time.time()\n",
    "        mols = smiles2mols(test_smiles, config)\n",
    "        tf32_time = time.time() - start\n",
    "        print(f\"With TF32: {tf32_time:.2f}s\")\n",
    "    else:\n",
    "        print(f\"GPU compute capability {capability[0]}.{capability[1]} - TF32 not available\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Batch Size Optimization\n",
    "\n",
    "Tune batch size for your GPU memory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# GPU memory recommendations:\n",
    "# 8 GB:  batchsize_atoms=512\n",
    "# 16 GB: batchsize_atoms=1024 (default)\n",
    "# 24 GB: batchsize_atoms=1536\n",
    "# 40+ GB: batchsize_atoms=2048\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    total_mem = torch.cuda.get_device_properties(0).total_memory / 1e9\n",
    "    \n",
    "    if total_mem < 12:\n",
    "        recommended = 512\n",
    "    elif total_mem < 20:\n",
    "        recommended = 1024\n",
    "    elif total_mem < 32:\n",
    "        recommended = 1536\n",
    "    else:\n",
    "        recommended = 2048\n",
    "    \n",
    "    print(f\"GPU memory: {total_mem:.1f} GB\")\n",
    "    print(f\"Recommended batchsize_atoms: {recommended}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Convergence Tuning\n",
    "\n",
    "Trade off accuracy for speed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fast (screening)\n",
    "fast_config = Auto3DOptions(\n",
    "    k=1,\n",
    "    convergence_threshold=0.02,  # Loose threshold\n",
    "    patience=100,                # Quick dropout\n",
    "    opt_steps=1000,              # Fewer max steps\n",
    "    use_gpu=False,\n",
    ")\n",
    "\n",
    "# Default (balanced)\n",
    "default_config = Auto3DOptions(\n",
    "    k=1,\n",
    "    convergence_threshold=0.01,  # Default\n",
    "    patience=250,\n",
    "    opt_steps=2000,\n",
    "    use_gpu=False,\n",
    ")\n",
    "\n",
    "# Accurate (production)\n",
    "accurate_config = Auto3DOptions(\n",
    "    k=1,\n",
    "    convergence_threshold=0.003,  # Tight threshold\n",
    "    patience=500,\n",
    "    opt_steps=5000,\n",
    "    use_gpu=False,\n",
    ")\n",
    "\n",
    "# Compare\n",
    "for name, config in [(\"Fast\", fast_config), (\"Default\", default_config), (\"Accurate\", accurate_config)]:\n",
    "    start = time.time()\n",
    "    mols = smiles2mols([\"CCCCc1ccccc1\"], config)\n",
    "    elapsed = time.time() - start\n",
    "    energy = float(mols[0].GetProp(\"E_tot\"))\n",
    "    print(f\"{name}: {elapsed:.2f}s, E={energy:.6f} Hartree\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Reducing Initial Conformers\n",
    "\n",
    "Limit the number of initial conformers generated:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For flexible molecules, limit initial conformers\n",
    "config = Auto3DOptions(\n",
    "    k=1,\n",
    "    max_confs=50,  # Limit initial conformers (None = dynamic)\n",
    "    use_gpu=False,\n",
    ")\n",
    "\n",
    "start = time.time()\n",
    "mols = smiles2mols([\"CCCCCCCCCC\"], config)  # decane\n",
    "print(f\"Time: {time.time() - start:.2f}s\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Multi-GPU Processing\n",
    "\n",
    "Distribute work across multiple GPUs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if torch.cuda.is_available():\n",
    "    n_gpus = torch.cuda.device_count()\n",
    "    print(f\"Available GPUs: {n_gpus}\")\n",
    "    \n",
    "    if n_gpus > 1:\n",
    "        # Use all GPUs\n",
    "        config = Auto3DOptions(\n",
    "            k=1,\n",
    "            use_gpu=True,\n",
    "            gpu_idx=list(range(n_gpus)),  # [0, 1, 2, ...]\n",
    "        )\n",
    "        print(f\"Using GPUs: {list(range(n_gpus))}\")\n",
    "    else:\n",
    "        print(\"Single GPU mode\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Memory Profiling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if torch.cuda.is_available():\n",
    "    torch.cuda.reset_peak_memory_stats()\n",
    "    \n",
    "    config = Auto3DOptions(k=1, use_gpu=True)\n",
    "    mols = smiles2mols([\"c1ccc2c(c1)cccc2\"], config)\n",
    "    \n",
    "    peak_memory = torch.cuda.max_memory_allocated() / 1e9\n",
    "    print(f\"Peak GPU memory: {peak_memory:.2f} GB\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## Summary: Optimal Settings by Use Case\n\n| Use Case | Engine | GPU | TF32 | batchsize | threshold |\n|----------|--------|-----|------|-----------|------------|\n| Quick screening | ANI2xt | Yes | Yes | 1024 | 0.02 |\n| Standard workflow | AIMNET | Yes | Yes | 1024 | 0.01 |\n| High accuracy | AIMNET + ensemble | Yes | No | 512 | 0.003 |\n| Limited GPU memory | ANI2xt | Yes | Yes | 256 | 0.01 |\n| CPU only | ANI2xt | No | N/A | 1024 | 0.01 |\n\n## CLI Quick Reference for Performance\n\n```bash\n# Quick screening (fastest)\nauto3d run input.smi --k=1 --engine=ANI2xt --gpu\n\n# Standard workflow\nauto3d run input.smi --k=1 --engine=AIMNET --gpu\n\n# High accuracy (with ensemble)\nAUTO3D_USE_ENSEMBLE=1 auto3d run input.smi --k=1 --engine=AIMNET --gpu\n\n# With torch.compile optimization (ANI models)\nAUTO3D_COMPILE_MODEL=1 auto3d run input.smi --k=1 --engine=ANI2x --gpu\n\n# Multi-GPU processing\nauto3d run large_dataset.smi --k=1 --gpu --gpu-idx=\"0,1,2,3\"\n\n# Generate preset configurations\nauto3d config init -p quick -o quick.yaml      # Fast screening\nauto3d config init -p balanced -o balanced.yaml # Default settings  \nauto3d config init -p thorough -o thorough.yaml # High accuracy\n\n# Run with configuration file\nauto3d run input.smi -c quick.yaml\n```\n\nExample `performance.yaml` for HPC:\n\n```yaml\noptimizing_engine: ANI2xt\nuse_gpu: true\ngpu_idx: [0, 1, 2, 3]\nbatchsize_atoms: 2048\nallow_tf32: true\nconvergence_threshold: 0.01\npatience: 200\n```"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}