Test MTK barostat implementation

run.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Union, IO\n",
    "import weakref\n",
    "from concurrent.futures import ThreadPoolExecutor, as_completed\n",
    "\n",
    "import numpy as np\n",
    "import asap3\n",
    "from ase.parallel import world\n",
    "from ase.md.md import MolecularDynamics\n",
    "from ase.md.logger import MDLogger\n",
    "from ase import Atoms\n",
    "import ase.units\n",
    "import ase.build\n",
    "from ase.calculators.emt import EMT\n",
    "from ase.md.velocitydistribution import MaxwellBoltzmannDistribution, Stationary\n",
    "\n",
    "from ase.md.nose_hoover_chain import NoseHooverChainNVT, IsotropicMTKNPT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomMDLogger(MDLogger):\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        dyn: MolecularDynamics,\n",
    "        atoms: Atoms,\n",
    "        logfile: Union[IO, str],\n",
    "    ):\n",
    "        self.dyn: MolecularDynamics = weakref.proxy(dyn)\n",
    "        self.atoms = atoms\n",
    "\n",
    "        self.logfile = self.openfile(file=logfile, mode='w', comm=world)\n",
    "\n",
    "        self.hdr = \"step,time(ps),epot(eV/atom),ekin(eV/atom),etot(eV/atom),temperature(K),volume(Ang3/atom),pxx(GPa),pyy(GPa),pzz(GPa),pyz(GPa),pzx(GPa),pxy(GPa)\"\n",
    "        self.fmt = \"%d\" + \",%.8f\" * 12 + \"\\n\"\n",
    "\n",
    "        self.logfile.write(self.hdr + \"\\n\")\n",
    "\n",
    "    def __call__(self):\n",
    "        num_atoms = len(self.atoms)\n",
    "        epot_eV_per_atom = self.atoms.get_potential_energy() / num_atoms\n",
    "        ekin_eV_per_atom = self.atoms.get_kinetic_energy() / num_atoms\n",
    "        etot_eV_per_atom = epot_eV_per_atom + ekin_eV_per_atom\n",
    "        temperature_K = self.atoms.get_temperature()\n",
    "        volume_Ang3_per_atom = self.atoms.get_volume() / num_atoms\n",
    "\n",
    "        stress_GPa = self.atoms.get_stress(include_ideal_gas=True, voigt=False) / ase.units.GPa\n",
    "        pxx_GPa = -stress_GPa[0, 0]\n",
    "        pyy_GPa = -stress_GPa[1, 1]\n",
    "        pzz_GPa = -stress_GPa[2, 2]\n",
    "        pyz_GPa = -stress_GPa[1, 2]\n",
    "        pzx_GPa = -stress_GPa[2, 0]\n",
    "        pxy_GPa = -stress_GPa[0, 1]\n",
    "\n",
    "        steps = self.dyn.nsteps\n",
    "        time_ps = self.dyn.get_time() / (1000 * ase.units.fs)\n",
    "\n",
    "        dat = (\n",
    "            steps,\n",
    "            time_ps,\n",
    "            epot_eV_per_atom,\n",
    "            ekin_eV_per_atom,\n",
    "            etot_eV_per_atom,\n",
    "            temperature_K,\n",
    "            volume_Ang3_per_atom,\n",
    "            pxx_GPa,\n",
    "            pyy_GPa,\n",
    "            pzz_GPa,\n",
    "            pyz_GPa,\n",
    "            pzx_GPa,\n",
    "            pxy_GPa,\n",
    "        )\n",
    "\n",
    "        self.logfile.write(self.fmt % dat)\n",
    "        self.logfile.flush()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## NVT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_nvt(tchain: int, tdamp_timestep: int):\n",
    "    atoms = ase.build.bulk(\"Cu\", crystalstructure='hcp', a=2.53, c=4.11).repeat((4, 4, 4))\n",
    "    atoms.calc = EMT()\n",
    "\n",
    "    temperature_K = 300\n",
    "    rng = np.random.default_rng(0)\n",
    "    MaxwellBoltzmannDistribution(atoms, temperature_K=temperature_K, force_temp=True, rng=rng)\n",
    "    Stationary(atoms)\n",
    "\n",
    "    timestep = 1.0 * ase.units.fs\n",
    "    md = NoseHooverChainNVT(\n",
    "        atoms,\n",
    "        timestep=timestep,\n",
    "        temperature_K=temperature_K,\n",
    "        tdamp=tdamp_timestep * timestep,\n",
    "        tchain=tchain,\n",
    "    )\n",
    "    logfile = f'nvt_tchain-{tchain}-{tdamp_timestep}.log'\n",
    "    md.attach(CustomMDLogger(md, atoms, logfile), interval=100)\n",
    "    md.run(100_000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "run(tchain=3, tdamp_timestep=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Isotropic NPT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_iso_npt(pdamp_timestep: int, pressure_GPa: int):\n",
    "    atoms = ase.build.bulk(\"Cu\", crystalstructure='hcp', a=2.53, c=4.11).repeat((4, 4, 4))\n",
    "    atoms.calc = asap3.EMT()\n",
    "\n",
    "    temperature_K = 300\n",
    "    rng = np.random.default_rng(0)\n",
    "    MaxwellBoltzmannDistribution(atoms, temperature_K=temperature_K, force_temp=True, rng=rng)\n",
    "    Stationary(atoms)\n",
    "\n",
    "    timestep = 1.0 * ase.units.fs\n",
    "    md = IsotropicMTKNPT(\n",
    "        atoms,\n",
    "        timestep=timestep,\n",
    "        temperature_K=temperature_K,\n",
    "        pressure_GPa=float(pressure_GPa),\n",
    "        tdamp=100 * timestep,\n",
    "        pdamp=pdamp_timestep * timestep,\n",
    "    )\n",
    "    logfile = f'iso-npt_pressure-{pressure_GPa}_pdamp-{pdamp_timestep}.log'\n",
    "    md.attach(CustomMDLogger(md, atoms, logfile), interval=100)\n",
    "    md.run(100_000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "futures = []\n",
    "with ThreadPoolExecutor(max_workers=8) as executor:\n",
    "    for pressure_GPa in [0, 10, 100]:\n",
    "        for pdamp_timestep in [100, 1000, 10000]:\n",
    "            future = executor.submit(\n",
    "                run_iso_npt,\n",
    "                pdamp_timestep=pdamp_timestep,\n",
    "                pressure_GPa=pressure_GPa,\n",
    "            )\n",
    "            futures.append(future)\n",
    "\n",
    "for future in as_completed(futures):\n",
    "    _ = future.result()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ase",
   "language": "python",
   "name": "python3"
  },
  "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}

view.ipynb