{ "cells": [ { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [], "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.utils import REGISTRY, MLIPEnum\n", "from mlip_arena.tasks import MD" ] }, { "cell_type": "code", "execution_count": null, "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": null, "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": null, "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": [], "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": "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": 3, "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", "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.32:35601 --name dummy-name --nthreads 1 --memory-limit 59.60GiB --nanny --death-timeout 86400\n", "\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", " # f\"--exclusive\",\n", " # \"--time-min=00:30:00\",\n", " # \"--comment=1-00:00:00\",\n", " # \"--signal=B:USR1@60\",\n", " # \"--requeue\",\n", " # \"--open-mode=append\"\n", " ],\n", " death_timeout=86400\n", ")\n", "\n", "\n", "cluster = SLURMCluster(**cluster_kwargs)\n", "print(cluster.job_script())\n", "cluster.adapt(minimum_jobs=5, maximum_jobs=50)\n", "client = Client(cluster)" ] }, { "cell_type": "code", "execution_count": 4, "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", " mb_velocity_seed=0,\n", " traj_file=Path(REGISTRY[model.name][\"family\"])\n", " / f\"{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": [ { "data": { "text/html": [ "
10:53:24.417 | INFO    | prefect.engine - Created flow run 'masterful-panther' for flow 'combustion'\n",
       "
\n" ], "text/plain": [ "10:53:24.417 | \u001b[36mINFO\u001b[0m | prefect.engine - Created flow run\u001b[35m 'masterful-panther'\u001b[0m for flow\u001b[1;35m 'combustion'\u001b[0m\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
10:53:24.422 | INFO    | Flow run 'masterful-panther' - View at https://app.prefect.cloud/account/f7d40474-9362-4bfa-8950-ee6a43ec00f3/workspace/d4bb0913-5f5e-49f7-bfc5-06509088baeb/flow-runs/flow-run/4d0a29f8-4313-4dd3-b311-ea64abf82a83\n",
       "
\n" ], "text/plain": [ "10:53:24.422 | \u001b[36mINFO\u001b[0m | Flow run\u001b[35m 'masterful-panther'\u001b[0m - View at \u001b[94mhttps://app.prefect.cloud/account/f7d40474-9362-4bfa-8950-ee6a43ec00f3/workspace/d4bb0913-5f5e-49f7-bfc5-06509088baeb/flow-runs/flow-run/4d0a29f8-4313-4dd3-b311-ea64abf82a83\u001b[0m\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
10:53:24.423 | INFO    | prefect.task_runner.dask - Connecting to an existing Dask cluster at tcp://128.55.64.32:35601\n",
       "
\n" ], "text/plain": [ "10:53:24.423 | \u001b[36mINFO\u001b[0m | prefect.task_runner.dask - Connecting to an existing Dask cluster at tcp://128.55.64.32:35601\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
10:53:24.430 | INFO    | prefect.task_runner.dask - The Dask dashboard is available at http://128.55.64.32:8787/status\n",
       "
\n" ], "text/plain": [ "10:53:24.430 | \u001b[36mINFO\u001b[0m | prefect.task_runner.dask - The Dask dashboard is available at \u001b[94mhttp://128.55.64.32:8787/status\u001b[0m\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
10:53:24.927 | INFO    | Flow run 'masterful-panther' - Created task run 'md-4' for task 'md'\n",
       "
\n" ], "text/plain": [ "10:53:24.927 | \u001b[36mINFO\u001b[0m | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Created task run 'md-4' for task 'md'\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
10:53:24.975 | INFO    | Flow run 'masterful-panther' - Created task run 'md-6' for task 'md'\n",
       "
\n" ], "text/plain": [ "10:53:24.975 | \u001b[36mINFO\u001b[0m | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Created task run 'md-6' for task 'md'\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
10:53:24.981 | INFO    | Flow run 'masterful-panther' - Created task run 'md-0' for task 'md'\n",
       "
\n" ], "text/plain": [ "10:53:24.981 | \u001b[36mINFO\u001b[0m | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Created task run 'md-0' for task 'md'\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
10:53:24.989 | INFO    | Flow run 'masterful-panther' - Created task run 'md-1' for task 'md'\n",
       "
\n" ], "text/plain": [ "10:53:24.989 | \u001b[36mINFO\u001b[0m | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Created task run 'md-1' for task 'md'\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
10:53:24.995 | INFO    | Flow run 'masterful-panther' - Created task run 'md-5' for task 'md'\n",
       "
\n" ], "text/plain": [ "10:53:24.995 | \u001b[36mINFO\u001b[0m | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Created task run 'md-5' for task 'md'\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
10:53:24.999 | INFO    | Flow run 'masterful-panther' - Created task run 'md-7' for task 'md'\n",
       "
\n" ], "text/plain": [ "10:53:24.999 | \u001b[36mINFO\u001b[0m | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Created task run 'md-7' for task 'md'\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
10:53:25.002 | INFO    | Flow run 'masterful-panther' - Created task run 'md-8' for task 'md'\n",
       "
\n" ], "text/plain": [ "10:53:25.002 | \u001b[36mINFO\u001b[0m | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Created task run 'md-8' for task 'md'\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
10:53:25.005 | INFO    | Flow run 'masterful-panther' - Created task run 'md-3' for task 'md'\n",
       "
\n" ], "text/plain": [ "10:53:25.005 | \u001b[36mINFO\u001b[0m | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Created task run 'md-3' for task 'md'\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
10:53:25.008 | INFO    | Flow run 'masterful-panther' - Created task run 'md-2' for task 'md'\n",
       "
\n" ], "text/plain": [ "10:53:25.008 | \u001b[36mINFO\u001b[0m | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Created task run 'md-2' for task 'md'\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
10:53:25.577 | INFO    | Flow run 'masterful-panther' - Submitted task run 'md-0' for execution.\n",
       "
\n" ], "text/plain": [ "10:53:25.577 | \u001b[36mINFO\u001b[0m | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Submitted task run 'md-0' for execution.\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
10:53:25.592 | INFO    | Flow run 'masterful-panther' - Submitted task run 'md-1' for execution.\n",
       "
\n" ], "text/plain": [ "10:53:25.592 | \u001b[36mINFO\u001b[0m | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Submitted task run 'md-1' for execution.\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
10:53:25.606 | INFO    | Flow run 'masterful-panther' - Submitted task run 'md-8' for execution.\n",
       "
\n" ], "text/plain": [ "10:53:25.606 | \u001b[36mINFO\u001b[0m | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Submitted task run 'md-8' for execution.\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
10:53:25.613 | INFO    | Flow run 'masterful-panther' - Submitted task run 'md-3' for execution.\n",
       "
\n" ], "text/plain": [ "10:53:25.613 | \u001b[36mINFO\u001b[0m | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Submitted task run 'md-3' for execution.\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
10:53:25.623 | INFO    | Flow run 'masterful-panther' - Submitted task run 'md-2' for execution.\n",
       "
\n" ], "text/plain": [ "10:53:25.623 | \u001b[36mINFO\u001b[0m | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Submitted task run 'md-2' for execution.\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
10:53:25.638 | INFO    | Flow run 'masterful-panther' - Submitted task run 'md-4' for execution.\n",
       "
\n" ], "text/plain": [ "10:53:25.638 | \u001b[36mINFO\u001b[0m | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Submitted task run 'md-4' for execution.\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
10:53:25.713 | INFO    | Flow run 'masterful-panther' - Submitted task run 'md-5' for execution.\n",
       "
\n" ], "text/plain": [ "10:53:25.713 | \u001b[36mINFO\u001b[0m | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Submitted task run 'md-5' for execution.\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
10:53:25.831 | INFO    | Flow run 'masterful-panther' - Submitted task run 'md-6' for execution.\n",
       "
\n" ], "text/plain": [ "10:53:25.831 | \u001b[36mINFO\u001b[0m | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Submitted task run 'md-6' for execution.\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
10:53:25.903 | INFO    | Flow run 'masterful-panther' - Submitted task run 'md-7' for execution.\n",
       "
\n" ], "text/plain": [ "10:53:25.903 | \u001b[36mINFO\u001b[0m | Flow run\u001b[35m 'masterful-panther'\u001b[0m - Submitted task run 'md-7' for execution.\n" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "results = combustion(atoms)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] } ], "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 }