prepare.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6851ea37",
   "metadata": {},
   "source": [
    "# ASMSA: Prepare and check input files\n",
    "\n",
    "**Next steps**\n",
    "- [tune.ipynb](tune.ipynb): Perform initial hyperparameter tuning for this molecule\n",
    "- [train.ipynb](train.ipynb): Use results of previous tuning in more thorough training\n",
    "- [md.ipynb](md.ipynb): Use a trained model in MD simulation with Gromacs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c922b53-87a2-4018-8972-2707a06021c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "#avoid TF to consume GPU memory\n",
    "import tensorflow as tf\n",
    "tf.config.set_visible_devices([], 'GPU')\n",
    "tf.config.list_logical_devices()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b70ab11",
   "metadata": {},
   "outputs": [],
   "source": [
    "import mdtraj as md\n",
    "import numpy as np\n",
    "import urllib.request\n",
    "import asmsa\n",
    "import os\n",
    "import re\n",
    "import tensorflow as tf\n",
    "import gromacs as gmx\n",
    "import gromacs.fileformats as gf\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aaba1b8f-d0ff-4264-9712-d06ad9ac964f",
   "metadata": {},
   "source": [
    "## Prepare input files\n",
    "\n",
    "Tryptophan cage files are downloaded in this section from our Google drive. \n",
    "\n",
    "This is for demonstration purpose, in real use the inputs should be placed here, and _conf, traj, topol, index_ variables set to their filenames names."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b505da2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# download trpcage files\n",
    "urllib.request.urlretrieve(\"https://drive.google.com/uc?export=download&id=1ddOJPQxXCw3jY3Yds6bm-EwGps8ClVGQ\",\"trpcage_correct.pdb\")\n",
    "urllib.request.urlretrieve(\"https://drive.google.com/uc?export=download&id=1FRM3-bCdbesShVcyRfk0cj0ILBL1ycbb\",\"trpcage_red.xtc\")\n",
    "\n",
    "#urllib.request.urlretrieve(\"https://drive.google.com/uc?export=download&id=19RZmGgz9goXEAbreUd0OBCKWr5UAVN5d\",\"index_correct.ndx\")\n",
    "#urllib.request.urlretrieve(\"https://drive.google.com/uc?export=download&id=1G-4HRnc1R-LqAArFh0DTY-e5ULd8Goly\",\"topol_correct.top\")\n",
    "#urllib.request.urlretrieve(\"https://drive.google.com/uc?export=download&id=1A28ik84vDhIzyHPHrL8nTxgjzOQAmzLU\",\"trpcage_correct.gro\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3b1a498-01db-4f56-8acc-fad309e987b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define input files\n",
    "\n",
    "base = 'trpcage'\n",
    "\n",
    "# input conformation, it should not contain hydrogens\n",
    "conf = base + '.pdb'\n",
    "\n",
    "# input trajectory\n",
    "# atom numbering must be consistent with {conf}, no hydrogens as well\n",
    "\n",
    "traj = \"trpcage_red.xtc\"\n",
    "\n",
    "# everything else is generated with pdb2gmx to make sure the files are consistent"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8316ea0f-936e-49ff-994b-59c89af73b95",
   "metadata": {},
   "source": [
    "#### Density of additional internal coordinates\n",
    "\n",
    "In how many randomly sampled distances from all atom-to-atom one atom should appear in average"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d63d71f3-eef0-4875-b875-cbd9402be267",
   "metadata": {},
   "outputs": [],
   "source": [
    "nb_density = 2 # integer in [1, n_atoms-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52ab5846-7498-4e6e-88b3-9fc9c4ffd706",
   "metadata": {},
   "outputs": [],
   "source": [
    "topol = base + '.top'\n",
    "index = base + '.ndx'\n",
    "gro = base + '.gro'\n",
    "\n",
    "with open('inputs.py','w') as i:\n",
    "    i.write(f'''\n",
    "base = '{base}'\n",
    "conf = '{conf}'\n",
    "traj = '{traj}'\n",
    "topol = '{topol}'\n",
    "index = '{index}'\n",
    "gro = '{gro}'\n",
    "\n",
    "nb_density = {nb_density}\n",
    "''')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4393cf6-7f51-4a0a-b627-7a510ef48737",
   "metadata": {},
   "outputs": [],
   "source": [
    "gmx.pdb2gmx(f=conf,p=topol,n=index,o=gro,water='tip3p',ff='amber99sb-ildn')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46bad735-2917-4645-8d9b-32e0bf8bbcaf",
   "metadata": {},
   "source": [
    "## Sanity checks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef96d527",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the trajectory, it should report expected numbers of frames and atoms/residua\n",
    "\n",
    "tr = md.load(traj,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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17683c42",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visual check, all frames should look \"reasonable\"\n",
    "\n",
    "# Because of different conventions of numbering atoms in proteins,\n",
    "# PDB file {conf} and the trajectory {traj} can become inconsistent, and this would appear here \n",
    "# as rather weird shapes of the molecule\n",
    "\n",
    "import nglview as nv\n",
    "\n",
    "v = nv.show_mdtraj(tr)\n",
    "v.clear()\n",
    "v.add_representation(\"licorice\")\n",
    "v"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "875fb6f9-d5f7-457f-8362-f6eac58889ff",
   "metadata": {},
   "source": [
    "## Split datasets\n",
    "\n",
    "Split trajectory into 3 parts. Each part will represent training, validation and testing dataset respectively. The workflow is following:\n",
    "1. Shuffle configurations in trajectory\n",
    "2. Select proportions to divide the trajectory\n",
    "3. Divide the trajectory\n",
    "4. Compute RMSD between\n",
    "   * **train x validation** trajectory and filter similar structures in train trajectory\n",
    "   * **train x test** trajectory and filter similar structures in train trajectory\n",
    "   * **test x validation** trajectory and filter similar structures in test trajectory\n",
    "5. Transform into internal coordinates\n",
    "6. Save internal coordinates as datasets which can be loaded in **train.ipynb** and **tune.ipynb** notebooks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e6439d8-3423-40a1-a98b-259e6d0b1446",
   "metadata": {},
   "outputs": [],
   "source": [
    "# shuffle the trajectory so the configurations are dispersed across all datasets\n",
    "np.random.shuffle(tr.xyz)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d742966c-94b3-475c-aecc-0c727e10e2a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# - set proportions for train, validation and test datasets\n",
    "# - proportions must be equal to 1 when added together\n",
    "train = .7\n",
    "validation = .15\n",
    "test = .15\n",
    "\n",
    "assert train + validation + test == .9999999999999999 or 1\n",
    "\n",
    "tr_i = len(tr) * train\n",
    "X_train = tr.slice(slice(0,int(tr_i)))\n",
    "\n",
    "va_i = len(tr) * validation\n",
    "X_validate = tr.slice(slice(int(tr_i),int(tr_i)+int(va_i)))\n",
    "\n",
    "te_i = len(tr) * test\n",
    "X_test = tr.slice(slice(int(tr_i)+int(va_i),len(tr)))\n",
    "\n",
    "X_train.xyz.shape, X_validate.xyz.shape, X_test.xyz.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ac9c322-9427-4771-8da1-0558b834c705",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train.save_xtc('train.xtc')\n",
    "X_validate.save_xtc('validate.xtc')\n",
    "X_test.save_xtc('test.xtc')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a15fec5-5e5f-4385-b043-728b8a36fd05",
   "metadata": {},
   "outputs": [],
   "source": [
    "# eventual recovery\n",
    "\"\"\"\n",
    "X_train = md.load_xtc('train.xtc',conf)\n",
    "X_validate = md.load_xtc('validate.xtc',conf)\n",
    "X_test = md.load_xtc('test.xtc',conf)\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b0089b2-3b3c-4683-99a7-c4cc0bed26fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get RMSD from train trajectory compared to validation trajectory\n",
    "gmx.select(s=conf,on='backbone.ndx',select='Backbone')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a1ca5b1-a885-46ae-a30d-fd2ff2fb63f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "gmx.rms(s=conf,f='train.xtc',f2='validate.xtc',n='backbone.ndx',m='trainxval_rmsd.xpm')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bda4330f-6e2f-4ded-a555-a9a164a1a467",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the RMDS matrix\n",
    "txv = gf.XPM('trainxval_rmsd.xpm')\n",
    "txv.array.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8e0259e-f166-4d3b-8e1b-70eb2a5741ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "# minima per row -- for each configuration in train, how far is the nearest one from validation\n",
    "txv_min = np.min(txv.array,axis=1)\n",
    "txv_min.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fededfb-e5ce-43a3-85ea-4edec1e6c709",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(txv_min,bins=50)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f8d6f56-05c9-4b8b-aa2d-da67e16d77a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# drop similar structures (to validation trajectory) in train trajectory to avoid dataset being biased\n",
    "txv_difference = 0.05\n",
    "\n",
    "train_tr = X_train[np.argwhere(txv_min > txv_difference).flat]\n",
    "train_tr.xyz.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22f7ca04-c7fe-4df5-b9fa-bfb64118ab3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_tr.save_xtc('tmp_train.xtc')\n",
    "gmx.rms(s=conf,f='tmp_train.xtc',f2='test.xtc',n='backbone.ndx',m='trainxtest_rmsd.xpm')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55c22d17-7bdd-4c4c-a0c1-c745ddacfe6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "txt = gf.XPM('trainxtest_rmsd.xpm')\n",
    "txt.array.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b2d11b6-1ffe-41bb-a1e5-29856287cb5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "txt_min = np.min(txt.array,axis=1)\n",
    "txt_min.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a14e533f-f83f-41ed-8d4d-c20aafc67fba",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(txt_min,bins=50)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17541f72-e8e5-400c-96ca-9d2372c4d3b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ... one more time with test trajectory & test x validation...\n",
    "txt_difference = 0.05\n",
    "\n",
    "x_train = train_tr[np.argwhere(txt_min > txt_difference).flat]\n",
    "x_train.save_xtc('x_train.xtc')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b027ccef-fd97-494b-bfce-d5769762336f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# test x validation\n",
    "gmx.rms(f='test.xtc',f2='validate.xtc',s=conf,n='backbone.ndx',m='testxvalidate_rmsd.xpm')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26d995e8-361f-441c-986d-f8235efe9a11",
   "metadata": {},
   "outputs": [],
   "source": [
    "txv = gf.XPM('testxvalidate_rmsd.xpm')\n",
    "txv.array.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d986593-432f-45c6-82e9-b3f6a3dbf031",
   "metadata": {},
   "outputs": [],
   "source": [
    "txv_min = np.min(txv.array,axis=1)\n",
    "txv_min.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea949db5-9b85-4508-8e51-414992e060e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(txv_min,bins=50)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6bf59b4-b974-4d22-bed7-9402aa296271",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ... one more time with test trajectory & test x validation...\n",
    "txv_difference = 0.05\n",
    "\n",
    "x_test = X_test[np.argwhere(txv_min > txv_difference).flat]\n",
    "x_test.save_xtc('x_test.xtc')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f294de0-0325-45f7-b14e-5715c5354c10",
   "metadata": {},
   "outputs": [],
   "source": [
    "# recovery\n",
    "'''\n",
    "x_train = md.load('x_train.xtc', top=conf)\n",
    "x_test = md.load('x_test.xtc', top=conf)\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "097a6ea9-58a6-47c6-a5d1-6464d53f946c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get shapes of filtered trajectories that are to be used as datasets\n",
    "validate_tr = md.load('validate.xtc', top=conf)\n",
    "\n",
    "trajs = [x_train, validate_tr, x_test]\n",
    "x_train.xyz.shape, validate_tr.xyz.shape, x_test.xyz.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51e03cf1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# reshuffle the geometries to get frame last so that we can use vectorized calculations later on\n",
    "geoms = []\n",
    "\n",
    "for i in range(len(trajs)):\n",
    "    geoms.append(np.moveaxis(trajs[i].xyz,0,-1))\n",
    "    print(geoms[i].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fff2ff1c-4222-45c9-b650-111d4f05c101",
   "metadata": {},
   "outputs": [],
   "source": [
    "# save geometries\n",
    "\n",
    "tf.data.Dataset.from_tensor_slices(geoms[0]).save('datasets/geoms/train')\n",
    "tf.data.Dataset.from_tensor_slices(geoms[1]).save('datasets/geoms/validate')\n",
    "tf.data.Dataset.from_tensor_slices(geoms[2]).save('datasets/geoms/test')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab4902fd-e2b9-4812-8b38-59bc444d7d1a",
   "metadata": {},
   "source": [
    "### Internal coordinates computation\n",
    "\n",
    "Exercise the ASMSA library on your input. Just check everything appears to work.\n",
    "\n",
    "There are multiple options that can be combined:\n",
    "- use traditional internal coordinates (bond distances, angles, and dihedrals) or not\n",
    "- include additional distances between atoms that may not be bound to express protein folding state more directly\n",
    "   - dense (all-to-all) atom distances, feasible for very small peptides only\n",
    "   - sparse atom distances (only some pairs are chosen)\n",
    "   \n",
    "**Choose the suitable one in the cell bellow, and copy the same to [tune.ipynb](tune.ipynb) and [train.ipynb](train.ipynb)**, they must be consistent\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f16daa6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# traditional internal coordinates (bond distances, angles, and torsions) only\n",
    "\n",
    "# mol = asmsa.Molecule(conf,topol)\n",
    "\n",
    "# internal coordinates and sparse any-any atom distances (not restricted to bonds)\n",
    "# eventually, top (and index) can be left out to use sparse distances only\n",
    "\n",
    "\n",
    "mols = []\n",
    "for i in range(len(geoms)):\n",
    "    sparse_dists = asmsa.NBDistancesSparse(geoms[i].shape[0], density=nb_density)\n",
    "    mols.append(asmsa.Molecule(pdb=conf,top=topol,ndx=index,fms=[sparse_dists]))\n",
    "\n",
    "# dense distances are feasible for very small (upto 5 residua) peptides only\n",
    "\n",
    "# dense_dists = asmsa.NBDistancesDense(geom.shape[0])\n",
    "# mol = asmsa.Molecule(pdb=conf,top=topol,ndx=index,fms=[dense_dists])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14a5bb18",
   "metadata": {},
   "outputs": [],
   "source": [
    "intcoords = []\n",
    "for i in range(len(mols)):\n",
    "    intcoords.append(mols[i].intcoord(geoms[i]).T)\n",
    "    print(intcoords[i].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b239f0e-0cf2-42f8-a3e4-d35afe93ed8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "[train,validate,test] = intcoords"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1986bd6c-7f23-4925-bb34-9464032ab9d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# normalize training set\n",
    "train_mean = np.mean(train,axis=0)\n",
    "train -= train_mean\n",
    "train_scale = np.std(train,axis=0)\n",
    "train /= train_scale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68a06c84-0164-4f4d-a797-ca0a85153442",
   "metadata": {},
   "outputs": [],
   "source": [
    "# normalize test and validation sets\n",
    "test -= train_mean\n",
    "test /= train_scale\n",
    "validate -= train_mean\n",
    "validate /= train_scale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9927a64-0e0b-457e-851d-156737a9c828",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# save for usage in tune/train/test phase\n",
    "\n",
    "tf.data.Dataset.from_tensor_slices(train).save('datasets/intcoords/train')\n",
    "tf.data.Dataset.from_tensor_slices(validate).save('datasets/intcoords/validate')\n",
    "tf.data.Dataset.from_tensor_slices(test).save('datasets/intcoords/test')\n",
    "\n",
    "np.savetxt('datasets/intcoords/mean.txt',train_mean)\n",
    "np.savetxt('datasets/intcoords/scale.txt',train_scale)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94748d7a-b868-4124-97f7-fed4e8fcc24b",
   "metadata": {},
   "source": [
    "### Density of the conformational space\n",
    "\n",
    "- Sample the training trajectory randomly\n",
    "- For each point in the trajectory:\n",
    "  - calculate RMSD to all points in the sample\n",
    "  - pick some number $n$ of nearest ones\n",
    "  - calculate the _density_ at this point as $$ d = \\sum_{i=1}^n e^{-d_i} / n $$  i.e. the nearer the sample points are, the higher the density\n",
    " \n",
    "Altogether, $d$ roughly corresponds to the probability that the molecule during simulation ends up in this area of the conformational space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebb4eee9-eb63-4fdb-9926-3c0cb3b756cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_size = 5000\n",
    "x_train = md.load('x_train.xtc', top=conf)\n",
    "tr_sample = x_train[np.random.choice(len(x_train),sample_size,False)]\n",
    "tr_sample.save('sample.xtc')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b91c4bc-62d5-4923-92dc-b1bf44a0e2b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "gmx.rms(f='x_train.xtc',f2='sample.xtc',s=conf,n='backbone.ndx',m='sample_rmsd.xpm')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4ee579a-b896-4187-b693-7caf26c7daf8",
   "metadata": {},
   "outputs": [],
   "source": [
    "rms = gf.XPM('sample_rmsd.xpm')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84024073-4300-4c10-8d97-c50c040c5091",
   "metadata": {},
   "source": [
    "#### Visual check to verify the sample size is representative\n",
    "- typically, not many distances should be less than 0.1 nm and more than 1 nm \n",
    "(the latter depends on the molecule, can be more for e.g. big disordered proteins)\n",
    "- the histogram should be semi-smooth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "688bdff2-1c1f-48b9-903c-3a119f7950c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(rms.array.flatten(),bins=50)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ebd6aae-f788-444b-809e-291556ad809a",
   "metadata": {},
   "outputs": [],
   "source": [
    "k_nearest = 200\n",
    "rms_sort = np.sort(rms.array.astype(np.float32))\n",
    "erms = np.exp(-rms_sort[:,:k_nearest])\n",
    "dens = (np.sum(erms,axis=1)-1.) / (erms.shape[1] - 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d14e9e66-9bb8-42dd-b37d-7cd9f00b7b44",
   "metadata": {},
   "source": [
    "#### Histogram of densities\n",
    "- quite high number of points should fall above 0.8, those are low energy basins\n",
    "- the interval [0.5, 1.0] should be reasonably covered\n",
    "- on the contrary, too many points below 0.4 would indicate either insufficient sampling above or too sparse trajectory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc17812a-3fba-42cc-aec6-feab2f28a319",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(dens,bins=20)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a483d963-10a7-401d-a053-0ce234abbddd",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(dens),len(x_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86054b3d-f79a-4e6b-af1f-f308043bf16c",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.savetxt('datasets/train_density.txt',dens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93063ecd-f512-4350-9b86-01a0aeaa7234",
   "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
}