train.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6851ea37",
   "metadata": {},
   "source": [
    "# ASMSA: Train AAE model with the tuned hyperparameters\n",
    "\n",
    "**Previous steps**\n",
    "- [prepare.ipynb](prepare.ipynb): Download and sanity check input files\n",
    "- [tune.ipynb](tune.ipynb): Perform initial hyperparameter tuning for this molecule\n",
    "\n",
    "**Next step**\n",
    "- [md.ipynb](md.ipynb): Use a trained model in MD simulation with Gromacs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "235ca1f6",
   "metadata": {},
   "source": [
    "## Notebook setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91559377-60e1-421a-a51e-5e78f0c1b99b",
   "metadata": {},
   "outputs": [],
   "source": [
    "threads = 2\n",
    "import os\n",
    "os.environ['OMP_NUM_THREADS']=str(threads)\n",
    "import tensorflow as tf\n",
    "\n",
    "# PyTorch favours OMP_NUM_THREADS in environment\n",
    "import torch\n",
    "\n",
    "# Tensorflow needs explicit cofig calls\n",
    "tf.config.threading.set_inter_op_parallelism_threads(threads)\n",
    "tf.config.threading.set_intra_op_parallelism_threads(threads)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b70ab11",
   "metadata": {},
   "outputs": [],
   "source": [
    "from asmsa.tuning_analyzer import TuningAnalyzer\n",
    "import tensorflow_probability as tfp\n",
    "import matplotlib.pyplot as plt\n",
    "import mdtraj as md\n",
    "import numpy as np\n",
    "import urllib.request\n",
    "from tensorflow import keras\n",
    "import keras_tuner\n",
    "import asmsa.visualizer as visualizer\n",
    "import asmsa"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aaba1b8f-d0ff-4264-9712-d06ad9ac964f",
   "metadata": {},
   "source": [
    "## Input files\n",
    "\n",
    "All input files are prepared (up- or downloaded) in [prepare.ipynb](prepare.ipynb). \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fcd3c43e",
   "metadata": {},
   "outputs": [],
   "source": [
    "exec(open('inputs.py').read())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "149cc516-0303-47ea-8b73-b84666bed151",
   "metadata": {},
   "source": [
    "## Apply the tuning results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a56f21c6-eeca-4268-bbb7-7232c9e1bd6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pick from plots in tune.ipynb\n",
    "\n",
    "best_enc_seed=96\n",
    "best_disc_seed=32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2054d30-a16c-40ba-8369-cbd3d6ce2fba",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Get best HP from latest tuning\n",
    "analyzer = TuningAnalyzer()\n",
    "analyzer.get_best_hp(num_trials=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ff484e1-b334-491b-9179-1a9484fa0045",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Select HP to use by specifying trial_id\n",
    "#  e.g: trial_id = '483883b929b3445bff6dee9759c4d50ee3a4ba7f0db22e665c49b5f942d9693b'\n",
    "# ... or don't specify, by default use the trial with the lowest score\n",
    "trial_id = ''\n",
    "\n",
    "hps = None\n",
    "for trial in analyzer.sorted_trials:\n",
    "    if trial['trial_id'] == trial_id:\n",
    "        hps = trial['hp']\n",
    "    \n",
    "if not hps:\n",
    "    print(f'Could not find trial with specified ID, using one with the lowest score - {analyzer.sorted_trials[0][\"trial_id\"]}')\n",
    "    hps = analyzer.sorted_trials[0]['hp']\n",
    "    \n",
    "print(hps)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69a77b87-93ed-4ed1-a613-666e268676d4",
   "metadata": {},
   "source": [
    "## Load datasets\n",
    "Load filtered trajectory datasets that were processed in **prepare.ipynb**. Trajectories are in internal coordinates format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9c602fd-a6c9-4e97-ad1c-d27e99bab9c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load train dataset\n",
    "X_train = tf.data.Dataset.load('datasets/intcoords/train')\n",
    "\n",
    "# get batched version of dataset to feed to AAE model for training\n",
    "X_train_batched = X_train.batch(hps['batch_size'],drop_remainder=True)\n",
    "\n",
    "# get numpy version for visualization purposes\n",
    "X_train_np = np.stack(list(X_train))\n",
    "X_train_np.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8bff1f9-7cd4-4369-ab22-63310897c4a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load test dataset\n",
    "X_test = tf.data.Dataset.load('datasets/intcoords/test')\n",
    "\n",
    "# get batched version of dataset to feed to AAE model for prediction\n",
    "X_test_batched = X_test.batch(hps['batch_size'],drop_remainder=True)\n",
    "\n",
    "# get numpy version for testing purposes\n",
    "X_test_np = np.stack(list(X_test))\n",
    "X_test_np.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0da97b2f-c256-4c18-9267-e2ed819b6067",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Merge (zip) the trajectory density of the training set points\n",
    "# those will be aligned with the probability density of the prior distribution \n",
    "\n",
    "dens = tf.data.Dataset.from_tensor_slices(np.loadtxt('datasets/train_density.txt'))\n",
    "X_train_dens = tf.data.Dataset.zip((X_train,dens)).batch(hps['batch_size'],drop_remainder=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abc0190f-0b3e-4628-ac2d-fd6996a559ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "for e in X_train_dens.as_numpy_iterator():\n",
    "    break\n",
    "e[0].shape,e[1].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea158bb9-8ccf-48ba-9cb7-ccd7ba8fdfd1",
   "metadata": {},
   "source": [
    "## Train"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "967a2ec3-7104-4e4e-895b-dfb4b454913b",
   "metadata": {},
   "source": [
    "### Distribution prior\n",
    "Train with common prior distributions. See https://www.tensorflow.org/probability/api_docs/python/tfp/distributions for all available distributions. It is ideal to use tuned Hyperparameters for training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7854a732-5841-4a4f-84d3-8239370fc5c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set used prior\n",
    "\n",
    "# this one is (more or less) required to with the density alignment\n",
    "prior = tfp.distributions.MultivariateNormalDiag(loc=[0.,0.])\n",
    "\n",
    "#prior = tfp.distributions.Normal(loc=0, scale=1)\n",
    "# prior = tfp.distributions.Uniform()\n",
    "# prior = tfp.distributions.Weibull(1,0.5)\n",
    "# prior = tfp.distributions.Cauchy(loc=0, scale=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef728783-cfb4-4f5e-9dc5-ba23769496ab",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# prepare model using the best hyperparameters\n",
    "testm = asmsa.AAEModel((X_train_np.shape[1],),\n",
    "                       prior=prior,\n",
    "                       hp=hps,\n",
    "                       enc_seed=best_enc_seed,\n",
    "                       disc_seed=best_disc_seed,\n",
    "                       with_density=True\n",
    "                      )\n",
    "testm.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "706f3e67-02ca-4671-9390-8c19b62ff5ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "# specify earlystopping callback to avoid overfitting\n",
    "monitored_metric = \"AE loss min\"\n",
    "\n",
    "early_stop_cb = tf.keras.callbacks.EarlyStopping(\n",
    "    monitor=monitored_metric,\n",
    "    min_delta=0.0001,\n",
    "    patience=15,\n",
    "    verbose=1,\n",
    "    mode=\"min\",\n",
    "    restore_best_weights=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "074bb15e-9c8c-4c57-b996-6b04beafb9e3",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# train it (can be repeated several times to add more epochs)\n",
    "\n",
    "testm.fit(X_train_dens, # X_train_batched,\n",
    "          epochs=600,\n",
    "          verbose=2, # this flag is essential due to connection with EarlyStopping callback (epoch vs batch)\n",
    "          callbacks=[\n",
    "              early_stop_cb,\n",
    "              visualizer.VisualizeCallback(testm,freq=25,inputs=X_train_np[15000:25000],figsize=(12,3))\n",
    "          ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e21481a-0655-48d1-84b2-8924302cc50c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# - plot AE loss min during training\n",
    "# - specify \"since_epoch\" for better plot scaling (ignore outliers)\n",
    "# - note that numbering of epochs starts at 1, 0th epoch does not exist\n",
    "since_epoch = 1\n",
    "\n",
    "assert since_epoch > 0\n",
    "history = np.array(testm.history.history[monitored_metric])\n",
    "y = history[since_epoch-1:]\n",
    "x = list(range(since_epoch, len(y)+since_epoch))\n",
    "result = np.array(list(map(lambda x: x+1, np.where(history == history.min())[0]))) # add +1 to convert index to epoch\n",
    "\n",
    "[plt.axvline(_x, linewidth=0.5, color='r', ls=':') for _x in result]\n",
    "plt.plot(x, y)\n",
    "plt.title(f'Best weights for metric [{monitored_metric}] at epoch/s {result}')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94ed9886-43e7-470c-9d4f-8470673f0321",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# whatever test\n",
    "''' \n",
    "batch_size = 256\n",
    "\n",
    "val_result = testm.predict(X_test_batched)\n",
    "mse = keras.losses.MeanSquaredError()\n",
    "dataset_size = X_test_np.shape[0]\n",
    "print(dataset_size)\n",
    "mse_result=[]\n",
    "for i in range(0, dataset_size, batch_size):\n",
    "    if i+batch_size > dataset_size:\n",
    "        batch_size = batch_size-(i+batch_size-dataset_size)\n",
    "    batch_mse = mse(X_test_np[i:i+batch_size],val_result[i:i+batch_size]).numpy()\n",
    "    mse_result.append(batch_mse)\n",
    "\n",
    "mse_result'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b11b211d-2c01-4b5c-ad35-9f641b3b4ff0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# final visualization, pick a slice of the input data for demo purposes\n",
    "#visualizer.Visualizer(figsize=(12,3)).make_visualization(testm.call_enc(X_train_np[15000:20000]).numpy())\n",
    "\n",
    "# on test data\n",
    "visualizer.Visualizer(figsize=(12,3)).make_visualization(testm.call_enc(X_test_np).numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16507e04-5671-4aa1-bea0-746130c0ca02",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load testing trajectory for further visualizations and computations\n",
    "tr = md.load('x_test.xtc',top=conf)\n",
    "idx=tr[0].top.select(\"name CA\")\n",
    "\n",
    "# for trivial cases like Ala-Ala, where superposing on CAs fails\n",
    "#idx=tr[0].top.select(\"element != H\") \n",
    "\n",
    "tr.superpose(tr[0],atom_indices=idx)\n",
    "\n",
    "# reshuffle the geometry to get frame last so that we can use vectorized calculations\n",
    "geom = np.moveaxis(tr.xyz ,0,-1)\n",
    "geom.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbbe19b1-ada7-4cb1-8757-6e7e532e81fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Rgyr and rmsd color coded in low dim (rough view)\n",
    "\n",
    "lows = testm.call_enc(X_test_np).numpy()\n",
    "rg = md.compute_rg(tr)\n",
    "base = md.load(conf)\n",
    "rmsd = md.rmsd(tr,base[0])\n",
    "cmap = plt.get_cmap('rainbow')\n",
    "plt.figure(figsize=(12,4))\n",
    "plt.subplot(121)\n",
    "plt.scatter(lows[:,0],lows[:,1],marker='.',c=rg,cmap=cmap)\n",
    "plt.colorbar(cmap=cmap)\n",
    "plt.title(\"Rg\")\n",
    "plt.subplot(122)\n",
    "plt.scatter(lows[:,0],lows[:,1],marker='.',c=rmsd,cmap=cmap)\n",
    "plt.colorbar(cmap=cmap)\n",
    "plt.title(\"RMSD\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f8ebddb-f9d5-4485-a54e-db8516baca02",
   "metadata": {},
   "outputs": [],
   "source": [
    "# not used\n",
    "'''testm.enc.save('enc.keras')\n",
    "testm.dec.save('dec.keras')\n",
    "testm.disc.save('dec.keras')'''"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54b29f3c-0702-41e8-8b57-1d8937d12134",
   "metadata": {},
   "source": [
    "### Image prior\n",
    "\n",
    "**Almost surely broken now with the density alignment**\n",
    "\n",
    "Use Image as a prior distribution. Again use tuned Hyperparameters for better training performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74a6eee1-7ac3-44f3-821b-11df13c83002",
   "metadata": {},
   "outputs": [],
   "source": [
    "urllib.request.urlretrieve(\"https://drive.google.com/uc?export=download&id=1I2WP92MMWS5s5vin_4cvmruuV-1W77Hl\", \"mushroom_bw.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfb59272-f629-4aa3-9f8c-5c9c7b8e410e",
   "metadata": {},
   "outputs": [],
   "source": [
    "mmush = asmsa.AAEModel((X_train_np.shape[1],),\n",
    "                       hp=hps,\n",
    "                       enc_seed=best_enc_seed,\n",
    "                       disc_seed=best_disc_seed,\n",
    "                       prior='mushroom_bw.png'\n",
    "                      )\n",
    "mmush.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5eb74e7-0123-41fc-88dd-91270ac533a7",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "mmush.fit(X_train_batched,\n",
    "          epochs=500,\n",
    "          verbose=2,\n",
    "          callbacks=[\n",
    "              early_stop_cb,\n",
    "              visualizer.VisualizeCallback(mmush,freq=25,inputs=X_train_np[15000:25000],figsize=(12,3))\n",
    "          ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b42cfa5f-739c-4d74-be58-bcfe64703a06",
   "metadata": {},
   "outputs": [],
   "source": [
    "# - plot AE loss min during training\n",
    "# - specify \"since_epoch\" for better plot scaling (ignore outliers)\n",
    "# - note that numbering of epochs starts at 1, 0th epoch does not exist\n",
    "since_epoch = 1\n",
    "monitored_metric = 'AE loss min'\n",
    "\n",
    "assert since_epoch > 0\n",
    "history = np.array(mmush.history.history[monitored_metric])\n",
    "y = history[since_epoch-1:]\n",
    "x = list(range(since_epoch, len(y)+since_epoch))\n",
    "result = np.array(list(map(lambda x: x+1, np.where(history == history.min())[0]))) # add +1 to convert index to epoch\n",
    "\n",
    "[plt.axvline(_x, linewidth=0.5, color='r', ls=':') for _x in result]\n",
    "plt.plot(x, y)\n",
    "plt.title(f'Best weights for metric [{monitored_metric}] at epoch/s {result}')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6031485a-64dd-48c5-b25e-6bbc1dc06ccb",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 256\n",
    "\n",
    "val_result = testm.predict(X_test_batched)\n",
    "mse = keras.losses.MeanSquaredError()\n",
    "dataset_size = X_test_np.shape[0]\n",
    "print(dataset_size)\n",
    "mse_result=[]\n",
    "for i in range(0, dataset_size, batch_size):\n",
    "    if i+batch_size > dataset_size:\n",
    "        batch_size = batch_size-(i+batch_size-dataset_size)\n",
    "    batch_mse = mse(X_test_np[i:i+batch_size],val_result[i:i+batch_size]).numpy()\n",
    "    mse_result.append(batch_mse)\n",
    "\n",
    "mse_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c206ea3-5e10-4930-af49-45b1a35b92ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "step=4\n",
    "tr2 = tr[::step]\n",
    "lows = mmush.call_enc(X_test_np[::step]).numpy()\n",
    "rg = md.compute_rg(tr2)\n",
    "base = md.load(conf)\n",
    "rmsd = md.rmsd(tr2,base[0])\n",
    "cmap = plt.get_cmap('rainbow')\n",
    "plt.figure(figsize=(12,4))\n",
    "plt.subplot(121)\n",
    "plt.scatter(lows[:,0],lows[:,1],marker='.',c=rg,cmap=cmap)\n",
    "plt.colorbar(cmap=cmap)\n",
    "plt.title(\"Rg\")\n",
    "plt.subplot(122)\n",
    "plt.scatter(lows[:,0],lows[:,1],marker='.',c=rmsd,cmap=cmap)\n",
    "plt.colorbar(cmap=cmap)\n",
    "plt.title(\"RMSD\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d11c28e3-3d3e-4cca-8fbf-dbf6d8098368",
   "metadata": {},
   "source": [
    "## Save the model for Gromacs\n",
    "\n",
    "*Another wave of magics ...*\n",
    "\n",
    "There are multiple ways how atoms are numbered in PDB, GRO, etc. files. \n",
    "\n",
    "So far we worked with atoms numbered as in the `conf` PDB file, assuming `traj` XTC file was consistent with those.\n",
    "If the topology was used, it might have had different numbering, as Gromacs likes. \n",
    "\n",
    "In the subsequent simulations, we assume the usual protocol starting with `pdb2gmx` to generate topology,\n",
    "hence Gromacsish atom numbering will be followed afterwards.\n",
    "Therefore we need `plumed.dat` to pick the atoms according to the PDB file order, and skip hydrogens added by Gromacs. \n",
    "\n",
    "Many things can go wrong, therefore we strongly encorage to check the results manually. For example, the first residuum (ASP) of tryptophan cage may look like the following in PDB file:\n",
    "\n",
    "    ATOM      1  N   ASP     1      28.538  39.747  31.722  1.00  1.00           N\n",
    "    ATOM      2  CA  ASP     1      28.463  39.427  33.168  1.00  1.00           C\n",
    "    ATOM      3  C   ASP     1      29.059  37.987  33.422  1.00  1.00           C\n",
    "    ATOM      4  O   ASP     1      30.226  37.748  33.735  1.00  1.00           O\n",
    "    ATOM      5  CB  ASP     1      26.995  39.482  33.630  1.00  1.00           C\n",
    "    ATOM      6  CG  ASP     1      26.889  39.307  35.101  1.00  1.00           C\n",
    "    ATOM      7  OD1 ASP     1      27.749  39.962  35.773  1.00  1.00           O\n",
    "    ATOM      8  OD2 ASP     1      26.012  38.510  35.611  1.00  1.00           O\n",
    "    \n",
    "Which turns Gromacs topology: \n",
    "\n",
    "     1         N3      1    ASP      N      1     0.0782      14.01   ; qtot 0.0782\n",
    "     2          H      1    ASP     H1      2       0.22      1.008   ; qtot 0.2982\n",
    "     3          H      1    ASP     H2      3       0.22      1.008   ; qtot 0.5182\n",
    "     4          H      1    ASP     H3      4       0.22      1.008   ; qtot 0.7382\n",
    "     5         CT      1    ASP     CA      5     0.0292      12.01   ; qtot 0.7674\n",
    "     6         HP      1    ASP     HA      6     0.1141      1.008   ; qtot 0.8815\n",
    "     7         CT      1    ASP     CB      7    -0.0235      12.01   ; qtot 0.858\n",
    "     8         HC      1    ASP    HB1      8    -0.0169      1.008   ; qtot 0.8411\n",
    "     9         HC      1    ASP    HB2      9    -0.0169      1.008   ; qtot 0.8242\n",
    "    10          C      1    ASP     CG     10     0.8194      12.01   ; qtot 1.644\n",
    "    11         O2      1    ASP    OD1     11    -0.8084         16   ; qtot 0.8352\n",
    "    12         O2      1    ASP    OD2     12    -0.8084         16   ; qtot 0.0268\n",
    "    13          C      1    ASP      C     13     0.5621      12.01   ; qtot 0.5889\n",
    "    14          O      1    ASP      O     14    -0.5889         16   ; qtot 0\n",
    "    \n",
    "Besides adding hydrogens, the carboxyl group of the protein backbone (atoms 3,4 in PDB) is pushed down (to become 13,14 in the topology).\n",
    "\n",
    "Consequently, the ATOMS setting in the generated `plumed.dat` must be:\n",
    "\n",
    "    model: PYTORCH_MODEL_CV FILE=model.pt ATOMS=1,5,13,14,7,10,11,12, ...\n",
    "    \n",
    "i.e., the atoms are enumerated *in the order* of PDB file but *referring to numbers* of topology file. \n",
    "\n",
    "If there is any mismatch, the MD simulations are likely to fail, or at least to produce meaningless results.\n",
    "\n",
    "It's also **critical** that `{conf}`, `{top}`, and `{gro}` correspond to one another, and that `{gro}` **includes hydrogens**.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97bb8fe9-d24f-4d97-9049-1e927b2e0255",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tf2onnx\n",
    "import onnx2torch\n",
    "import tempfile\n",
    "\n",
    "def _convert_to_onnx(model, destination_path):\n",
    "#    model = keras.models.load_model(source_path)\n",
    "\n",
    "    input_tensor = model.layers[0]._input_tensor\n",
    "#    input_tensor = model.inputs[0]\n",
    "    input_signature = tf.TensorSpec(\n",
    "        name=input_tensor.name, shape=input_tensor.shape, dtype=input_tensor.dtype\n",
    "    )\n",
    "    output_name = model.layers[-1].name\n",
    "\n",
    "    @tf.function(input_signature=[input_signature])\n",
    "    def _wrapped_model(input_data):\n",
    "        return {output_name: model(input_data)}\n",
    "\n",
    "    tf2onnx.convert.from_function(\n",
    "        _wrapped_model, input_signature=[input_signature], output_path=destination_path\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "635522d7-9f81-4acd-9734-80904c72cfba",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = testm\n",
    "\n",
    "with tempfile.NamedTemporaryFile() as onnx:\n",
    "#    tf2onnx.convert.from_keras(model.enc,output_path=onnx.name)\n",
    "    _convert_to_onnx(model.enc,onnx.name)\n",
    "    torch_encoder = onnx2torch.convert(onnx.name)\n",
    "\n",
    "# load test geometry dataset\n",
    "geom = np.stack(list(tf.data.Dataset.load('datasets/geoms/test')))\n",
    "\n",
    "# XXX: we rely on determinism of the model creation, it must be the same as in prepare.ipynb\n",
    "# better to store it there in onnx, and reload here\n",
    "\n",
    "sparse_dists = asmsa.NBDistancesSparse(geom.shape[0], density=nb_density)\n",
    "mol = asmsa.Molecule(pdb=conf,top=topol,ndx=index,fms=[sparse_dists])\n",
    "    \n",
    "mol_model = mol.get_model()\n",
    "\n",
    "train_mean = np.loadtxt('datasets/intcoords/mean.txt',dtype=np.float32)\n",
    "train_scale = np.loadtxt('datasets/intcoords/scale.txt',dtype=np.float32)\n",
    "\n",
    "def complete_model(x):\n",
    "    return torch_encoder(\n",
    "        ((mol_model(x) - torch.from_numpy(np.reshape(train_mean,(-1,1)))) / torch.from_numpy(np.reshape(train_scale,(-1,1)))).reshape(-1)\n",
    "    )\n",
    "\n",
    "# Save Torch model using TorchScript trace\n",
    "example_input = torch.randn([geom.shape[0], geom.shape[1], 1])\n",
    "traced_script_module = torch.jit.trace(complete_model, example_input)\n",
    "\n",
    "model_file_name = \"model.pt\"\n",
    "traced_script_module.save(model_file_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fac7455e-7acc-401e-af96-117985b43404",
   "metadata": {},
   "outputs": [],
   "source": [
    "# validate\n",
    "\n",
    "example_geom = np.random.rand(geom.shape[0], geom.shape[1], 1)\n",
    "#X = mol.intcoord(example_geom).T\n",
    "X = ((mol.intcoord(example_geom) - np.reshape(train_mean,(-1,1))) / np.reshape(train_scale,(-1,1))).T\n",
    "tf_low = np.array(model.enc(X))\n",
    "\n",
    "torch_geom = torch.tensor(example_geom.reshape(-1), dtype=torch.float32, requires_grad=True)\n",
    "torch_low = traced_script_module(torch_geom)\n",
    "\n",
    "for out in torch_low:\n",
    "    grad = torch.autograd.grad(out, torch_geom, retain_graph=True)\n",
    "\n",
    "    \n",
    "# should be very small, eg. less than 1e-5\n",
    "np.max(np.abs(tf_low - torch_low.detach().numpy()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeaa1bde-50b1-4b9f-b42a-6f0f141eb0f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Atom numbering magic with Gromacs, see above\n",
    "\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",
    "\n",
    "# maybe double check manually wrt. the files\n",
    "# pdb2gmx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2aa07611-6cec-4d3c-b3ba-c0a83af9d496",
   "metadata": {},
   "outputs": [],
   "source": [
    "ndx"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa17907d-21b9-4ad3-80ca-e425e6efc94b",
   "metadata": {},
   "source": [
    "#### Determine range of CVs for simulation\n",
    "\n",
    "Plumed maintains a grid to approximate accumulated bias potential, which size must be known in advance.\n",
    "\n",
    "Making it wider is safe, the simulation is less likely to escape and crash, but there is perfomance penalty.\n",
    "\n",
    "Calculate the CVs on the testset, determine their range, and add some margins\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c26ffa3c-3c65-4cac-8949-9760f5f0ff6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "grid_margin = 1.  # that many times the actual computed size added on both sides\n",
    "\n",
    "lows = model.call_enc(X_test_np).numpy()\n",
    "lmin = np.min(lows,axis=0)\n",
    "lmax = np.max(lows,axis=0)\n",
    "llen = lmax-lmin\n",
    "lmin -= llen * grid_margin\n",
    "lmax += llen * grid_margin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f1692c6-974b-4089-a72b-78b924001fa0",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"plumed.dat\",\"w\") as p:\n",
    "    p.write(f\"\"\"\\\n",
    "RESTART\n",
    "WHOLEMOLECULES ENTITY0=1-{grotr.xyz.shape[1]}\n",
    "model: PYTORCH_MODEL_CV FILE={model_file_name} ATOMS={','.join(map(str,pdb2gmx))}\n",
    "metad: METAD ARG=model.node-0,model.node-1 PACE=1000 HEIGHT=1 BIASFACTOR=15 SIGMA=0.1,0.1 GRID_MIN={lmin[0]},{lmin[1]} GRID_MAX={lmax[0]},{lmax[1]} FILE=HILLS\n",
    "PRINT FILE=COLVAR ARG=model.node-0,model.node-1,metad.bias STRIDE=100\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18bada48-2877-4323-88d1-eaf3353e2ea1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "121abd33-7442-4a33-9141-6e60a30a6bd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# XXX: these were some tests for the density alingment, I'm not sure anymore; not important either"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "482079e2-65c0-4725-9913-59a0a19530cb",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "255404dd-f035-48fb-a200-ba22130c8596",
   "metadata": {},
   "outputs": [],
   "source": [
    "traj = 'validate.xtc'\n",
    "tr = md.load(traj,top=conf)\n",
    "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": "52396d51-0e34-4214-9783-746d60704ec6",
   "metadata": {},
   "outputs": [],
   "source": [
    "geom = np.moveaxis(tr.xyz,0,-1)\n",
    "geom.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95aadc0a-526a-4ac3-b706-5b701e77c937",
   "metadata": {},
   "outputs": [],
   "source": [
    "density = 2\n",
    "sparse_dists = asmsa.NBDistancesSparse(geom.shape[0], density=density)\n",
    "mol = asmsa.Molecule(pdb=conf,top=topol,ndx=index,fms=[sparse_dists])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d75bd03d-c1e7-4848-858c-6940ee15e4f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "intcoord = mol.intcoord(geom).T\n",
    "intcoord -= train_mean\n",
    "intcoord /= train_scale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b335f3f-b9a6-4a56-8131-bda385e86e7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gromacs.fileformats as gf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "744b89c1-8e8c-49d7-8901-ac799d0e49e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "rms = gf.XPM('rmsd.xpm')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb2bcd61-1674-44df-a540-e1e298601671",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(rms.array.flatten(),bins=50)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "497438bc-c6bb-4fec-9c98-79e42d57f8d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "rms_sort = np.sort(rms.array.astype(np.float32))\n",
    "erms = np.exp(-rms_sort[:,:50])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "381e3d09-9544-4739-9292-31a796dc1f25",
   "metadata": {},
   "outputs": [],
   "source": [
    "dens = (np.sum(erms,axis=1)-1.) / (erms.shape[1] - 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6c85c58-3e29-4358-9f42-0f6bb3dff4e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(dens,bins=20)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4a3372f-65c0-424a-bafd-9b1e613c0cea",
   "metadata": {},
   "outputs": [],
   "source": [
    "lows = testm.enc(intcoord)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "706d5984-1d9e-4ed6-b7f2-9ad48a7541ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "lows.shape, dens.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "670ea4d4-de83-49c4-80c5-80cbef9c7452",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.scatter(lows[:,0],lows[:,1],c=dens,marker='.',cmap='rainbow')\n",
    "plt.colorbar()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f22a2af1-ae92-43a8-bd0b-82981afb97c1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2aa5ac50-04c2-49d9-b46b-a9ce7ed19a6e",
   "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
}