ase_example.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "# Set up logger\n",
    "logger = logging.getLogger()\n",
    "logger.handlers = []\n",
    "ch = logging.StreamHandler()\n",
    "formatter = logging.Formatter(\n",
    "    fmt=\"%(asctime)s (%(levelname)s): %(message)s\", datefmt=\"%Y-%m-%d %H:%M:%S\"\n",
    ")\n",
    "ch.setFormatter(formatter)\n",
    "logger.addHandler(ch)\n",
    "logger.setLevel(\"INFO\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from gemnet.model.gemnet import GemNet\n",
    "from gemnet.model.utils import read_json\n",
    "\n",
    "from ase_calculator import Molecule, MDSimulator\n",
    "from ase.build import molecule as ase_molecule_db\n",
    "\n",
    "# for visualization\n",
    "from ase.io.trajectory import Trajectory\n",
    "import nglview"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"GemNet-Q\"\n",
    "# model_name = \"GemNet-T\"\n",
    "\n",
    "pretrained_models_path = \"./pretrained\"\n",
    "weights_file = f\"{pretrained_models_path}/{model_name}/weights/ckpt\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num_spherical: 7\n",
      "num_radial: 6\n",
      "num_blocks: 4\n",
      "emb_size_atom: 128\n",
      "emb_size_edge: 128\n",
      "emb_size_trip: 64\n",
      "emb_size_quad: 32\n",
      "emb_size_rbf: 16\n",
      "emb_size_cbf: 16\n",
      "emb_size_sbf: 32\n",
      "emb_size_bil_trip: 64\n",
      "emb_size_bil_quad: 32\n",
      "num_before_skip: 1\n",
      "num_after_skip: 1\n",
      "num_concat: 1\n",
      "num_atom: 2\n",
      "triplets_only: False\n",
      "num_targets: 1\n",
      "direct_forces: False\n",
      "cutoff: 5.0\n",
      "int_cutoff: 10.0\n",
      "envelope_exponent: 5\n",
      "extensive: True\n",
      "forces_coupled: False\n",
      "output_init: HeOrthogonal\n",
      "activation: swish\n",
      "scale_file: ./pretrained/scaling_factors.json\n"
     ]
    }
   ],
   "source": [
    "model_kwargs = read_json(f\"{pretrained_models_path}/{model_name}/model_kwargs.json\")\n",
    "model_kwargs[\"scale_file\"] = f\"{pretrained_models_path}/\" + model_kwargs[\"scale_file\"]\n",
    "\n",
    "for key, value in model_kwargs.items():\n",
    "    print(f\"{key}: {value}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.training.tracking.util.CheckpointLoadStatus at 0x7f80fee37520>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# # uncomment to run in eager mode\n",
    "# tf.config.experimental_run_functions_eagerly(True) \n",
    "\n",
    "model = GemNet(**model_kwargs)\n",
    "model.load_weights(weights_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Molecule setup\n",
    "Load from database or build your own by specifying R and Z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "mol = ase_molecule_db('C7NH5')\n",
    "R   = mol.get_positions()\n",
    "Z   = mol.get_atomic_numbers()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MD simulation settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "traj_path = \"./md_sim.traj\"\n",
    "logfile = \"-\"               # “-” refers to standard output.\n",
    "dynamics = \"langevin\"       # Name of the MD integrator. Implemented: 'langevin' or 'verlet'.\n",
    "max_steps = 10              # Maximum number of simulation steps.\n",
    "time = 0.5                  # Integration time step for Newton's law in femtoseconds.\n",
    "interval = 2                # Write only every <interval> time step to trajectory file.\n",
    "temperature = 1500          # The temperature in Kelvin.\n",
    "langevin_friction = 0.002   # Friction coefficient (only used when dynamics is langevin)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup and run the simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "cutoff = model_kwargs[\"cutoff\"]\n",
    "int_cutoff = model_kwargs[\"int_cutoff\"]\n",
    "triplets_only = model_kwargs[\"triplets_only\"]\n",
    "molecule = Molecule(\n",
    "    R, Z, cutoff=cutoff, int_cutoff=int_cutoff, triplets_only=triplets_only\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-03-15 19:41:11 (INFO): Selected MD integrator: Langevin\n",
      "2022-03-15 19:41:11 (INFO): Save trajectory to ./md_sim.traj\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time[ps]      Etot[eV]     Epot[eV]     Ekin[eV]    T[K]\n",
      "0.0000         -74.6765     -77.4343       2.7578  1641.2\n",
      "0.0010         -74.6480     -77.2796       2.6316  1566.1\n",
      "0.0020         -74.6280     -76.9330       2.3051  1371.8\n",
      "0.0030         -74.6367     -76.5432       1.9065  1134.6\n",
      "0.0040         -74.6339     -76.2188       1.5849   943.2\n",
      "0.0050         -74.6246     -75.9957       1.3711   816.0\n"
     ]
    }
   ],
   "source": [
    "simulation = MDSimulator(\n",
    "    molecule, model, \n",
    "    dynamics=dynamics, max_steps=max_steps, time=time, temperature=temperature, langevin_friction=langevin_friction,\n",
    "    interval=interval, traj_path=traj_path, logfile=logfile\n",
    ")\n",
    "simulation.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualize simulated trajectory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "67c5efc7e941477b84d79df586e34989",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "NGLWidget(count=6)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "traj = Trajectory(traj_path)\n",
    "nglview.show_asetraj(traj)"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "73d4dc6ffc134dc5e05ee963c4039b14792ec4f63c8d27e3dd67b524fa7b1d65"
  },
  "kernelspec": {
   "display_name": "Python 3.8.0 ('tf')",
   "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.8.0"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}