{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No module named 'deepmd'\n"
     ]
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "from ase import units, Atoms\n",
    "from ase.build import molecule\n",
    "from ase.io import read, write\n",
    "from dask.distributed import Client\n",
    "from dask_jobqueue import SLURMCluster\n",
    "from prefect import flow\n",
    "from prefect_dask import DaskTaskRunner\n",
    "from pymatgen.core import Molecule\n",
    "from pymatgen.io.packmol import PackmolBoxGen\n",
    "\n",
    "from mlip_arena.models import REGISTRY, MLIPEnum\n",
    "from mlip_arena.tasks.md import run as MD"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Intial configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "h2 = molecule(\"H2\")\n",
    "o2 = molecule(\"O2\")\n",
    "h2o = molecule(\"H2O\")\n",
    "\n",
    "write(\"h2.xyz\", h2)\n",
    "write(\"o2.xyz\", o2)\n",
    "write(\"h2o.xyz\", h2o)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "h2 = Molecule.from_file(\"h2.xyz\")\n",
    "o2 = Molecule.from_file(\"o2.xyz\")\n",
    "h2o = Molecule.from_file(\"h2o.xyz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "molecules = []\n",
    "\n",
    "for m, number in zip([h2, o2], [128, 64]):\n",
    "    molecules.append(\n",
    "        {\n",
    "            \"name\": m.composition.to_pretty_string(),\n",
    "            \"number\": number,\n",
    "            \"coords\": m,\n",
    "        }\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Atoms(symbols='H256O128', pbc=True, cell=[30.0, 30.0, 30.0])\n"
     ]
    }
   ],
   "source": [
    "tolerance = 2.0\n",
    "input_gen = PackmolBoxGen(\n",
    "    tolerance=tolerance,\n",
    "    seed=1,\n",
    ")\n",
    "margin = 0.5 * tolerance\n",
    "\n",
    "a = 30\n",
    "\n",
    "packmol_set = input_gen.get_input_set(\n",
    "    molecules=molecules,\n",
    "    box=[margin, margin, margin, a - margin, a - margin, a - margin],\n",
    ")\n",
    "packmol_set.write_input(\".\")\n",
    "packmol_set.run(\".\")\n",
    "\n",
    "atoms = read(\"packmol_out.xyz\")\n",
    "atoms.cell = [a, a, a]\n",
    "atoms.pbc = True\n",
    "\n",
    "print(atoms)\n",
    "\n",
    "write(f\"{atoms.get_chemical_formula()}.extxyz\", atoms)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run workflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Atoms(symbols='H256O128', pbc=True, cell=[30.0, 30.0, 30.0])\n"
     ]
    }
   ],
   "source": [
    "atoms = read(\"H256O128.extxyz\")\n",
    "print(atoms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#!/bin/bash\n",
      "\n",
      "#SBATCH -A matgen\n",
      "#SBATCH --mem=0\n",
      "#SBATCH -t 02:00:00\n",
      "#SBATCH -J combustion-water\n",
      "#SBATCH -q regular\n",
      "#SBATCH -N 1\n",
      "#SBATCH -C gpu\n",
      "#SBATCH -G 4\n",
      "#SBATCH --exclusive\n",
      "source ~/.bashrc\n",
      "module load python\n",
      "source activate /pscratch/sd/c/cyrusyc/.conda/mlip-arena\n",
      "/pscratch/sd/c/cyrusyc/.conda/mlip-arena/bin/python -m distributed.cli.dask_worker tcp://128.55.64.15:38781 --name dummy-name --nthreads 1 --memory-limit 59.60GiB --nanny --death-timeout 86400\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/node.py:187: UserWarning: Port 8787 is already in use.\n",
      "Perhaps you already have a cluster running?\n",
      "Hosting the HTTP server on port 44831 instead\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "nodes_per_alloc = 1\n",
    "gpus_per_alloc = 4\n",
    "ntasks = 1\n",
    "\n",
    "cluster_kwargs = dict(\n",
    "    cores=1,\n",
    "    memory=\"64 GB\",\n",
    "    shebang=\"#!/bin/bash\",\n",
    "    account=\"matgen\",\n",
    "    walltime=\"02:00:00\",\n",
    "    job_mem=\"0\",\n",
    "    job_script_prologue=[\n",
    "        \"source ~/.bashrc\",\n",
    "        \"module load python\",\n",
    "        \"source activate /pscratch/sd/c/cyrusyc/.conda/mlip-arena\",\n",
    "    ],\n",
    "    job_directives_skip=[\"-n\", \"--cpus-per-task\", \"-J\"],\n",
    "    job_extra_directives=[\n",
    "        \"-J combustion-water\",\n",
    "        \"-q regular\",\n",
    "        f\"-N {nodes_per_alloc}\",\n",
    "        \"-C gpu\",\n",
    "        f\"-G {gpus_per_alloc}\",\n",
    "        \"--exclusive\",\n",
    "    ],\n",
    "    death_timeout=86400,\n",
    ")\n",
    "\n",
    "cluster = SLURMCluster(**cluster_kwargs)\n",
    "\n",
    "\n",
    "print(cluster.job_script())\n",
    "cluster.adapt(minimum_jobs=2, maximum_jobs=2)\n",
    "client = Client(cluster)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@flow(task_runner=DaskTaskRunner(address=client.scheduler.address), log_prints=True)\n",
    "def combustion(atoms: Atoms):\n",
    "    futures = []\n",
    "\n",
    "    for model in MLIPEnum:\n",
    "        future = MD.submit(\n",
    "            atoms=atoms,\n",
    "            calculator_name=model,\n",
    "            calculator_kwargs=None,\n",
    "            ensemble=\"nvt\",\n",
    "            dynamics=\"nose-hoover\",\n",
    "            time_step=None,\n",
    "            ase_md_kwargs=dict(ttime=25 * units.fs, pfactor=None),\n",
    "            total_time=1000_000,\n",
    "            temperature=[300, 3000, 3000, 300],\n",
    "            pressure=None,\n",
    "            md_velocity_seed=0,\n",
    "            traj_file=Path(REGISTRY[model.name][\"family\"])\n",
    "            / f\"{model.name}_{atoms.get_chemical_formula()}.traj\",\n",
    "            traj_interval=1000,\n",
    "            restart=True,\n",
    "        )\n",
    "\n",
    "        futures.append(future)\n",
    "\n",
    "    return [future.result() for future in futures]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "results = combustion(atoms)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mlip-arena",
   "language": "python",
   "name": "mlip-arena"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.8"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}