reconstruct.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1aa28469-d6ee-4b2d-8268-e294041f3b43",
   "metadata": {},
   "outputs": [],
   "source": [
    "import asmsa\n",
    "import mdtraj as md\n",
    "import nglview as nv\n",
    "import numpy as np\n",
    "from tensorflow.keras.saving import load_model\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5eb6a3b1-21ea-4df6-920f-4c6c32dff337",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['REQUEST_CPU'] = '8'\n",
    "os.environ['REQUEST_RAM'] = '32'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19e6b9ad-0462-4aa3-b847-c4adeb0f7fbe",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"conf = \"trpcage_correct.pdb\"\n",
    "topol = \"topol_correct.top\"\n",
    "index = 'index_correct.ndx'\n",
    "gro = 'trpcage_correct.gro' \"\"\"\n",
    "exec(open('inputs.py').read())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6a06f60-3791-4f12-8f07-f5bcd8eee09a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# outputs of prepare.ipynb\n",
    "\n",
    "#train_tr = md.load('x_train.xtc',top=conf)\n",
    "#test_tr = md.load('x_test.xtc',top=conf)\n",
    "\n",
    "train_tr = md.load('train.xtc',top=conf)\n",
    "test_tr = md.load('test.xtc',top=conf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03926d18-1a73-41c3-9c27-cc21ac2a506a",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_g = np.moveaxis(train_tr.xyz,0,-1)\n",
    "test_g = np.moveaxis(test_tr.xyz,0,-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c4d1412-192e-456f-a7ee-d266161cd4f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_g.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a989797d-4fa5-4b26-9727-015a6709460e",
   "metadata": {},
   "outputs": [],
   "source": [
    "sparse_dists = asmsa.NBDistancesSparse(train_g.shape[0], density=2)\n",
    "mol = asmsa.Molecule(pdb=conf,top=topol,ndx=index,fms=[sparse_dists])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba11905f-09be-4267-972d-5d095b26c07b",
   "metadata": {},
   "outputs": [],
   "source": [
    "mol.model.dihed4_model.forward(torch.tensor(test_g)).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35704c30-55b3-4677-beb7-4358ecaf5d8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_int = mol.intcoord(train_g)\n",
    "train_int.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cb6b80b-e7da-448f-8982-2d36eb93d7cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_int = mol.intcoord(test_g)\n",
    "test_int.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01907b6c-247e-46b8-a19c-c4e8ec653bb7",
   "metadata": {},
   "outputs": [],
   "source": [
    "topol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6a077b7-7b5e-4218-a482-0ff0e089eb10",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "mol.dihed4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a23ffc5-a1b9-4707-a372-56cc21e0c243",
   "metadata": {},
   "outputs": [],
   "source": [
    "mol.dihed9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ec0ba28-5ff2-4ebb-8101-106dd55b3e7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "latent = np.loadtxt('sample_latent.txt')\n",
    "decout = np.loadtxt('sample_int.txt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b45e628-4e8c-4d38-9135-d7bfcede87dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "latent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2008aeb-ea3b-472b-a5bc-c2cb30448ace",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc4cc5bb-b1eb-4adc-830b-503f0b49c596",
   "metadata": {},
   "outputs": [],
   "source": [
    "# not anymore (faking) ... true output of decoder predicting something should come here\n",
    "\n",
    "out_idx = 113\n",
    "\n",
    "# dec_out = test_int[:,out_idx]\n",
    "dec_out = decout[out_idx,:]\n",
    "dec_out.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "085c400e-1a99-4059-8cd8-718b470f459d",
   "metadata": {},
   "outputs": [],
   "source": [
    "diff = train_int.T - dec_out\n",
    "msd = np.sum(diff * diff,axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d47de5e-7ad0-409c-bdca-48d236d96b6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "minidx = np.argmin(msd)\n",
    "minidx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ebfe027-47a4-4053-a71b-9be366e61bcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_mean = np.loadtxt('datasets/intcoords/mean.txt',dtype=np.float32)\n",
    "train_scale = np.loadtxt('datasets/intcoords/scale.txt',dtype=np.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3cfa128-ee10-46e1-b02b-58bc903dc953",
   "metadata": {},
   "outputs": [],
   "source": [
    "dec_out_scaled = dec_out * train_scale + train_mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e91c8399-4e5b-49bc-be93-a57ab654dacb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# see train.ipynb\n",
    "grotr = md.load(gro)\n",
    "nhs = grotr.topology.select('element != H')\n",
    "\n",
    "with open(index) as f:\n",
    "    f.readline()\n",
    "    ndx = np.fromstring(\" \".join(f),dtype=np.int32,sep=' ')-1\n",
    "\n",
    "pdb2gmx = nhs[np.argsort(ndx)]+1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09e7de53-8b63-459a-b59e-60097dde07d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "pdb2gmx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11155cbd-f570-4359-a59a-fecc3d96d1cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "dihed9_gmx = [\n",
    "    [ pdb2gmx[a] for a in d ]\n",
    "    for d in mol.dihed9\n",
    "]\n",
    "dihed4_gmx = [\n",
    "    [ pdb2gmx[a] for a in d ]\n",
    "    for d in mol.dihed4\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9f82726-e6d1-4587-a265-b8e83aac0688",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "dihed9_gmx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b806c27e-526a-4524-a640-be093145a2ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gromacs as gmx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66235a25-c9a3-4da9-ab2f-3d3c648041b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "gmx.select(s=gro,on='bb.ndx',select='Backbone')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcb1b93a-8f40-46b0-8a6a-880d53fd88b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('bb.ndx') as bb:\n",
    "    bb.readline()\n",
    "    l = \" \".join(bb.readlines())\n",
    "    bbndx = np.array(l.split(),np.int32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a9bcb9a-35fb-4923-ab4e-ca90e4d0f3f8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f396e1e2-169f-4f2d-bcc2-f8a30a320326",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db86d864-e1c9-4735-8d4f-bd3431ff0b97",
   "metadata": {},
   "outputs": [],
   "source": [
    "mintr = train_tr[minidx]\n",
    "#mintr = train_tr[32000]\n",
    "mintr.save('min.pdb')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf528f89-1335-4f10-82a2-c129d0116e76",
   "metadata": {},
   "outputs": [],
   "source": [
    "nv.show_mdtraj(mintr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8bb7494-85cf-4388-92b8-f013b638a3b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "gmx.pdb2gmx(f='min.pdb',o='min.gro',p='min.top',water='tip3p',ff='amber94')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69a27e1d-383d-4916-b321-17d9c04b0add",
   "metadata": {},
   "outputs": [],
   "source": [
    "mdbox=2.0\n",
    "gmx.editconf(f='min.gro',o=\"min-box.gro\",c=True,d=str(mdbox),bt=\"dodecahedron\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97764850-a7be-412e-ae63-69a2ff6c8920",
   "metadata": {},
   "outputs": [],
   "source": [
    "bbndx[0:4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10fdf0ef-74dc-4c13-9f85-8b785fb570a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "dih9_ndx = [ np.where(np.all(dihed9_gmx == bbndx[i:i+4],axis=1))[0][0] for i in range(len(bbndx)-3) ]\n",
    "\n",
    "# XXX: backbone dihedrals seem to be all in dih9 and none in dih4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3478d817-5aef-4621-89fc-31a02334df3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(dih9_ndx), bbndx.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "725339f1-ab69-48ad-8f1b-6d5dcb14b860",
   "metadata": {},
   "outputs": [],
   "source": [
    "off = mol.dihed4.shape[0]*2+mol.angles.shape[0]+mol.bonds.shape[0]\n",
    "size = mol.dihed9.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a9180ce-4429-4f83-a60a-5216e2ed837f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "dec_dih9_sc = dec_out_scaled[off:off+size*2]\n",
    "dec_dih9_sc.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db730439-92bb-453b-835b-455bb1fa5c27",
   "metadata": {},
   "outputs": [],
   "source": [
    "dec_dih9 = np.arctan2(dec_dih9_sc[0:size],dec_dih9_sc[size:]) / np.pi * 180."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f76c4f41-e434-4a36-a128-5d457afbe04a",
   "metadata": {},
   "outputs": [],
   "source": [
    "bb_angles = [ dec_dih9[i] for i in dih9_ndx ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee721a47-47a4-4b46-94ed-de08d2ef1a5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('posre.itp','w') as p:\n",
    "    p.write('[ dihedral_restraints ]\\n')\n",
    "    for i in range(0,len(bbndx)-3):\n",
    "        p.write('  '.join(map(str,bbndx[i:i+4])))\n",
    "        p.write(f' 1 {bb_angles[i]} 0 5000\\n')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2ffcb49-7436-4588-a243-e0f18dde3f70",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('min.mdp','w') as m:\n",
    "    m.write('''\n",
    "integrator  = steep         ; Algorithm (steep = steepest descent minimization)\n",
    "emtol       = 1000.0        ; Stop minimization when the maximum force < 1000.0 kJ/mol/nm\n",
    "emstep      = 0.01          ; Minimization step size\n",
    "nsteps      = 50000         ; Maximum number of (minimization) steps to perform\n",
    "\n",
    "nstxout                 = 1         \n",
    "nstvout                 = 0         \n",
    "nstfout                 = 0         \n",
    "nstlog                  = 5\n",
    "nstxout-compressed      = 1\n",
    "\n",
    "; Parameters describing how to find the neighbors of each atom and how to calculate the interactions\n",
    "nstlist         = 1         ; Frequency to update the neighbor list and long range forces\n",
    "cutoff-scheme   = Verlet    ; Buffered neighbor searching\n",
    "ns_type         = grid      ; Method to determine neighbor list (simple, grid)\n",
    "coulombtype     = Cut-off   ; Treatment of long range electrostatic interactions\n",
    "rcoulomb        = 1.0       ; Short-range electrostatic cut-off\n",
    "rvdw            = 1.0       ; Short-range Van der Waals cut-off\n",
    "pbc             = xyz       ; Periodic Boundary Conditions in all 3 dimensions\n",
    "\n",
    "disre           = Simple\n",
    "define                  = -DPOSRES \n",
    "''')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c0ceb0b-a84f-48ef-bf72-b7b19a3e125a",
   "metadata": {},
   "outputs": [],
   "source": [
    "gmx.grompp(f='min.mdp',c='min-box.gro',p='min.top',o='min.tpr')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aabfb276-c695-45a3-ba12-f0cdb263b205",
   "metadata": {},
   "outputs": [],
   "source": [
    "nv.show_file('min-box.gro')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06a381c1-d625-4015-868b-9f22990a0efc",
   "metadata": {},
   "outputs": [],
   "source": [
    "gmx.mdrun(deffnm='min')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1759a08f-6f52-4b8c-8a6e-b032a444c2f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "!tail -30 min.log"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe831450-2811-4bbd-8f87-26974d5fc044",
   "metadata": {},
   "outputs": [],
   "source": [
    "tr = md.load('min.trr',top='min-box.gro')[:100]\n",
    "v=nv.show_mdtraj(tr)\n",
    "#v.clear()\n",
    "v.add_representation('licorice')\n",
    "v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5a344d9-ddf7-43f6-ac77-43ed2c6a8331",
   "metadata": {},
   "outputs": [],
   "source": [
    "nv.show_file('min-box.gro')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79a5ada6-590a-42d2-b80d-cba1eae09e0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "!cp min.top restrained.top\n",
    "\n",
    "off = mol.dihed4.shape[0]*2+mol.dihed9.shape[0]*2+mol.angles.shape[0]+mol.bonds.shape[0]\n",
    "\n",
    "dec_dist = dec_out_scaled[off:]\n",
    "\n",
    "\n",
    "with open('restrained.top','a') as t:\n",
    "    t.write('''\n",
    "[ distance_restraints ]\n",
    "''')\n",
    "    for i,d in enumerate(sparse_dists.bonds):\n",
    "        t.write(f'{pdb2gmx[d[0]]} {pdb2gmx[d[1]]} 1 {i} 2 {dec_dist[i]*.99} {dec_dist[i]*1.01} 42.0 10.0\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63dc1616-971a-4096-aea4-6d574bd38624",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('min.mdp','w') as m:\n",
    "    m.write('''\n",
    "integrator  = steep         ; Algorithm (steep = steepest descent minimization)\n",
    "emtol       = 1000.0        ; Stop minimization when the maximum force < 1000.0 kJ/mol/nm\n",
    "emstep      = 0.01          ; Minimization step size\n",
    "nsteps      = 50000         ; Maximum number of (minimization) steps to perform\n",
    "\n",
    "nstxout                 = 0         \n",
    "nstvout                 = 0         \n",
    "nstfout                 = 0         \n",
    "nstlog                  = 5\n",
    "nstxout-compressed      = 0\n",
    "\n",
    "; Parameters describing how to find the neighbors of each atom and how to calculate the interactions\n",
    "nstlist         = 1         ; Frequency to update the neighbor list and long range forces\n",
    "cutoff-scheme   = Verlet    ; Buffered neighbor searching\n",
    "ns_type         = grid      ; Method to determine neighbor list (simple, grid)\n",
    "coulombtype     = Cut-off   ; Treatment of long range electrostatic interactions\n",
    "rcoulomb        = 1.0       ; Short-range electrostatic cut-off\n",
    "rvdw            = 1.0       ; Short-range Van der Waals cut-off\n",
    "pbc             = xyz       ; Periodic Boundary Conditions in all 3 dimensions\n",
    "\n",
    "disre           = Simple\n",
    "''')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5f5fb2a-3741-4790-a4a9-300c25301d15",
   "metadata": {},
   "outputs": [],
   "source": [
    "gmx.grompp(f=\"min.mdp\",c=\"min-box.gro\",p='restrained.top',o=\"min.tpr\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e33821ba-a6e4-4d19-94d1-768e1fd422fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "mdrun(deffnm=\"min\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5168f4f1-78f4-4c5e-bb1f-c5eb09688e35",
   "metadata": {},
   "outputs": [],
   "source": [
    "solm = md.load('min.gro')\n",
    "nhs = solm.topology.select('element != H')\n",
    "solm.atom_slice(nhs,inplace=True)\n",
    "v = nv.show_mdtraj(solm)\n",
    "\n",
    "refm = md.load_pdb(conf)\n",
    "refm.xyz = train_g[:,:,minidx]\n",
    "refm.superpose(solm)\n",
    "\n",
    "v.add_component(refm)\n",
    "v.clear(component=0)\n",
    "v.clear(component=1)\n",
    "\n",
    "v.add_representation('licorice',color='green',component=1)\n",
    "v.add_representation('licorice',color='red',component=0)\n",
    "\n",
    "v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb6f1bf1-4318-4faf-bb98-5594d2a18da5",
   "metadata": {},
   "outputs": [],
   "source": [
    "gmx.energy(f='min.edr',input=['Potential','Dis.-Rest.'],o='min.xvg')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "084963a7-c147-468d-8607-2d8fe6e77de0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "energ = np.loadtxt('min.xvg',comments=['#','@'])\n",
    "plt.plot(energ[:,1],label='Dis.-Rest.')\n",
    "plt.plot(energ[:,2],label='Potential')\n",
    "plt.plot(energ[:,2]-energ[:,1],label='Net')\n",
    "plt.yscale('log')\n",
    "plt.legend()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65805c1f-8e9c-4636-b1bf-4d1681ee7333",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e28f8664-1499-4ae6-9a09-ecde9a0ad84b",
   "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.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}