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