chicken-and-egg.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chicken and egg in molecular metadynamics\n",
    "\n",
    "### Create protein folding landmarks from scratch\n",
    "\n",
    "\n",
    "The whole workflow is implemented as a single notebook to allow batch run. It is structured to these major steps:\n",
    "\n",
    "1. generate random landmarks by twisting the input structrure backbone dihedrals and steep-descend minimization of the result to avoid side chain clash etc.\n",
    "1. compute low-dimensional embedding (isomap) of the landmarks to create collective variables\n",
    "1. train a neural networks to estimate the colvars from the structure, implement the resulting network as plumed input\n",
    "1. run both vanilla molecular dynamics and the metadynamics to compare the trajectories\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The input\n",
    "\n",
    "# pdbfile = \"2f21.pdb\"\n",
    "pdbfile=\"1L2Y.pdb\"\n",
    "\n",
    "# Or PDB download:\n",
    "# pdbid = '1l2y'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# we need all this fancy shit\n",
    "import anncolvar\n",
    "\n",
    "import os\n",
    "import shutil\n",
    "from contextlib import redirect_stdout\n",
    "import re\n",
    "\n",
    "import concurrent.futures\n",
    "\n",
    "import numpy as np\n",
    "import math\n",
    "from scipy.sparse import coo_matrix\n",
    "\n",
    "from pyDOE import lhs\n",
    "from scipy.sparse.csgraph import shortest_path\n",
    "from sklearn.manifold import MDS\n",
    "\n",
    "import PeptideBuilder as pb\n",
    "import Bio.PDB as pdb\n",
    "import Bio.SeqUtils as sequtil\n",
    "\n",
    "import mdtraj as md\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import nglview as nv\n",
    "\n",
    "from xvg import read_xvg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set the defaults \n",
    "# for batch processing, those can be set in config.py so that the notebook itself needn't be modified\n",
    "\n",
    "# number of steps to twist a dihedral when generating landmarks\n",
    "nsteps = 12\n",
    "\n",
    "# number of iterations of landmark generation; alltogether niter * nsteps random conformers are generated\n",
    "niter = 200\n",
    "\n",
    "# number of steps of the production MD run\n",
    "# md.mdp template sets 2 fs step\n",
    "mdsteps = 100000  # 200 ps just to test it works, production runs are expected to be extended sepearately\n",
    "\n",
    "# bounding box size to add for mininimization and \n",
    "minbox = 1.5\n",
    "mdbox = 1.5\n",
    "\n",
    "\n",
    "# count available cores -- check the output and set ncores to something else if necessary\n",
    "\n",
    "# TODO: capture PBS settings when running in batch\n",
    "\n",
    "if os.environ.get('OMP_NUM_THREADS') is None:\n",
    "    ncores = int(os.popen('./ncores.sh').read())\n",
    "    print('OMP_NUM_THREADS not set, using all (%d) available cores' % ncores)\n",
    "else:\n",
    "    ncores = int(os.environ.get('OMP_NUM_THREADS'))\n",
    "    print('Using OMP_NUM_THREADS = %d cores' % ncores)\n",
    "\n",
    "\n",
    "# XXX: poor man approach, expected to be tuned in config.py\n",
    "ntomp = 4\n",
    "ntmpi = ncores // ntomp\n",
    "try:\n",
    "    exec(open('config.py').read())\n",
    "    !cat config.py\n",
    "except FileNotFoundError:\n",
    "    print('config.py not found, hope it\\'s OK')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def scale(cores):\n",
    "    !kubectl scale deployment.apps/chicken-and-egg{os.environ.get('K8S_LABEL')}-placeholder --replicas={cores}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load a PDB file, store it locally, create a matching workdir\n",
    "\n",
    "try:\n",
    "    workdir\n",
    "    raise Exception(\"This cell should be run only once (workdir = %s)\" % workdir)\n",
    "except NameError:\n",
    "    pass\n",
    "\n",
    "basedir=os.getcwd()\n",
    "\n",
    "\n",
    "# load the PDB file, create workdir, set global variables to be used later\n",
    "do_load = False\n",
    "try:\n",
    "    pdbid\n",
    "    do_load = True\n",
    "    pdbfile = pdbid + '.pdb'\n",
    "\n",
    "except NameError:\n",
    "    pdbid = os.path.splitext(os.path.basename(pdbfile))[0]\n",
    "    \n",
    "workdir=os.path.join(basedir,pdbid)\n",
    "gmx = f\"{basedir}/gmx-k8s -w {pdbid}\"\n",
    "minim = f\"{basedir}/minim-k8s -w {pdbid}\"\n",
    " \n",
    "if not os.path.exists(workdir):\n",
    "    os.mkdir(workdir)\n",
    "\n",
    "os.chdir(workdir)\n",
    "\n",
    "\n",
    "if do_load:\n",
    "    pdbl = pdb.PDBList()\n",
    "    pdbl.retrieve_pdb_file(pdbid,file_format='pdb')\n",
    "    shutil.move(pdbid[1:3] + \"/pdb\" + pdbid + \".ent\", pdbid + \".pdb\")\n",
    "else:\n",
    "    shutil.copy(os.path.join(basedir,pdbfile),os.path.join(workdir,pdbfile))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1 Generate landmarks by random twisting PDB structure\n",
    "\n",
    "### 1.1 Initial preprocessing and visual check\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# preprocess the file with gromacs to get consistent atom naming and numbering\n",
    "!{gmx} pdb2gmx -f {pdbfile} -o {pdbid}-new.pdb -water tip3p -ff amber94 -ignh\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mv -f {pdbid}-new.pdb {pdbfile}\n",
    "\n",
    "m = md.load(pdbfile)\n",
    "heavy_idx = m[0].top.select(\"element != H\")\n",
    "heavy_atoms = len(heavy_idx)\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# inspect the loaded file\n",
    "# it must be a sane structure, no missing heavy atoms and/or hydrogens etc., \n",
    "# suitable as the starting point of usual MD protocol\n",
    "os.chdir(basedir)\n",
    "v = nv.NGLWidget()\n",
    "v.add_component(os.path.join(workdir,pdbfile))\n",
    "v.clear()\n",
    "v.add_representation('cartoon', selection='all')\n",
    "\n",
    "os.chdir(workdir)\n",
    "v"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2 Generate randomly twisted conformations\n",
    "\n",
    "$\\phi$ and $\\psi$ backbone dihedral angles of all but first and last residue of the loaded structure are twisted randomly.\n",
    "\n",
    "Systematic approach (e.g. 30 degree sampling of all angles) would yield too many conformations.\n",
    "Instead we use random latin hypercube sampling to get uniform coverage of all values of all angles.\n",
    "\n",
    "Empirically, running 10 times no. of residues (`niter` parameter bellow) seems to be sufficient to cover the whole conformational space while keeping number of landmarks still reasonable.\n",
    "\n",
    "Expect approx. 1 s per 300 residues. It is worth to inspect some of the outputs visually (the following cell)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = pdb.PDBParser()\n",
    "instruct = p.get_structure('in',pdbfile)\n",
    "\n",
    "# XXX: assuming one model and one chain, the method would be rather weird for more\n",
    "\n",
    "resl = list(map(lambda r: sequtil.seq1(r.get_resname()),instruct.get_residues()))\n",
    "nres = len(resl)\n",
    "\n",
    "out='conf%d.pdb'\n",
    "\n",
    "# make it really reproducible\n",
    "np.random.seed(123456789)\n",
    "    \n",
    "def random_twist(itrn):    \n",
    "    np.random.seed(itrn + 123456789)\n",
    "\n",
    "    phi = lhs(nres - 2, nsteps)\n",
    "    psi = lhs(nres - 2, nsteps)\n",
    "    outf = pdb.PDBIO()\n",
    "\n",
    "    for s in range(nsteps):\n",
    "        first = pb.Geometry.geometry(resl[0])\n",
    "        struct = pb.initialize_res(first)\n",
    "        \n",
    "        for r in range(1,nres-1):\n",
    "            if resl[r] == 'P':\n",
    "                pb.add_residue(struct,resl[r])\n",
    "            else:\n",
    "                pb.add_residue(struct,resl[r],phi[s][r-1]*360,psi[2][r-1]*360)\n",
    "                \n",
    "        pb.add_residue(struct,resl[nres-1])\n",
    "            \n",
    "        fn = out % (itrn * nsteps + s + 1)\n",
    "        outf.set_structure(struct)\n",
    "        outf.save(fn)\n",
    "        \n",
    "scale(ncores)\n",
    "# XXX: better with hyperthreading but we don't want to eat up 2x cores when running in batch mode\n",
    "with concurrent.futures.ProcessPoolExecutor(max_workers=ncores) as executor:\n",
    "    for _ in executor.map(random_twist,range(niter)):\n",
    "        pass\n",
    "\n",
    "scale(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tr = md.load([ \"conf%d.pdb\" % i for i in range(1,nsteps*niter+1)])\n",
    "idx=tr[0].top.select(\"name CA\")\n",
    "tr.superpose(tr[0],atom_indices=idx)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v=nv.show_mdtraj(tr)\n",
    "v.clear()\n",
    "v.add_representation(\"licorice\")\n",
    "v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tr.save('premin.xtc')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.3 Minimize the generated structures\n",
    "\n",
    "Run Gromacs steepest descend energy minimization in vacuo on all the generated structures. This is sufficient to fix colliding sidechains etc. while not changing the backbone dihedrals, hence preserving the conformational space coverage.\n",
    "\n",
    "Expect approx. 25 structures per minute per core in case of small protein like trpcage (1L2Y)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the most likely minimization parameters to change; rest is in the template file\n",
    "minim_mdp = '''\n",
    "emtol       = 500.0        ; Stop minimization when the maximum force is lower (kJ/mol/nm)\n",
    "emstep      = 0.05          ; Minimization step size\n",
    "nsteps      = 500         ; Maximum number of (minimization) steps to perform\n",
    "'''\n",
    "\n",
    "template = os.path.join(basedir,'minim.mdp.template')\n",
    "\n",
    "!cp {template} minim.mdp\n",
    "f=open('minim.mdp','a')\n",
    "f.write(minim_mdp)\n",
    "f.close()\n",
    "\n",
    "!bash {minim} -n {ncores} -b {minbox}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# filter the results to the reasonable ones only\n",
    "\n",
    "conflist = []\n",
    "frames = []\n",
    "energies = []\n",
    "maxenergy = 1e8\n",
    "\n",
    "for i in range(1,nsteps*niter+1):\n",
    "    try:\n",
    "        with open('conf%d.minen' % i) as ef:\n",
    "            l = ef.readline()\n",
    "            _,energy = l.split()\n",
    "            energy = float(energy)\n",
    "    except FileNotFoundError:\n",
    "        print(i, \"not found, something went wrong\")\n",
    "        continue\n",
    "        \n",
    "    fn = \"conf%d-min.gro\" % i\n",
    "    if os.path.isfile(fn):\n",
    "        one = md.load(fn)\n",
    "        heavy_idx = one[0].top.select(\"element != H\")\n",
    "        one.atom_slice(heavy_idx,inplace=True)\n",
    "        if one.n_atoms == heavy_atoms:\n",
    "            frames.append(one)\n",
    "        else:\n",
    "            print(\"%d number of heavy atoms (%d) should be %d, ignoring\" % (i,one.n_atoms, heavy_atoms))\n",
    "            continue\n",
    "    else:\n",
    "        print(fn, \"not found, ignoring\")\n",
    "        continue\n",
    "        \n",
    "    if energy < maxenergy:\n",
    "        conflist.append(i)\n",
    "        energies.append(energy)\n",
    "    else:\n",
    "        print(i,\"energy too high:\", energy, 'ignoring')\n",
    "        \n",
    "print('remaining conformers', len(conflist))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.4 Inspect the results\n",
    "\n",
    "Minimized structures are merged into virtual trajectory and displayed as animation.\n",
    "\n",
    "Histograms of their radius of gyration and energies (following cells) gives some evidence on conformational space coverage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tr=md.join(frames)\n",
    "tr[0].center_coordinates()\n",
    "idx=tr[0].top.select(\"name CA\")\n",
    "tr.superpose(tr[0],atom_indices=idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v=nv.show_mdtraj(tr,gui=False)\n",
    "v.clear()\n",
    "v.add_representation(\"licorice\")\n",
    "v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rgs=md.compute_rg(tr)\n",
    "\n",
    "plt.rcParams.update({'font.size': 17})\n",
    "plt.figure(figsize=(10,5))\n",
    "plt.hist(rgs,200)\n",
    "plt.xlabel('Radius of gyration (nm)')\n",
    "plt.ylabel('# of samples')\n",
    "plt.savefig('minim-rg.pdf')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10,5))\n",
    "plt.hist(energies,200)\n",
    "plt.yscale('log')\n",
    "plt.ylabel('# of samples')\n",
    "plt.xlabel('Energy (kJ/mol)')\n",
    "plt.xticks([0,5000,10000,15000,20000])\n",
    "plt.savefig('minim-energ.pdf')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save superposed landmarks as Gromacs trajectory\n",
    "tr.save_pdb('landmarks.pdb')\n",
    "tr.save_xtc('landmarks.xtc')\n",
    "tr[0].save_pdb('landmark1.pdb')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# cluster the minimized landmarks\n",
    "# the purpose is checking density of sampling the conformational space\n",
    "# the number of clusters should roughly match the number of landmarks, too few (dozens) clusters \n",
    "# indicates the minimization went too far\n",
    "\n",
    "!{gmx} -i 2,1 cluster -s landmark1.pdb -f landmarks.xtc -o clusters.xpm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Compute isomap projection of the landmarks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# number of nearest neighbours to consider (aka _k_)\n",
    "neighs = 5\n",
    "# targed no. of dimensions \n",
    "dims = 2\n",
    "\n",
    "try:\n",
    "    tr\n",
    "except NameError:\n",
    "    tr = md.load('landmarks.pdb')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scale(ncores)\n",
    "# nconf = niter * nsteps\n",
    "nconf = len(tr)\n",
    "\n",
    "# compute all-to-all RMSD and select _k_ closest neighbours\n",
    "row=[]\n",
    "col=[]\n",
    "dat=[]\n",
    "\n",
    "for i in range(nconf):\n",
    "    d = md.rmsd(tr,tr,frame=i)\n",
    "    d[range(i+1)] = np.inf\n",
    "    for _ in range(neighs):\n",
    "        j = np.argmin(d)\n",
    "        if d[j] < np.inf:\n",
    "            row.append(i)\n",
    "            col.append(j)\n",
    "            dat.append(d[j])\n",
    "            row.append(j)\n",
    "            col.append(i)\n",
    "            dat.append(d[j])\n",
    "            d[j] = np.inf\n",
    "\n",
    "# store results in sparse matrix\n",
    "dist = coo_matrix((dat,(row,col)),shape=(nconf,nconf)) \n",
    "scale(0)\n",
    "\n",
    "# check sanity\n",
    "print(\"conformations (original dimensions): \", nconf)\n",
    "print(\"non-zero distances: \", dist.getnnz())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# isomap itself: compute shortest paths in the k-neighbours graph, \n",
    "# and multi-dimensional scaling on the resulting all-to-all distances\n",
    "scale(ncores)\n",
    "sp = shortest_path(dist,directed=False)\n",
    "mds = MDS(n_components=dims,dissimilarity='precomputed',n_jobs=ncores)\n",
    "emb = mds.fit_transform(sp)\n",
    "scale(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(12,8))\n",
    "plt.scatter(*emb.transpose(),marker='.')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### XXX: assumes negative min and positive max\n",
    "embmin=np.min(emb,axis=0)*3.\n",
    "embmax=np.max(emb,axis=0)*3."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save collective variables\n",
    "np.savetxt('colvar.txt',emb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3 Prepare metadynamics\n",
    "\n",
    "### 3.1 Train the neural net\n",
    "\n",
    "Create artificial neural network and train it to produce the above isomap embedding from superposed heavy atom coordinates. The ANN is encoded in `plumed.dat` to be used by metadynamic run later.\n",
    "\n",
    "Technically, this is done for both coordinates of the embedding independently, the resulting `plumed.dat` files are merged.\n",
    "\n",
    "Uses [Anncolvar](https://github.com/spiwokv/anncolvar)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    emb\n",
    "except NameError:\n",
    "    emb=np.loadtxt('colvar.txt')\n",
    "    embmin=np.min(emb,axis=0)*3.\n",
    "    embmax=np.max(emb,axis=0)*3."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run anncolvar twice, for each isomap coordinate independently\n",
    "epochs = 500\n",
    "\n",
    "# anncolvar defaults\n",
    "nlayers = 1\n",
    "layers = [32, 0, 0]\n",
    "# actfun = ['sigmoid','linear','linear']\n",
    "# XXX: nothing else supported\n",
    "actfun = ['tanh']*3\n",
    "\n",
    "optim = 'adam'\n",
    "loss = 'mean_squared_error'\n",
    "batch = 256\n",
    "\n",
    "first = md.load('landmark1.pdb')\n",
    "rg = md.compute_rg(first)\n",
    "\n",
    "# XXX magic -- seems to be safe\n",
    "shift = rg[0] * 5\n",
    "box = shift * 2\n",
    "\n",
    "one = md.load(pdbfile)\n",
    "heavy_idx = one[0].top.select(\"element != H\")\n",
    "\n",
    "!{gmx} editconf -f landmark1.pdb -o landmark1-box.pdb -translate {shift} {shift} {shift} -box {box} {box} {box} -c >editconf.log 2>&1\n",
    "\n",
    "with open('landmark1-box.pdb') as boxf, open('reference.pdb','w') as ref:\n",
    "    lines = boxf.readlines()\n",
    "    for l in lines:\n",
    "        if l[:4] == 'ATOM':\n",
    "            newi = heavy_idx[int(l[4:11])-1]+1\n",
    "            ref.write('ATOM%7d' % newi)\n",
    "            ref.write(l[11:])\n",
    "        else:\n",
    "            ref.write(l)\n",
    "\n",
    "scale(ncores)\n",
    "# XXX: too much stdout \n",
    "for col in [1,2]:\n",
    "    with open(\"anncolvar-%d.log\" % col,\"w\") as log, redirect_stdout(log):\n",
    "        anncolvar.anncollectivevariable('landmarks.xtc','reference.pdb','colvar.txt',col,\n",
    "                                    box,box,box,.1,0,0,\n",
    "                                    nlayers,*layers,\n",
    "                                    *actfun,\n",
    "                                    optim,loss,epochs,batch,\n",
    "                                    '','',fannfile='cv%d-plumed.dat' % col,\n",
    "                                    plumedfile='old-cv%d-plumed.dat' % col)\n",
    "            \n",
    "scale(0)\n",
    "!tail editconf.log anncolvar-[12].log"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# merge plumed[12].dat from the previous cell\n",
    "\n",
    "onlyone = ['WHOLEMOLECULES', 'FIT_TO_TEMPLATE']\n",
    "\n",
    "# XXX entirely\n",
    "with open('plumed-ann.dat','w') as fout:\n",
    "    fout.write('RESTART\\n')\n",
    "    for col in [1,2]:\n",
    "        with open('cv%d-plumed.dat' % col) as fin:\n",
    "            for line in fin:\n",
    "                '''not necessary, reference.pdb is already renumbered\n",
    "                w = line.split()\n",
    "                if w[1] == 'POSITION':\n",
    "                    a = w[2].split('=')\n",
    "                    line = w[0] + (' POSITION ATOM=%d ' % (heavy_idx[int(a[1])-1]+1)) + ' '.join(w[3:]) + '\\n'\n",
    "                '''\n",
    "                shutup = False\n",
    "                for o in onlyone:\n",
    "                    if line[:len(o)] == o:\n",
    "                        if col != 1:\n",
    "                            shutup = True\n",
    "                if shutup:\n",
    "                    continue\n",
    "                    \n",
    "                if line[:5] == 'PRINT':\n",
    "                    continue\n",
    "                            \n",
    "                if re.match('[pl][0-9_rxyz]',line) or line[:4] == 'ARG=':\n",
    "                    line = re.sub('[pl][0-9_rxyz]','cv%d_\\g<0>' % col,line)\n",
    "                    fout.write(line)\n",
    "                elif line == 'LABEL=ann\\n':\n",
    "                    fout.write('LABEL=ann_cv%d\\n' % col)\n",
    "                else:\n",
    "                    fout.write(line)\n",
    "                \n",
    "  \n",
    "    # XXX: hardcoded\n",
    "    fout.write('PRINT ARG=ann_cv1.node-0,ann_cv2.node-0 STRIDE=100 FILE=COLVAR-ann\\n')\n",
    "    fout.write('METAD ARG=ann_cv1.node-0,ann_cv2.node-0 SIGMA=0.1,0.1 HEIGHT=1.0 FILE=HILLS-ann PACE=1000 BIASFACTOR=15 TEMP=300 LABEL=restraint')\n",
    "    fout.write(' GRID_MIN=%f,%f GRID_MAX=%f,%f\\n' % (*embmin,*embmax)) \n",
    "    \n",
    "# XXX: hack plumed.dat, FIT_TO_TEMPLATE TYPE=OPTIMAL is broken\n",
    "!sed '/^FIT_TO_TEMPLATE/s/TYPE=OPTIMAL/TYPE=SIMPLE/' plumed-ann.dat >plumed-ann.dat.$$ && mv plumed-ann.dat.$$ plumed-ann.dat"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2 Prepare PCV with the same landmarks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cvs = np.loadtxt('colvar.txt')\n",
    "i=0\n",
    "start=True\n",
    "with open('landmarks.pdb') as inp, open('landmarks-pcv.pdb','w') as out:\n",
    "    for line in inp:\n",
    "        if line[:5] == 'MODEL':\n",
    "            if i > 0: \n",
    "                out.write('END\\n')\n",
    "                \n",
    "            out.write(\"REMARK X=%f Y=%f\\n\" % tuple(cvs[i]))\n",
    "            i += 1\n",
    "        if line[:4] == 'ATOM':\n",
    "            out.write(line)\n",
    "    out.write('END\\n')\n",
    "\n",
    "!echo RESTART >plumed-pcv.dat\n",
    "!grep WHOLEMOLECULES plumed-ann.dat >>plumed-pcv.dat\n",
    "with open('plumed-pcv.dat','a') as plmd:\n",
    "    plmd.write(\"p1: PROPERTYMAP REFERENCE=landmarks-pcv.pdb PROPERTY=X,Y LAMBDA=50.0 NEIGH_SIZE=50 NEIGH_STRIDE=50 EPSILON=0.01\\n\")\n",
    "    plmd.write('METAD ARG=p1.X,p1.Y SIGMA=0.1,0.1 HEIGHT=1.0 FILE=HILLS-pcv PACE=1000 BIASFACTOR=15 TEMP=300 LABEL=restraint')\n",
    "    plmd.write(' GRID_MIN=%f,%f GRID_MAX=%f,%f\\n' % (*embmin,*embmax)) \n",
    "    plmd.write('PRINT ARG=p1.X,p1.Y,p1.zzz,restraint.bias STRIDE=100 FILE=COLVAR-pcv FMT=%8.4f\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Run MD\n",
    "\n",
    "Run quite standard molecular dynamics protocol, adapted from [Lysosome tutorial](http://www.mdtutorials.com/gmx/lysozyme/index.html), i.e. solvate, add counterions, minimize, equilibrate, and run production.\n",
    "\n",
    "Preparation phases are common, then we run vanilla and metadynamic simulations to compare the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2 fs steps\n",
    "mdsteps = 500*1000*200"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.1 Prepare, minimize, and equilibrate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# elementary preparation\n",
    "\n",
    "# XXX hardcoded defaults for the time being, replace with template eventually\n",
    "\n",
    "os.chdir(basedir)\n",
    "!cp ions.mdp minim-sol.mdp {workdir}\n",
    "os.chdir(workdir)\n",
    "\n",
    "!{gmx} pdb2gmx -f {pdbfile} -o {pdbid}.gro -water tip3p -ff amber94 -ignh -p {pdbid}.top && \\\n",
    "{gmx} editconf -f {pdbid}.gro -o {pdbid}-box.gro -c -d {mdbox} -bt dodecahedron && \\\n",
    "{gmx} solvate -cp {pdbid}-box.gro -cs spc216.gro -o {pdbid}-solv.gro -p {pdbid}.top && \\\n",
    "{gmx} grompp -f ions.mdp -c {pdbid}-solv.gro -p {pdbid}.top -o ions.tpr && \\\n",
    "{gmx} -i 13 genion -s ions.tpr -o {pdbid}-ions.gro -p {pdbid}.top -pname NA -nname CL -neutral"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# minimize with steepest descend\n",
    "!{gmx} grompp -f minim-sol.mdp -c {pdbid}-ions.gro -p {pdbid}.top -o em.tpr &&\\\n",
    "unset OMP_NUM_THREADS && {gmx} -n {ntmpi} mdrun -v -deffnm em -ntomp {ntomp} -pin on &&\\\n",
    "{gmx} -i 10 energy -f em.edr -o em.xvg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x,y=read_xvg(os.path.join(workdir,'em.xvg'))\n",
    "\n",
    "plt.figure(figsize=(15,5))\n",
    "plt.plot(x,y)\n",
    "plt.grid()\n",
    "plt.xlabel('step')\n",
    "plt.ylabel('potential (kJ/mol)')\n",
    "plt.title('Energy minimization')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# isothermal - isochoric equilibration\n",
    "!cp {basedir}/nvt.mdp .\n",
    "\n",
    "!{gmx} grompp -f nvt.mdp -c em.gro -r em.gro -p {pdbid}.top -o nvt.tpr && \\\n",
    "unset OMP_NUM_THREADS && {gmx} -n {ntmpi} mdrun -ntomp {ntomp}  -pin on -deffnm nvt && \\\n",
    "{gmx} -i 16 energy -f nvt.edr -o temp.xvg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x,y=read_xvg(os.path.join(workdir,'temp.xvg'))\n",
    "\n",
    "plt.figure(figsize=(15,5))\n",
    "plt.plot(x,y)\n",
    "plt.grid()\n",
    "plt.xlabel('time (ps)')\n",
    "plt.ylabel('temperature (K)')\n",
    "plt.title('isothermal-isochoric equilibration')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# isothermal - isobaric equilibration\n",
    "!cp {basedir}/npt.mdp .\n",
    "\n",
    "!{gmx} grompp -f npt.mdp -c nvt.gro -r nvt.gro -t nvt.cpt -p {pdbid}.top -o npt.tpr && \\\n",
    "unset OMP_NUM_THREADS && {gmx} -n {ntmpi} mdrun -ntomp {ntomp} -pin on -deffnm npt && \\\n",
    "{gmx} -i 18 energy -f npt.edr -o press.xvg && \\\n",
    "{gmx} -i 24 energy -f npt.edr -o dens.xvg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xp,yp=read_xvg(os.path.join(workdir,'press.xvg'))\n",
    "xd,yd=read_xvg(os.path.join(workdir,'dens.xvg'))\n",
    "\n",
    "plt.figure(figsize=(15,8))\n",
    "plt.subplot(211)\n",
    "plt.plot(xp,yp)\n",
    "plt.title('isothermal-isobaric equilibration')\n",
    "plt.grid()\n",
    "#plt.xlabel('time (ps)')\n",
    "plt.ylabel(\"pressure (bar)\")\n",
    "\n",
    "\n",
    "plt.subplot(212)\n",
    "plt.xlabel('time (ps)')\n",
    "plt.ylabel('density (kg/m3)')\n",
    "plt.grid()\n",
    "plt.plot(xd,yd)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.2 Run vanilla MD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!cp {basedir}/md.mdp.template md.mdp\n",
    "with open('md.mdp','a') as mdp:\n",
    "    mdp.write(\"nsteps = %d\\n\" % mdsteps)\n",
    "\n",
    "!{gmx} grompp -f md.mdp -c npt.gro -t npt.cpt -p {pdbid}.top -o md-vanilla.tpr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!unset OMP_NUM_THREADS && {gmx} -n {ntmpi} mdrun -ntomp {ntomp} -pin on -deffnm md-vanilla -cpi md-vanilla.cpt -cpt 15"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.3 Run Anncolvar metadynamics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cp {basedir}/md.mdp.template md.mdp\n",
    "\n",
    "# mdsteps=1000000\n",
    "with open('md.mdp','a') as mdp:\n",
    "    mdp.write(\"nsteps = %d\\n\" % mdsteps)\n",
    "\n",
    "!{gmx} grompp -f md.mdp -c npt.gro -t npt.cpt -p {pdbid}.top -o md-ann.tpr "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!unset OMP_NUM_THREADS && {gmx} -n {ntmpi} mdrun -ntomp {ntomp} -pin on -deffnm md-ann -plumed plumed-ann.dat -cpi md-ann.cpt -cpt 15"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.4 Run PCV metadynamics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!{gmx} grompp -f md.mdp -c npt.gro -t npt.cpt -p {pdbid}.top -o md-pcv.tpr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!unset OMP_NUM_THREADS && {gmx} -n {ntmpi} mdrun -ntomp {ntomp} -pin on -deffnm md-pcv -plumed plumed-pcv.dat -cpi md-psv.cpt -cpt 15"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.5 Extend trajectories arbitrarily\n",
    "\n",
    "You can skip this section entirely, and go to 5. to examine results of the short trajectories if it is the purpose. But typically longer trajectories must be computed.\n",
    "\n",
    "This is commented out now, we run computation in K8s on strong enough resources.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Backup short trajectories first\n",
    "# for suffix in ['vanilla','ann','pcv']:\n",
    "#     !cp md-{suffix}.tpr md-{suffix}-short.tpr\n",
    "#     !cp md-{suffix}.xtc md-{suffix}-short.xtc\n",
    "#     !cp md-{suffix}.edr md-{suffix}-short.edr\n",
    "#     !cp COLVAR-{suffix} COLVAR-{suffix}-short\n",
    "#     !cp HILLS-{suffix} HILLS-{suffix}-short"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# wanted=100000000 # in steps\n",
    "\n",
    "# !{gmx} convert-tpr -s md1.tpr -o md1-long.tpr -nsteps {wanted} && mv md1-long.tpr md1.tpr\n",
    "# !{gmx} convert-tpr -s md2.tpr -o md2-long.tpr -nsteps {wanted} && mv md2-long.tpr md2.tpr\n",
    "# !{gmx} convert-tpr -s md3.tpr -o md3-long.tpr -nsteps {wanted} && mv md3-long.tpr md3.tpr\n",
    "# !echo RESTART >plumed-ann-restart.dat && cat plumed-ann.dat >>plumed-ann-restart.dat\n",
    "# !echo RESTART >plumed-pcv-restart.dat && cat plumed-pcv.dat >>plumed-pcv-restart.dat"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Transfer the workdir to a more powerful node, and run re-run the MD computation arbitrarily to the desired trajectory lenght.\n",
    "\n",
    "This involves grabbing files md[123].\\*, plumed-restart.dat, COLVAR, HILLS, and reference.pdb from the working directory, and running commands in the following cells on the powerful node (with appropriate {ntmpi} and {ntomp} settings). The script qsub-extend.sh from the same repo can be used to submit to PBS (after editing to match your environment).\n",
    "\n",
    "Alternatively, just uncomment the following cells and run them. It will take quite long time to finish. If the computation gets killed, just run it again, it restarts from a checkpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !unset OMP_NUM_THREADS && {gmx} -n {ntmpi} mdrun -ntomp {ntomp} -pin on -deffnm md1 -cpi md1.cpt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !unset OMP_NUM_THREADS && {gmx} -n {ntmpi} mdrun -ntomp {ntomp} -pin on -deffnm md2 -cpi md2.cpt -plumed plumed-ann-restart"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !unset OMP_NUM_THREADS && {gmx} -n {ntmpi} mdrun -ntomp {ntomp} -pin on -deffnm md2 -cpi md2.cpt -plumed plumed-pcv-restart"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 5 Analyze results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pbc_and_fit(base):\n",
    "    xtc = base + \".xtc\"\n",
    "    pbc = base + \"-pbc.xtc\"\n",
    "    !{gmx} -i 1 trjconv -f {xtc} -s npt.gro -pbc nojump -o {pbc} 2>&1 | tail -100\n",
    "    tr = md.load_xtc(pbc,top=pdbid+'.gro')\n",
    "    idx=tr[0].top.select(\"name CA\")\n",
    "    tr.superpose(tr[0],atom_indices=idx)\n",
    "    return tr\n",
    "\n",
    "def plot_rmsd_rgyr(tr):\n",
    "    rmsd = md.rmsd(tr,tr)\n",
    "    rg = md.compute_rg(tr)\n",
    "    plt.figure(figsize=(15,8))\n",
    "    plt.subplot(211)\n",
    "    plt.plot(rmsd)\n",
    "    plt.grid()\n",
    "    plt.ylabel('RMSD wrt. frame 0 (nm)')\n",
    "    plt.subplot(212)\n",
    "    plt.plot(rg)\n",
    "    plt.grid()\n",
    "    plt.ylabel('Radius of gyration (nm)')\n",
    "    plt.xlabel('time (10 ps steps)')\n",
    "    plt.show()\n",
    "    \n",
    "def plot_stability(tr,part=.7):\n",
    "    idx=tr.top.select(\"protein and element != H\")\n",
    "    heavy = tr.atom_slice(idx,inplace=False)\n",
    "\n",
    "    xyz = np.reshape(heavy.xyz,(heavy.xyz.shape[0],heavy.xyz.shape[1]*3))\n",
    "    xyz_avg = np.average(xyz,axis=0)\n",
    "    xyz -= xyz_avg\n",
    "\n",
    "    cor = np.matmul(np.transpose(xyz),xyz)\n",
    "    cor /= xyz.shape[0]\n",
    "\n",
    "    full = np.abs(np.sort(np.linalg.eigvalsh(cor)))\n",
    "\n",
    "    sum = np.sum(full)\n",
    "    full /= sum\n",
    "\n",
    "    num = len(full)\n",
    "    full = np.flip(full[-num:])\n",
    "    full = np.cumsum(full)\n",
    "\n",
    "    part_heavy = heavy[:int(len(heavy) * part)]\n",
    "    xyz = np.reshape(part_heavy.xyz,(part_heavy.xyz.shape[0],part_heavy.xyz.shape[1]*3))\n",
    "    xyz_avg = np.average(xyz,axis=0)\n",
    "    xyz -= xyz_avg\n",
    "\n",
    "    cor = np.matmul(np.transpose(xyz),xyz)\n",
    "    cor /= xyz.shape[0]\n",
    "\n",
    "    start = np.abs(np.sort(np.linalg.eigvalsh(cor)))\n",
    "\n",
    "    sum = np.sum(start)\n",
    "    start /= sum\n",
    "\n",
    "    start = np.flip(start[-num:])\n",
    "    start = np.cumsum(start)\n",
    "\n",
    "    show = np.count_nonzero(full < .95)\n",
    "\n",
    "    plt.figure(figsize=((15,8)))\n",
    "    plt.plot(range(show),full[:show])\n",
    "    plt.plot(range(show),start[:show])\n",
    "    plt.grid()\n",
    "    plt.legend(['full','initial %d%%' % int(part*100)])\n",
    "    plt.ylabel(\"cummulative eigenvalues\")\n",
    "    plt.show()\n",
    "    \n",
    "    \n",
    "def plot_energy(base):\n",
    "    edr = base + \".edr\"\n",
    "#    !echo 11 | {gmx} energy -f {edr} 2>&1 | tail\n",
    "    !{gmx} -i 11 energy -f {edr} 2>&1 | tail\n",
    "    energ = read_xvg('energy.xvg')\n",
    "    \n",
    "    plt.figure(figsize=(15,8))\n",
    "    plt.plot(*energ)\n",
    "    plt.grid()\n",
    "    plt.ylabel('Potential energy (kJ/mol)')\n",
    "    plt.xlabel('time (ps)')\n",
    "    plt.show()\n",
    "    return energ\n",
    "    \n",
    "def plot_colvar(suffix,stride=20):\n",
    "    lms=np.loadtxt('colvar.txt').T\n",
    "    cv = np.transpose(np.loadtxt('COLVAR-' + suffix)[::stride])\n",
    "    cv2 = cv[1:3]\n",
    "    plt.figure(figsize=(12,12))\n",
    "    plt.scatter(*cv2,c=range(cv2.shape[1]),marker='.',cmap=plt.get_cmap('rainbow'))\n",
    "    plt.colorbar()\n",
    "    plt.scatter(*lms,c='black',marker='+')\n",
    "    plt.show()\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5.1 Unbiased MD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base1 = \"md-vanilla\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tr1=pbc_and_fit(base1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "v = nv.show_mdtraj(tr1,gui=False)\n",
    "v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_rmsd_rgyr(tr1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "energ1=plot_energy(base1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_stability(tr1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5.2 Biased MD with Anncolvar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base2 = \"md-ann\"\n",
    "#base2 = \"md-fann\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tr2 = pbc_and_fit(base2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v = nv.show_mdtraj(tr2,gui=False)\n",
    "v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_rmsd_rgyr(tr2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "energ2 = plot_energy(base2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_stability(tr2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_colvar('ann')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5.3 Biased MD with PCV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base3='md-pcv'\n",
    "tr3 = pbc_and_fit(base3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v = nv.show_mdtraj(tr3,gui=False)\n",
    "v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_rmsd_rgyr(tr3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "energ3=plot_energy(base3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_stability(tr3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_colvar('pcv')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5.4 Progress alltogether"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "rmsd1 = md.rmsd(tr1,tr1)\n",
    "rg1 = md.compute_rg(tr1)\n",
    "rmsd2 = md.rmsd(tr2,tr2)\n",
    "rg2 = md.compute_rg(tr2)\n",
    "rmsd3 = md.rmsd(tr3,tr3)\n",
    "rg3 = md.compute_rg(tr3)\n",
    "\n",
    "# XXX: same lenth expected\n",
    "l = len(rmsd1)\n",
    "l8 = l // 8;\n",
    "ticks = np.arange(0,l,l8)\n",
    "labels = ticks / 100\n",
    "\n",
    "plt.rcParams.update({'font.size': 14})\n",
    "_,ax = plt.subplots(3,1,figsize=(15,8))\n",
    "#plt.subplot(311)\n",
    "ax[0].plot(rmsd1)\n",
    "ax[0].plot(rmsd2)\n",
    "ax[0].plot(rmsd3)\n",
    "ax[0].grid()\n",
    "ax[0].set_ylabel('RMSD (nm)')\n",
    "ax[0].set_xticks(ticks)\n",
    "ax[0].set_xticklabels(labels)\n",
    "ax[0].legend(['unbiased','ANN','PCV'])\n",
    "#plt.subplot(312)\n",
    "ax[1].plot(rg1)\n",
    "ax[1].plot(rg2)\n",
    "ax[1].plot(rg3)\n",
    "ax[1].grid()\n",
    "ax[1].set_ylabel('R. gyr. (nm)')\n",
    "ax[1].set_xticks(ticks)\n",
    "ax[1].set_xticklabels(labels)\n",
    "#plt.subplot(313)\n",
    "ax[2].plot(energ1[1])\n",
    "ax[2].plot(energ2[1])\n",
    "ax[2].plot(energ3[1])\n",
    "ax[2].grid()\n",
    "ax[2].set_ylabel('Epot (kJ/mol)')\n",
    "ax[2].set_xticks(ticks)\n",
    "ax[2].set_xticklabels(labels)\n",
    "ax[2].set_xlabel('time (ns)')\n",
    "plt.savefig('graphs.pdf')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5.5 Crosscheck of both CV calculations\n",
    "\n",
    "Use plumed driver to calculate CVs with PCV CV definition on ANN trajectory and vice versa.\n",
    "ANN tends to explore wider regions, beyond the space covered by landmarks (should be visible on the maps above). \n",
    "The reason is that PCV approach zero when farther from any landmark. \n",
    "This should be visible as concentrating these segments of ANN trajectory (color) around (1,1) when evaluated on PCV.\n",
    "\n",
    "On the contrary, PCV trajectories should look similar on ANN CVs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plumed=f\"{gmx} -p plumed\"\n",
    "\n",
    "!grep WHOLEMOLECULES plumed-ann.dat >plumed-pcv-driver.dat\n",
    "with open('plumed-pcv-driver.dat','a') as plmd:\n",
    "    plmd.write(\"p1: PROPERTYMAP REFERENCE=landmarks-pcv.pdb PROPERTY=X,Y LAMBDA=50.0 NEIGH_SIZE=50 NEIGH_STRIDE=1\\n\")\n",
    "    plmd.write('PRINT ARG=p1.X,p1.Y,p1.zzz STRIDE=1 FILE=COLVAR-pcv-driver FMT=%8.4f\\n')\n",
    "\n",
    "# md.mdp: dt = 2fs, nstxout = 5000 => one frame per 10 ps\n",
    "\n",
    "!{plumed} driver --mf_xtc {base2}.xtc --plumed plumed-pcv-driver.dat --timestep 10 --trajectory-stride 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# COLVAR-ann: PRINT STRIDE=100 ~ 200 fs  => 50x finer than .xtc\n",
    "stride = 50\n",
    "pureann=np.loadtxt('COLVAR-ann')[::stride].T[1:3]\n",
    "\n",
    "# driver pukes one more value\n",
    "pcvonann=np.loadtxt('COLVAR-pcv-driver').T[1:3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_cv2(first,second,stride=1,start=0,stop=-1):\n",
    "    both=np.concatenate((first,second),axis=1)\n",
    "    xymin = np.min(both,axis=1) * 1.1\n",
    "    xymax = np.max(both,axis=1) * 1.1\n",
    "    first = first[:,::stride]\n",
    "    second = second[:,::stride]\n",
    "\n",
    "    if stop == -1:\n",
    "        stop = first.shape[1]\n",
    "        \n",
    "    norm = plt.Normalize(0,first.shape[1])\n",
    "    cmap = plt.get_cmap('rainbow')\n",
    "    first = first[:,start:stop]\n",
    "    second = second[:,start:stop]\n",
    "    cb = range(start,stop)\n",
    "    \n",
    "    plt.figure(figsize=(14,12))\n",
    "    plt.subplot(221)\n",
    "    plt.xlim((xymin[0],xymax[0]))\n",
    "    plt.ylim((xymin[1],xymax[1]))\n",
    "    plt.scatter(*first,c=cb,marker='.',cmap=cmap,norm=norm)\n",
    "    plt.colorbar(cmap=cmap,norm=norm)\n",
    "    \n",
    "    plt.subplot(222)\n",
    "    plt.xlim((xymin[0],xymax[0]))\n",
    "    plt.ylim((xymin[1],xymax[1]))\n",
    "    plt.scatter(*second,c=cb,marker='.',cmap=cmap,norm=norm)\n",
    "    plt.colorbar(cmap=cmap,norm=norm)\n",
    "    \n",
    "    plt.subplot(223)\n",
    "    plt.plot(range(start,stop),first[0],label='first')\n",
    "    plt.plot(range(start,stop),second[0],label='second')\n",
    "    plt.legend()\n",
    "    \n",
    "    plt.subplot(224)\n",
    "    plt.plot(range(start,stop),first[0],label='first')\n",
    "    plt.plot(range(start,stop),second[1],label='second')\n",
    "    plt.legend()\n",
    "    \n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# whole trajectory first\n",
    "plot_cv2(pureann,pcvonann,stride=50,start=0,stop=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# segments out of landmarks\n",
    "plot_cv2(pureann,pcvonann,stride=50,start=220,stop=320)\n",
    "# plot_cv2(pureann,pcvonann,stride=50,start=100,stop=150)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO\n",
    "# porovnani proti literature (jine CV)\n",
    "# spojit trajektorie, udelat ruzne projekce -- esencialni souradnice i isomap, podivat se, jak tam vypadaji jednotlive trajektorie"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}