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": "b4d8932e-711d-4789-b4c6-c9eb4893b7ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "#%cd villin"
   ]
  },
  {
   "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=128\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": "b7efd2bf-75d3-4d82-98a8-ad4cfc4c4c93",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_val = tf.data.Dataset.load('datasets/intcoords/validate').batch(hps['batch_size'],drop_remainder=True)\n",
    "X_val"
   ]
  },
  {
   "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=False\n",
    "                      )\n",
    "testm.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f95c599-aebb-4195-9e04-c66a33a0479e",
   "metadata": {},
   "outputs": [],
   "source": [
    "testm.summary()"
   ]
  },
  {
   "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",
    "#monitored_metric = 'val_loss'\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_batched, # 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",
    "#          validation_data = X_val,\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": "ef39d5b9-1ecc-455d-9fd3-2526e2be108b",
   "metadata": {},
   "outputs": [],
   "source": [
    "visualizer.Visualizer(figsize=(12,3)).make_visualization(testm.call_enc(X_train_np).numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fb90de1-e327-47f4-adcb-09461492573c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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": "6e9d3156-039e-4918-97d2-828a15970fe1",
   "metadata": {},
   "source": [
    "## Save the encoder and decoder models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e7d45a0-5bdd-4cce-913b-cdd2565e31d6",
   "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": "324e8e1a-dd89-4867-9878-34d2bc53c2de",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5916afd3-a4ec-4cd2-a961-a354bc988e24",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = testm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63987325-2b30-4960-ae21-e9533ead27b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "with tempfile.NamedTemporaryFile() as onnx:\n",
    "    _convert_to_onnx(model.enc,onnx.name)\n",
    "    torch_enc = onnx2torch.convert(onnx.name)\n",
    "\n",
    "example_input = torch.randn([X_train_np.shape[1]])\n",
    "traced_script_module = torch.jit.trace(torch_enc, example_input)\n",
    "\n",
    "traced_script_module.save('encoder.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28a84ebb-e11d-40ad-af88-f029f24ef9e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "example_input.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abcae519-478c-47ee-a7d3-1dbc0ac94e62",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.enc.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7097e35c-a6b2-4833-b59e-0e40fe6632b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "lenc = torch.jit.load('encoder.pt')\n",
    "example_input = np.random.rand(10000,X_train_np.shape[1])\n",
    "rtf = model.enc(example_input)\n",
    "rpt = lenc(torch.tensor(example_input,dtype=torch.float32))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01f55169-d93c-4905-bc79-b0109a919290",
   "metadata": {},
   "outputs": [],
   "source": [
    "maxerr = np.max(np.abs(rtf - rpt.detach().numpy()))\n",
    "maxerr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39c80f40-cca7-4441-94f7-d74e32c12039",
   "metadata": {},
   "outputs": [],
   "source": [
    "with tempfile.NamedTemporaryFile() as onnx:\n",
    "    _convert_to_onnx(model.dec,onnx.name)\n",
    "    torch_dec = onnx2torch.convert(onnx.name)\n",
    "\n",
    "example_input = torch.randn([2])\n",
    "traced_script_module = torch.jit.trace(torch_dec, example_input)\n",
    "\n",
    "traced_script_module.save('decoder.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69f370e3-8037-418e-9121-461a195f19d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "ldec = torch.jit.load('decoder.pt')\n",
    "example_input = np.random.rand(10000,2)\n",
    "rtf = model.dec(example_input)\n",
    "rpt = ldec(torch.tensor(example_input,dtype=torch.float32))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94f4d3c9-73be-4836-b6ac-04f4ba0a825e",
   "metadata": {},
   "outputs": [],
   "source": [
    "err = np.abs(rtf - rpt.detach().numpy())\n",
    "train_mean = np.loadtxt('datasets/intcoords/mean.txt',dtype=np.float32).reshape(1,1,-1)\n",
    "rerr = err/np.abs(train_mean)\n",
    "np.max(err),np.max(rerr)"
   ]
  }
 ],
 "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
}