From f1a3c1b70a761c27643e3f9c097c93205c2ceda2 Mon Sep 17 00:00:00 2001
From: Alexander Reinauer <areinauer@icp.uni-stuttgart.de>
Date: Mon, 9 Sep 2024 15:08:09 +0200
Subject: [PATCH] Some cleanup and figuring out environment issues

---
 doc/tutorials/mlip/mlip.ipynb | 82 +++++++++++++++++++++++++++--------
 1 file changed, 65 insertions(+), 17 deletions(-)

diff --git a/doc/tutorials/mlip/mlip.ipynb b/doc/tutorials/mlip/mlip.ipynb
index 5175ac8a58..bd11985566 100644
--- a/doc/tutorials/mlip/mlip.ipynb
+++ b/doc/tutorials/mlip/mlip.ipynb
@@ -29,6 +29,26 @@
     "We can use "
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# packmol"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# import os\n",
+    "# os.environ[\"PATH\"] = f\"{os.environ['PATH']}:/tikhome/ess_delme/Documents/test_for_ml_tutorial/espresso/build\"\n",
+    "# os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\""
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -58,7 +78,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "For MLIP the typical unit system is energies in eV, distances in Angstrom, time in fs and mass in atomic units. This is also the default unit system within ASE. We will use pint to ensure correct unit usage within ESPResSo. For setting the simulation temperature, we will also convert the boltzmann constant `boltzmann_k` to our desired unit system."
+    "For MLIP the typical unit system is energies in eV, distances in Angstrom, time in fs and mass in atomic units. This is also the default unit system within ASE. We will use pint to ensure correct unit usage within ESPResSo."
    ]
   },
   {
@@ -67,9 +87,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "ureg = pint.UnitRegistry()\n",
-    "boltzmann_k = 1.380649e-23 * ureg.J / ureg.K\n",
-    "boltzmann_k = boltzmann_k.to(ureg.eV / ureg.K)"
+    "ureg = pint.UnitRegistry()"
    ]
   },
   {
@@ -148,9 +166,8 @@
    "source": [
     "system = espressomd.System(box_l=[16] * 3) # A cubic box of 16 Angstroms\n",
     "system.time_step = (\n",
-    "    (0.05 * ureg.fs)\n",
-    "    .to(((1 * ureg.u * ureg.angstrom**2) / ureg.electron_volt) ** 0.5)\n",
-    "    .magnitude\n",
+    "    (0.5 * ureg.fs)\n",
+    "    .m_as(((1 * ureg.u * ureg.angstrom**2) / ureg.electron_volt) ** 0.5)\n",
     ") # we convert the time units from fs to be in line with our other units for mass, distance and energy.\n",
     "system.cell_system.skin = 0.4"
    ]
@@ -210,7 +227,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "tbar = tqdm.trange(500, ncols=120)\n",
+    "tbar = tqdm.trange(10, ncols=120)\n",
     "for idx in tbar:\n",
     "    forces = mace_mp_calc.get_forces(system.ase.get())\n",
     "    system.part.all().ext_force = forces\n",
@@ -230,7 +247,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### MD Simulations of Sulfuric Acid and Water"
+    "# MD Simulations of Sulfuric Acid and Water"
    ]
   },
   {
@@ -284,7 +301,7 @@
    "source": [
     "system.box_l = box_of_atoms.get_cell().diagonal()\n",
     "system.time_step = (\n",
-    "    (0.05 * ureg.fs)\n",
+    "    (0.25 * ureg.fs)\n",
     "    .to(((1 * ureg.u * ureg.angstrom**2) / ureg.electron_volt) ** 0.5)\n",
     "    .magnitude\n",
     ")\n",
@@ -313,9 +330,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "scrolled": true
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "del vis.zndraw[:]\n",
@@ -330,7 +345,16 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "tbar = tqdm.trange(500, ncols=120)\n",
+    "vis.zndraw.config.scene.vector_scale = 0"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tbar = tqdm.trange(50, ncols=120)\n",
     "for _ in tbar:\n",
     "    atoms = system.ase.get()\n",
     "    atoms.calc = mace_mp_calc\n",
@@ -367,7 +391,7 @@
    "source": [
     "# Langevin dynamics at 400 K\n",
     "system.integrator.set_vv()\n",
-    "system.thermostat.set_langevin(kT=(400 * ureg.K * boltzmann_k).magnitude, gamma=2, seed=42)"
+    "system.thermostat.set_langevin(kT=(400 * ureg.K * ureg.boltzmann_constant).m_as(\"eV\"), gamma=2, seed=42)"
    ]
   },
   {
@@ -376,7 +400,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "tbar = tqdm.trange(5000, ncols=120)\n",
+    "tbar = tqdm.trange(500, ncols=120)\n",
     "for idx in tbar:\n",
     "    atoms = system.ase.get()\n",
     "    atoms.calc = mace_mp_calc\n",
@@ -402,9 +426,11 @@
    "source": [
     "OH_1 = []\n",
     "OH_2 = []\n",
+    "energies = []\n",
     "for atoms in vis.zndraw:\n",
     "    OH_1.append(atoms.get_distance(300, 305, mic=True))\n",
     "    OH_2.append(atoms.get_distance(304, 306, mic=True))\n",
+    "    energies.append(atoms.get_potential_energy())\n",
     "\n",
     "df = pd.DataFrame({\"OH_1\": OH_1, \"OH_2\": OH_2})\n",
     "fig = px.line(df, y=[\"OH_1\", \"OH_2\"])"
@@ -426,6 +452,28 @@
     "vis.zndraw.figures = {\"distance\": fig.to_json()}"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "fig"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "df = pd.DataFrame({\"Energies\": energies})\n",
+    "fig = px.line(df, y=\"Energies\")\n",
+    "fig"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -467,7 +515,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.14"
+   "version": "3.10.12"
   }
  },
  "nbformat": 4,