examples/Ammonia_MD.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NNMD of ammonia using the pretrained models\n",
    "\n",
    "In this notebook, we perform a NN-based MD NVE simulation using ASE and the [SchNet NN potential](https://github.com/learningmatter-mit/NeuralForceField). We will be using the third generation of ammonia models, as shown in [our paper](https://arxiv.org/abs/2101.11588).\n",
    "\n",
    "The utilities at `nff` will be used to perform the MD simulation. `nglview` will be used to visualize the generated trajectories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0bfb91771f834f83b5c87be0c989506c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "from ase import Atoms, units\n",
    "from ase.md.verlet import VelocityVerlet\n",
    "\n",
    "from nff.io import NeuralFF, AtomsBatch, EnsembleNFF\n",
    "from nff.md.nve import Dynamics\n",
    "from nff.md.nvt import NoseHoover\n",
    "from nff.data import Dataset\n",
    "import nff.utils.constants as const\n",
    "\n",
    "import nglview as nv\n",
    "from ase.io import Trajectory, read"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading the models and parameters\n",
    "\n",
    "The dataset contains a PyTorch dataset with all the DFT data used to train the third generation of NN potentials. Here, we will use the pre-trained ensemble. For learning how to train the models using the SchNet architecture and the current dataset, check the tutorials at the original [NFF repo](https://github.com/learningmatter-mit/NeuralForceField)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "dset = Dataset.from_file('../data/ammonia.pth.tar')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`DEVICE` sets the GPU used for evaluating the model. If you want to evaluate on a CPU, set `DEVICE = 'cpu'`. The models are stored at `/models/ammonia`, where `/` is the root folder of this repository."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "DEVICE = 2\n",
    "\n",
    "def get_ensemble_calc(device=DEVICE):\n",
    "    PATH = '../models/ammonia'\n",
    "    models = []\n",
    "    for model_name in sorted(os.listdir(PATH)):\n",
    "        m = NeuralFF.from_file(os.path.join(PATH, model_name), device=device).model\n",
    "        models.append(m)\n",
    "\n",
    "    return EnsembleNFF(models, device=device)\n",
    "\n",
    "ensemble = get_ensemble_calc(device=DEVICE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we set the parameters for the MD simulation. For learning how to use these parameters within our code, check the tutorials at the original [NFF repo](https://github.com/learningmatter-mit/NeuralForceField)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def get_md_params(traj_filename, temperature=1000):\n",
    "    return {\n",
    "        'T_init': temperature,\n",
    "        'time_step': 1.0,\n",
    "        'thermostat': VelocityVerlet,  \n",
    "        'thermostat_params': {'timestep': 1.0 * units.fs},\n",
    "        'steps': 5 * 1000,\n",
    "        'save_frequency': 40,\n",
    "        'nbr_list_update_freq': 10,\n",
    "        'thermo_filename': 'thermo.log',\n",
    "        'traj_filename': traj_filename,\n",
    "        'skip': 0\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we use the lowest energy conformation within the existing dataset (i.e. the ground state of ammonia) as a starting configuration for the MD simulation. `AtomsBatch` is a [wrapper within our NFF repo](https://github.com/learningmatter-mit/NeuralForceField/blob/master/nff/io/ase.py) and can be used to interface an ASE atom with NFF."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "CUTOFF = 5.0\n",
    "\n",
    "def get_md_atoms(dset=dset, cutoff=CUTOFF, device=DEVICE):\n",
    "    props = dset[np.argmin(dset.props['energy'])]\n",
    "\n",
    "    atoms = AtomsBatch(\n",
    "        positions=props['nxyz'][:, 1:],\n",
    "        numbers=props['nxyz'][:, 0],\n",
    "        cutoff=cutoff,\n",
    "        props={'energy': 0, 'energy_grad': []},\n",
    "        calculator=ensemble,\n",
    "        device=device,\n",
    "    )\n",
    "    _ = atoms.update_nbr_list()\n",
    "    \n",
    "    return atoms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Performing the MD simulation\n",
    "\n",
    "Now, we perform the MD simulation using the parameters shown before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/dskoda/.conda/envs/nff/lib/python3.7/site-packages/ase/md/md.py:48: FutureWarning: Specify the temperature in K using the 'temperature_K' argument\n",
      "  warnings.warn(FutureWarning(w))\n",
      "/home/dskoda/.conda/envs/nff/lib/python3.7/site-packages/ase/atoms.py:968: VisibleDeprecationWarning: Use get_global_number_of_atoms() instead\n",
      "  np.VisibleDeprecationWarning)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time[ps]      Etot[eV]     Epot[eV]     Ekin[eV]    T[K]\n",
      "0.0000           0.9087      -0.0017       0.9104  1760.8\n",
      "\n",
      "0.0400           0.9214       0.2683       0.6532  1263.3\n",
      "\n",
      "0.0800           0.9166       0.3652       0.5513  1066.3\n",
      "\n",
      "0.1200           0.9147       0.3221       0.5925  1146.0\n",
      "\n",
      "0.1600           0.9282       0.7402       0.1880   363.6\n",
      "\n",
      "0.2000           0.9259       0.3814       0.5445  1053.1\n",
      "\n",
      "0.2400           0.9245       0.4986       0.4260   823.8\n",
      "\n",
      "0.2800           0.9210       0.2714       0.6496  1256.3\n",
      "\n",
      "0.3200           0.9289       0.3034       0.6255  1209.7\n",
      "\n",
      "0.3600           0.9235       0.2768       0.6467  1250.8\n",
      "\n",
      "0.4000           0.9313       0.4563       0.4750   918.6\n",
      "\n",
      "0.4400           0.9301       0.7316       0.1985   383.8\n",
      "\n",
      "0.4800           0.9296       0.4780       0.4516   873.5\n",
      "\n",
      "0.5200           0.9296       0.8318       0.0977   189.0\n",
      "\n",
      "0.5600           0.9210       0.4233       0.4976   962.5\n",
      "\n",
      "0.6000           0.9176       0.2918       0.6258  1210.3\n",
      "\n",
      "0.6400           0.9225       0.2324       0.6901  1334.8\n",
      "\n",
      "0.6800           0.9101       0.1030       0.8071  1561.0\n",
      "\n",
      "0.7200           0.9152       0.2738       0.6415  1240.7\n",
      "\n",
      "0.7600           0.9247       0.2986       0.6261  1211.0\n",
      "\n",
      "0.8000           0.9197       0.5639       0.3558   688.2\n",
      "\n",
      "0.8400           0.9066       0.2816       0.6250  1208.9\n",
      "\n",
      "0.8800           0.9278       0.5785       0.3493   675.5\n",
      "\n",
      "0.9200           0.9263       0.2901       0.6363  1230.6\n",
      "\n",
      "0.9600           0.9132       0.0646       0.8487  1641.4\n",
      "\n",
      "1.0000           0.9116       0.0554       0.8562  1655.9\n",
      "\n",
      "1.0400           0.9355       0.3203       0.6152  1189.9\n",
      "\n",
      "1.0800           0.9340       0.5438       0.3902   754.7\n",
      "\n",
      "1.1200           0.9238       0.4716       0.4522   874.7\n",
      "\n",
      "1.1600           0.9246       0.6515       0.2731   528.1\n",
      "\n",
      "1.2000           0.9253       0.3434       0.5819  1125.5\n",
      "\n",
      "1.2400           0.9141       0.2442       0.6698  1295.5\n",
      "\n",
      "1.2800           0.9207       0.1735       0.7472  1445.2\n",
      "\n",
      "1.3200           0.9096       0.1154       0.7942  1536.1\n",
      "\n",
      "1.3600           0.9094       0.2680       0.6414  1240.5\n",
      "\n",
      "1.4000           0.9170       0.4194       0.4977   962.5\n",
      "\n",
      "1.4400           0.9206       0.6615       0.2592   501.3\n",
      "\n",
      "1.4800           0.9180       0.3204       0.5977  1155.9\n",
      "\n",
      "1.5200           0.9218       0.2915       0.6302  1218.9\n",
      "\n",
      "1.5600           0.9198       0.1926       0.7272  1406.5\n",
      "\n",
      "1.6000           0.9285       0.3954       0.5331  1031.1\n",
      "\n",
      "1.6400           0.9260       0.5635       0.3626   701.2\n",
      "\n",
      "1.6800           0.9320       0.4972       0.4348   841.0\n",
      "\n",
      "1.7200           0.9328       0.8561       0.0767   148.3\n",
      "\n",
      "1.7600           0.9330       0.5354       0.3975   768.9\n",
      "\n",
      "1.8000           0.9223       0.2810       0.6412  1240.2\n",
      "\n",
      "1.8400           0.9278       0.3132       0.6146  1188.6\n",
      "\n",
      "1.8800           0.9269       0.4055       0.5213  1008.3\n",
      "\n",
      "1.9200           0.9315       0.7550       0.1765   341.3\n",
      "\n",
      "1.9600           0.9218       0.3835       0.5384  1041.2\n",
      "\n",
      "2.0000           0.9138       0.2979       0.6159  1191.3\n",
      "\n",
      "2.0400           0.9141       0.4604       0.4537   877.5\n",
      "\n",
      "2.0800           0.9186       0.3093       0.6093  1178.4\n",
      "\n",
      "2.1200           0.9078       0.3341       0.5738  1109.7\n",
      "\n",
      "2.1600           0.8998       0.3186       0.5812  1124.0\n",
      "\n",
      "2.2000           0.9042       0.0828       0.8214  1588.6\n",
      "\n",
      "2.2400           0.9062       0.4265       0.4797   927.8\n",
      "\n",
      "2.2800           0.9118       0.1289       0.7829  1514.2\n",
      "\n",
      "2.3200           0.9191       0.4702       0.4489   868.2\n",
      "\n",
      "2.3600           0.9430       0.4008       0.5421  1048.5\n",
      "\n",
      "2.4000           0.9405       0.6690       0.2715   525.2\n",
      "\n",
      "2.4400           0.9263       0.3836       0.5427  1049.6\n",
      "\n",
      "2.4800           0.9513       0.6345       0.3168   612.7\n",
      "\n",
      "2.5200           0.9306       0.6190       0.3116   602.6\n",
      "\n",
      "2.5600           0.9175       0.2622       0.6554  1267.5\n",
      "\n",
      "2.6000           0.9051       0.2637       0.6414  1240.5\n",
      "\n",
      "2.6400           0.9198       0.6736       0.2462   476.1\n",
      "\n",
      "2.6800           0.9158       0.3884       0.5274  1020.1\n",
      "\n",
      "2.7200           0.9173       0.2625       0.6548  1266.5\n",
      "\n",
      "2.7600           0.8954       0.0892       0.8062  1559.2\n",
      "\n",
      "2.8000           0.9057       0.1676       0.7381  1427.5\n",
      "\n",
      "2.8400           0.8964       0.2756       0.6208  1200.7\n",
      "\n",
      "2.8800           0.9062       0.4084       0.4979   962.9\n",
      "\n",
      "2.9200           0.9120       0.4014       0.5106   987.5\n",
      "\n",
      "2.9600           0.9225       0.4720       0.4505   871.3\n",
      "\n",
      "3.0000           0.9240       0.4398       0.4842   936.5\n",
      "\n",
      "3.0400           0.9281       0.6666       0.2614   505.6\n",
      "\n",
      "3.0800           0.9279       0.7390       0.1889   365.3\n",
      "\n",
      "3.1200           0.9177       0.5369       0.3808   736.5\n",
      "\n",
      "3.1600           0.9219       0.3746       0.5473  1058.5\n",
      "\n",
      "3.2000           0.9320       0.5801       0.3518   680.5\n",
      "\n",
      "3.2400           0.9297       0.7005       0.2292   443.4\n",
      "\n",
      "3.2800           0.9266       0.6602       0.2664   515.3\n",
      "\n",
      "3.3200           0.9325       0.7101       0.2224   430.1\n",
      "\n",
      "3.3600           0.9299       0.5498       0.3800   735.0\n",
      "\n",
      "3.4000           0.9264       0.4887       0.4377   846.5\n",
      "\n",
      "3.4400           0.9272       0.6994       0.2277   440.5\n",
      "\n",
      "3.4800           0.9295       0.7339       0.1956   378.3\n",
      "\n",
      "3.5200           0.9317       0.5630       0.3687   713.2\n",
      "\n",
      "3.5600           0.9250       0.4957       0.4293   830.3\n",
      "\n",
      "3.6000           0.9236       0.4842       0.4394   849.8\n",
      "\n",
      "3.6400           0.9300       0.7680       0.1620   313.3\n",
      "\n",
      "3.6800           0.9285       0.7321       0.1965   380.0\n",
      "\n",
      "3.7200           0.9269       0.6746       0.2523   488.0\n",
      "\n",
      "3.7600           0.9320       0.6078       0.3243   627.1\n",
      "\n",
      "3.8000           0.9297       0.6774       0.2523   488.0\n",
      "\n",
      "3.8400           0.9178       0.6046       0.3132   605.8\n",
      "\n",
      "3.8800           0.9306       0.5339       0.3967   767.2\n",
      "\n",
      "3.9200           0.9257       0.5905       0.3351   648.2\n",
      "\n",
      "3.9600           0.9078       0.2459       0.6619  1280.2\n",
      "\n",
      "4.0000           0.8969       0.2461       0.6508  1258.7\n",
      "\n",
      "4.0400           0.9055       0.4437       0.4619   893.3\n",
      "\n",
      "4.0800           0.9051       0.4308       0.4742   917.2\n",
      "\n",
      "4.1200           0.8886       0.3047       0.5839  1129.4\n",
      "\n",
      "4.1600           0.8984       0.2713       0.6271  1212.8\n",
      "\n",
      "4.2000           0.8930       0.1797       0.7132  1379.5\n",
      "\n",
      "4.2400           0.9087       0.2992       0.6095  1178.9\n",
      "\n",
      "4.2800           0.9258       0.5443       0.3815   737.9\n",
      "\n",
      "4.3200           0.9276       0.6574       0.2703   522.7\n",
      "\n",
      "4.3600           0.9241       0.5891       0.3350   647.9\n",
      "\n",
      "4.4000           0.9194       0.4911       0.4283   828.3\n",
      "\n",
      "4.4400           0.9253       0.4961       0.4292   830.2\n",
      "\n",
      "4.4800           0.9202       0.4338       0.4864   940.8\n",
      "\n",
      "4.5200           0.9157       0.4633       0.4524   874.9\n",
      "\n",
      "4.5600           0.9333       0.7549       0.1783   344.9\n",
      "\n",
      "4.6000           0.9273       0.6801       0.2472   478.2\n",
      "\n",
      "4.6400           0.9256       0.7345       0.1910   369.4\n",
      "\n",
      "4.6800           0.9332       0.8131       0.1201   232.3\n",
      "\n",
      "4.7200           0.9005       0.4375       0.4631   895.6\n",
      "\n",
      "4.7600           0.9141       0.3903       0.5238  1013.0\n",
      "\n",
      "4.8000           0.9195       0.4367       0.4828   933.8\n",
      "\n",
      "4.8400           0.9063       0.3267       0.5796  1121.0\n",
      "\n",
      "4.8800           0.9148       0.5708       0.3440   665.3\n",
      "\n",
      "4.9200           0.9153       0.6782       0.2371   458.6\n",
      "\n",
      "4.9600           0.9340       0.4754       0.4586   887.1\n",
      "\n",
      "5.0000           0.9114       0.6793       0.2321   449.0\n",
      "\n"
     ]
    }
   ],
   "source": [
    "TEMPERATURE = 1500\n",
    "\n",
    "name = f'NVE_{TEMPERATURE}.traj'\n",
    "\n",
    "atoms = get_md_atoms()\n",
    "md_params = get_md_params(name, temperature=TEMPERATURE)\n",
    "dyn = Dynamics(atoms, md_params)\n",
    "traj = dyn.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualizing the trajectory\n",
    "\n",
    "Finally, we can visualize the trajectory on the Jupyter server using `nglview`. The filename was selected prior to the MD simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4aaa6d820752460aa5ca11849daace0b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "NGLWidget(max_frame=125)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "filetraj = Trajectory(name)\n",
    "view = nv.show_asetraj(filetraj)\n",
    "view"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nff",
   "language": "python",
   "name": "nff"
  },
  "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}