{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c6fba6a6",
   "metadata": {},
   "source": [
    "# QM/MM in Gromos"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42e53815",
   "metadata": {},
   "outputs": [],
   "source": [
    "# necessary imports are the Gromos_System and convenience functions emin and md\n",
    "import pygromos\n",
    "from pygromos.files.simulation_parameters.imd import Imd\n",
    "from pygromos.files.qmmm.qmmm import QMMM\n",
    "from pygromos.files.gromos_system.gromos_system import Gromos_System\n",
    "from pygromos.simulations.modules.preset_simulation_modules import emin, md\n",
    "from pygromos.data.simulation_parameters_templates import template_emin\n",
    "from pygromos.simulations.hpc_queuing.submission_systems.local import LOCAL\n",
    "\n",
    "# for file paths\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea7f87a7-1869-4560-a6cc-5511179addee",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Support for QMMM functionality in `GROMOS` input files\n",
    "\n",
    "This notebook demonstrates support of `PyGromosTools` for QM/MM functionality.\n",
    "\n",
    "https://github.com/rinikerlab/PyGromosTools/blob/qmmm/examples/example_gromos_qmmm.ipynb (part of the `qmmm` branch and soon to be merged to `release3`)\n",
    "\n",
    "Author: Felix Pultar\n",
    "\n",
    "\n",
    "Features include:\n",
    "\n",
    "* QM/MM blocks in `imd` files\n",
    "* QM/MM specification files\n",
    "* Running QM/MM simulations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "821fe3c7",
   "metadata": {},
   "source": [
    "### Load an `imd` file containing a QMMM block <a class=\"anchor\" id=\"imd-files\"></a>\n",
    "Simple demonstration of how to handle `.imd` files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbbf090b-19c7-48b5-86ac-ca0f6d1be117",
   "metadata": {},
   "outputs": [],
   "source": [
    "imd_path = os.path.abspath(\"example_files/QMMM_files/md.imd\")\n",
    "imd_file = Imd(imd_path)\n",
    "imd_file.TITLE.content = \"Demonstration of a Gromos imd file containing a QMMM block\"\n",
    "imd_file"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f48a2e4c-9b54-4c9c-9064-df84826b7b19",
   "metadata": {},
   "source": [
    "### Print out different sections of the QMMM block"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e27e062",
   "metadata": {},
   "source": [
    "Print out selected parameters from the `QMMM` block or also the `TITLE` block."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0794168f-e5ee-402d-aad9-a456ff4df6f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(imd_file.QMMM.NTQMMM) # QM/MM toggled on/off\n",
    "print(imd_file.QMMM.NTQMSW) # which QM/MM engine\n",
    "print(imd_file.TITLE.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9ef2fb7-0205-4a34-8024-fff058a1bab8",
   "metadata": {},
   "source": [
    "### Change a block value and print again"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6db72f93",
   "metadata": {},
   "source": [
    "Just change values of the `QMMM` block like with other `PyGromosTools` blocks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34e14585-76c7-4509-a121-17c58b56262b",
   "metadata": {},
   "outputs": [],
   "source": [
    "imd_file.QMMM.NTQMSW = 4 # switch to ORCA as QM software\n",
    "imd_file.QMMM"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27bbd639-1a57-4be3-aea5-67517721a6ed",
   "metadata": {},
   "source": [
    "## Directly manipulate a QMMM specification file"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d70086b-f262-4ec7-b7d7-d636eb8a08e8",
   "metadata": {},
   "source": [
    "The QMMM object allows to directly interact with QM/MM specification files. Future releases of `PyGromosTools` will also support generation of `QMMM` files from coordinate files (`.cnf`, `.xyz`, `.pdb`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "025685f2-7888-43d1-8815-78bd6182d610",
   "metadata": {},
   "outputs": [],
   "source": [
    "# instantiate the file object\n",
    "qmmm_file_path = os.path.abspath(\"example_files/QMMM_files/menthol-methanol-dmf.qmmm\")\n",
    "qmmm_file = QMMM(qmmm_file_path)\n",
    "print(qmmm_file)\n",
    "# There will be warnings if more than one QM engine is selected"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ff05c79-a2a7-4759-a378-f8651eeb62e0",
   "metadata": {},
   "source": [
    "## Print out and change some blocks in the QMMM specification file"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "330b78da-a354-4986-a6fe-f0640b9c7456",
   "metadata": {},
   "source": [
    "### Title block"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "637b6909",
   "metadata": {},
   "source": [
    "The `QMMM` specification file can be handled like any other `GROMOS` file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed139717-b743-4bf8-99e4-3732cd02f1ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(qmmm_file.TITLE.content)\n",
    "qmmm_file.TITLE.content = \"Custom file header\"\n",
    "print(qmmm_file.TITLE.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "083699b9-0a96-4380-ba54-d2529767135e",
   "metadata": {},
   "source": [
    "### QMZONE block"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e78da513",
   "metadata": {},
   "source": [
    "Print out the `QMZONE` section that defines which atoms will be treated quantum-mechanically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df7e4b44-9562-48b6-8597-bd4a9468e951",
   "metadata": {},
   "outputs": [],
   "source": [
    "# as in other Gromos files, the first bunch of characters are ignored and used to comment, e.g. name of the atom\n",
    "# second value: index of te position (starting from 1)\n",
    "# third value: element number according to the PSE\n",
    "# fourth value: indicate whether bond can be broken or not, default = 0\n",
    "print(qmmm_file.QMZONE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4b47246-47df-4758-b529-5cf7a14e57cc",
   "metadata": {},
   "source": [
    "### QMUNIT block"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a75b5c8c",
   "metadata": {},
   "source": [
    "Print out the `QMUNIT` block that defines some unit conversions between the MD engine and the QM software."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2191cba-b23c-4ff4-9e44-0febea848195",
   "metadata": {},
   "outputs": [],
   "source": [
    "# usually, these conversion factors are hard-coded in Gromos; left for historical reasons\n",
    "# first value: QM length to Gromos length (e.g. Bohr to nm) \n",
    "# second value: QM energy to Gromos energy (e.g. Hartree to kJ / mol)\n",
    "# third value: Gromos charge to QM charge (the same in this case)\n",
    "# fourth value: QM input units to Gromos input units (e.g. Angstrom to nm)\n",
    "print(qmmm_file.QMUNIT)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29f0397e-dbc2-4226-bec9-a8cb1d1c07df",
   "metadata": {},
   "source": [
    "### XTBELEMENTS block"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c127e63-9652-4119-81a8-939cdf79b14b",
   "metadata": {},
   "source": [
    "Print and update the `XTBELEMENTS` block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bfec885-5481-45bf-ba22-aac8e5498be6",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(qmmm_file.XTBELEMENTS)\n",
    "print(qmmm_file.XTBELEMENTS.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1503c040-bed2-4461-8f93-2262fb6b934f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# replace element numbers manually with the first ten elements of the PSE\n",
    "xtbelements_new  = [[str(i) for j in range(1)] for i in range(1,11)]\n",
    "print(xtbelements_new)\n",
    "qmmm_file.XTBELEMENTS.content = xtbelements_new"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6eb98bd1-0af0-4db5-becd-2af9a6a8bdf0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# show the updated section in the file\n",
    "qmmm_file.XTBELEMENTS"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11615f62-005e-4e99-84a8-7e2e32badb15",
   "metadata": {},
   "source": [
    "### A helper function that returns all QM engines specified in the QM/MM specification file"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "466a4510",
   "metadata": {},
   "source": [
    "There is also a sanity check in the constructor of `QMMM` to see if you did not accidentally add more than one QM engine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "171dcd38-50b3-4cbc-9056-1a0a83b679b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(qmmm_file.get_qm_engines())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6762dd06",
   "metadata": {},
   "source": [
    "### Store your QMMM specification file with all your other simulation files in a `Gromos_System` object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b53cc612",
   "metadata": {},
   "outputs": [],
   "source": [
    "# that's what we want to simulate\n",
    "system_name = \"menthol-methanol-dmf\"\n",
    "\n",
    "# that's where we want to simulate it at\n",
    "work_folder = f\"example_files/{system_name}\"\n",
    "\n",
    "# create a Gromos_System object from scratch\n",
    "system = Gromos_System(work_folder, system_name)\n",
    "\n",
    "# specify prepared topology, configuration, QMMM specification file, and input file\n",
    "system.top = f\"example_files/QMMM_files/{system_name}-all-atom_54a7.top\"\n",
    "system.cnf = f\"example_files/QMMM_files/{system_name}-all-atom-init_54a7.cnf\"\n",
    "system.qmmm = qmmm_file\n",
    "system.imd = imd_file\n",
    "\n",
    "# clean up\n",
    "system.rebase_files()\n",
    "\n",
    "# all your simulation files now live in the work folder\n",
    "system"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b311787",
   "metadata": {},
   "outputs": [],
   "source": [
    "# the imd file has been adapted (force groups, multibath, etc.)\n",
    "system.imd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8a4dd03",
   "metadata": {},
   "source": [
    "## Run QM/MM Simulations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f4b1139",
   "metadata": {},
   "source": [
    "QM/MM calculations are possible with a special in-house build of `GROMOS`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "477f3d2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# binaries (not yet QM/MM)\n",
    "gromosPP = None\n",
    "gromosXX = None\n",
    "\n",
    "# folders and title\n",
    "system_name = \"menthol-methanol-dmf\"\n",
    "work_folder =  f\"example_files/{system_name}\"\n",
    "\n",
    "# files\n",
    "in_cnf_path  = f\"example_files/QMMM_files/menthol-methanol-dmf-all-atom-init_54a7.cnf\"\n",
    "in_top_path  = f\"example_files/QMMM_files/menthol-methanol-dmf-all-atom_54a7.top\"\n",
    "\n",
    "# system\n",
    "system = Gromos_System(\n",
    "    work_folder, \n",
    "    system_name, \n",
    "    in_top_path=in_top_path, \n",
    "    in_cnf_path=in_cnf_path,\n",
    "    in_gromosPP_bin_dir=gromosPP,\n",
    "    in_gromosXX_bin_dir=gromosXX\n",
    ")\n",
    "\n",
    "system"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9edc9650",
   "metadata": {},
   "outputs": [],
   "source": [
    "system.cnf.view"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec3896bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a local submission system and specify the number of cores\n",
    "submit = LOCAL(nomp=8)\n",
    "\n",
    "# energy minimize the system\n",
    "minimized_system = emin(system, submission_system=submit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cf50815",
   "metadata": {},
   "outputs": [],
   "source": [
    "# new imd file for equilibration\n",
    "in_imd_path  = f\"example_files/QMMM_files/menthol-methanol-dmf-eq.imd\"\n",
    "minimized_system.imd = in_imd_path\n",
    "minimized_system.imd.STEP.NSTLIM = 1000\n",
    "minimized_system.imd.STEP.DT = 0.002 # 2.0 fs\n",
    "minimized_system"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ef864f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# equilibrate the system\n",
    "equilibrated_system = md(minimized_system, step_name=\"eq\", submission_system=submit)\n",
    "equilibrated_system"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac863a60",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load your favorite GCC version (required for your special GROMOS QM/MM build)\n",
    "os.environ[\"PATH\"] += \"/opt/gcc-8.2.0/bin\" #Change here!\n",
    "os.environ[\"LD_LIBRARY_PATH\"] = \"/opt/gcc-8.2.0/lib:/home/fpultar/opt/gcc-8.2.0/lib64\" #Change here!\n",
    "os.environ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d86372ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "# new imd file for QM/MM run\n",
    "in_imd_path  = f\"example_files/QMMM_files/md.imd\"\n",
    "equilibrated_system.imd = in_imd_path\n",
    "equilibrated_system.imd.STEP.NSTLIM = 100\n",
    "equilibrated_system.imd.STEP.DT = 0.0005 # 0.5 fs\n",
    "# qmmm specification file\n",
    "equilibrated_system.qmmm = QMMM(f\"example_files/QMMM_files/menthol-methanol-dmf.qmmm\")\n",
    "\n",
    "# now you want to switch to your special build of GROMOS :)\n",
    "gromosXX = None # path to freshly compiled gromos \"build-gcc-8.2.0-release/program\"\n",
    "equilibrated_system.gromosXX = gromosXX\n",
    "\n",
    "# check if everying is correct\n",
    "# note that the new .imd file and .qmmm file still live in the old location\n",
    "# while the .cnf file and .top file result from a previous simulation (equilibration)\n",
    "equilibrated_system"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f558f69e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# print the relevant QMMM block in the new imd file\n",
    "equilibrated_system.imd.QMMM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f0974e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# go QM/MM!\n",
    "production_system = md(equilibrated_system, submission_system=submit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abdadc91",
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize the last .cnf - you are done!\n",
    "production_system.cnf.view"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1c5194e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "b1b7b2ea43b8e767316eee98e01335d045804d2d47db68b6a5827e187ee91a7e"
  },
  "kernelspec": {
   "display_name": "Python 3.9.7 ('pygro2')",
   "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}