{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# TI Calculation\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bb7af915f1e54549a21770e380f1d64c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#for analysis\n",
    "from matplotlib import pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "import os\n",
    "import numpy as np\n",
    "from pygromos.files.gromos_system import Gromos_System\n",
    "from pygromos.simulations.hpc_queuing.submission_systems.local import LOCAL as subSystem\n",
    "from pygromos.files.blocks.imd_blocks import PERTURBATION, WRITETRAJ, DISTANCERES\n",
    "\n",
    "from pygromos.utils import bash"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Input files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "root_dir = os.path.abspath(\"example_files/TI_Calculation\")\n",
    "root_in_dir = root_dir+\"/TI_input\"\n",
    "cnf_path = root_in_dir+\"/M030_6KET.cnf\"\n",
    "top_path = root_in_dir + \"/M030_6KET.top\"\n",
    "disres_path = root_in_dir+\"/M030_6KET.disres\"\n",
    "\n",
    "\n",
    "sys_name = \"M030_to_6KET\"\n",
    "lam = 0\n",
    "\n",
    "project_dir = bash.make_folder(root_dir+\"/\"+sys_name)\n",
    "input_dir = bash.make_folder(project_dir+\"/input\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Vacuum Simulation\n",
    "\n",
    "### Direction A->B\n",
    "\n",
    "#### Setup:\n",
    "\n",
    "\n",
    "##### Build pertubation file\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pygromos.files.topology.ptp import Pertubation_topology\n",
    "from pygromos.files.blocks.topology_blocks import pertubation_lam_state, atom_lam_pertubation_state, PERTATOMPARAM, TITLE\n",
    "\n",
    "\n",
    "#External imd_changes:\n",
    "grom_system = Gromos_System(in_cnf_path=cnf_path, in_top_path=top_path,\n",
    "                            in_disres_path=disres_path,\n",
    "                            system_name=sys_name, work_folder=input_dir)\n",
    "\n",
    "\n",
    "#Build up lambda - States\n",
    "pert_atoms=[]\n",
    "for atom_line in grom_system.top.SOLUTEATOM:\n",
    "    states = {}\n",
    "    phys_state = pertubation_lam_state(IAC=atom_line.IAC, MASS=atom_line.MASS, CHARGE=atom_line.CG)\n",
    "    states = {atom_line.MRES: phys_state }\n",
    "\n",
    "    pert_atom = atom_lam_pertubation_state(atom_line.ATNM,RES=atom_line.MRES,NAME=atom_line.PANM, STATES=states)\n",
    "    pert_atoms.append(pert_atom)\n",
    "pert_atom_block = PERTATOMPARAM(pert_atoms)\n",
    "\n",
    "# Generate ptp file\n",
    "grom_system.ptp = Pertubation_topology(in_value = None)\n",
    "grom_system.ptp.PERTATOMPARAM = pert_atom_block\n",
    "grom_system.ptp.TITLE = TITLE(\"Automatic generated pertubation file. \")\n",
    "\n",
    "grom_system.ptp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "##Write out all generated files\n",
    "grom_system.rebase_files()\n",
    "grom_system.work_folder = project_dir\n",
    "\n",
    "##save Input System\n",
    "grom_system.save(project_dir+\"/initial_startSys.obj\")\n",
    "grom_system"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### RUN Emin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# PREPARE EMIN\n",
    "## IMPORT\n",
    "from pygromos.data.simulation_parameters_templates import template_emin_vac\n",
    "from pygromos.simulations.modules.preset_simulation_modules import emin\n",
    "\n",
    "step_name  = \"a_emin\" #also the dir_name, out prefix etc.\n",
    "grom_system = Gromos_System.load(project_dir+\"/initial_startSys.obj\")\n",
    "grom_system.imd = template_emin_vac #read template imd\n",
    "\n",
    "#Pertubation for molecules to sim params\n",
    "pert_block  =  PERTURBATION(NTG=1, NRDGL=0, RLAM=lam, DLAMT=0,\n",
    "                            ALPHC=0.5, ALPHLJ=0.5, NLAM=2, NSCALE=0)\n",
    "grom_system.imd.add_block(block=pert_block)\n",
    "\n",
    "#add Distance Res:\n",
    "disres_block = DISTANCERES(NTDIR=1, TAUDIR=1,)\n",
    "grom_system.imd.add_block(block=disres_block)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#EXECUTE EMIN\n",
    "emin_gromos_system = emin(in_gromos_system=grom_system, \n",
    "                          step_name=step_name, submission_system=subSystem())\n",
    "\n",
    "emin_gromos_system.save(project_dir+\"/emin_out.obj\")\n",
    "emin_gromos_system"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RUN Test SD EQ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pygromos.data.simulation_parameters_templates import template_sd\n",
    "from pygromos.simulations.modules.preset_simulation_modules import sd\n",
    "\n",
    "step_name  = \"b_vacuum_sd\"\n",
    "grom_system = emin_gromos_system\n",
    "grom_system.system_name = step_name\n",
    "grom_system.imd = template_sd\n",
    "\n",
    "#Pertubation\n",
    "pert_block  =  PERTURBATION(NTG=1, NRDGL=0, RLAM=lam, DLAMT=0,\n",
    "                            ALPHC=0.5, ALPHLJ=0.5, NLAM=2, NSCALE=0)\n",
    "grom_system.imd.add_block(block=pert_block)\n",
    "\n",
    "\n",
    "#add Distance Res:\n",
    "disres_block = DISTANCERES(NTDIR=1, TAUDIR=1,)\n",
    "grom_system.imd.add_block(block=disres_block)\n",
    "\n",
    "\n",
    "#write out trajs:\n",
    "write_traj = WRITETRAJ(NTWX=100, NTWE=100)\n",
    "grom_system.imd.add_block(block=write_traj)\n",
    "\n",
    "#further mods:\n",
    "grom_system.imd.CONSTRAINT.NTC = 3\n",
    "grom_system.imd.FORCE.BONDS = 0\n",
    "\n",
    "grom_system.imd.STEP.NSTLIM = 30000\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sd_gromos_system  = sd(in_gromos_system=grom_system, step_name=step_name,\n",
    "                              submission_system=subSystem(), equilibration_runs=1, simulation_runs=1)\n",
    "sd_gromos_system.save(project_dir+\"/sd_out_system.obj\")\n",
    "sd_gromos_system"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Further Analysis:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#final analysis dir:\n",
    "from pygromos.utils import bash\n",
    "\n",
    "out_ana = project_dir+\"/c_ana\"\n",
    "if(not os.path.exists(out_ana)):\n",
    "    bash.make_folder(out_ana)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Coordinate analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pygromos.files.coord import Cnf\n",
    "in_path=project_dir+\"/a_emin/analysis/data/a_emin.cnf\"\n",
    "cnf_file = Cnf(in_path)\n",
    "cnf_file.write_pdb(in_path.replace(\"cnf\", \"pdb\"))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cnf_file.view"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pygromos.files.trajectory.trc import Trc\n",
    "\n",
    "in_path=project_dir+\"/b_vacuum_sd/analysis/data/b_vacuum_sd.trc.h5\"\n",
    "\n",
    "trc = Trc(traj_path=in_path, in_cnf=cnf_file)\n",
    "trc.write(out_ana+\"/sd_traj.pdb\")#grom_system.cnf.path)\n",
    "trc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trc.view\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Energy analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pygromos.files.trajectory.tre import Tre\n",
    "\n",
    "in_path=project_dir+\"/b_vacuum_sd/analysis/data/b_vacuum_sd.tre.h5\"\n",
    "\n",
    "tre = Tre(input_value=in_path)\n",
    "tre\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Plot Potential Energies\n",
    "V_tot = np.array(list(map(lambda x: x[2], tre.database.totals)))\n",
    "step = len(tre.database.time)//10\n",
    "\n",
    "plt.plot(tre.database.time, V_tot)\n",
    "plt.xticks(np.round(list(tre.database.time[::step]),2))\n",
    "plt.xlabel(\"$t~[ps]$\")\n",
    "plt.ylabel(\"$V~[kJ]$\")\n",
    "plt.title(\"V total timeseries\")\n",
    "plt.savefig(out_ana+\"/potential_energy_timeseries.png\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Lambda Sampling\n",
    "\n",
    "### Setup again"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "from pygromos.files.gromos_system import Gromos_System\n",
    "from pygromos.simulations.hpc_queuing.submission_systems.local import LOCAL as subSystem\n",
    "from pygromos.utils import bash\n",
    "sys_name = \"M030_to_6KET\"\n",
    "\n",
    "sd_gromos_system.imd.WRITETRAJ.NTWG =  sd_gromos_system.imd.WRITETRAJ.NTWX = sd_gromos_system.imd.WRITETRAJ.NTWE =10\n",
    "sd_gromos_system.imd.STEP.NSTLIM = 100\n",
    "\n",
    "sd_gromos_system"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Submission"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pygromos.simulations.modules.ti_modules import TI_sampling\n",
    "\n",
    "step_name  = \"d_lambda_sampling\"\n",
    "sd_gromos_system.name = step_name\n",
    "\n",
    "\n",
    "TI_sampling(in_gromos_system = sd_gromos_system, step_name=step_name,\n",
    "            lambda_values= np.arange(0, 1.2, 0.2), \n",
    "            subSystem=subSystem(), n_productions=3, n_equilibrations = 1)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Free Energy Calculation:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#read in trgs files for free energy data:\n",
    "import glob\n",
    "from matplotlib import pyplot as plt\n",
    "from pygromos.files.trajectory.trg import Trg\n",
    "%matplotlib inline\n",
    "\n",
    "step_name  = \"d_lambda_sampling\"\n",
    "path = project_dir+\"/\"+step_name\n",
    "\n",
    "trg_paths = sorted(glob.glob(path+\"/*/analysis/data/*trg.h5\"), key=lambda x: float(\"\".join(x.split(\"_\")[-1].split(\".\")[1])))\n",
    "trg_paths.append(trg_paths.pop(0))\n",
    "trgs = [Trg(tre_path) for tre_path in trg_paths]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get lambda window values:\n",
    "lams = [float(trg.get_lambdas().iloc[0]) for trg in trgs]\n",
    "dhdl_means = [float(trg.get_totals()[\"dHdl\"].mean()) for trg in trgs]\n",
    "\n",
    "lams, dhdl_means"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TI - Integration Curve"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(lams, dhdl_means)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculate Integration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.integrate import simpson\n",
    "dF = simpson(y=dhdl_means, x=lams)\n",
    "dF"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "dd72df1003b4b968da7a05a6ba63af3773812ec2262e6b7ecb2dafd4fd33c317"
  },
  "kernelspec": {
   "display_name": "Python 3.9.12 ('pygromosDev')",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}