{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Training a model on forces and energies\n",
    "\n",
    "In addition to the energy, machine learning models can also be used to model molecular forces.\n",
    "These are $N_\\mathrm{atoms} \\times 3$ arrays describing the Cartesian force acting on each atom due to the overall\n",
    "(potential) energy. They are formally defined as the negative gradient of the energy $E_\\mathrm{pot}$ with respect to\n",
    "the nuclear positions $\\mathbf{R}$\n",
    "\n",
    "\\begin{equation}\n",
    "\\mathbf{F}^{(\\alpha)} = -\\frac{\\partial E_\\mathrm{pot}}{\\partial \\mathbf{R}^{(\\alpha)}},\n",
    "\\end{equation}\n",
    "\n",
    "where $\\alpha$ is the index of the nucleus.\n",
    "\n",
    "The above expression offers a straightforward way to include forces in machine learning models by simply defining a\n",
    "model for the energy and taking the appropriate derivatives.\n",
    "The resulting model can directly be trained on energies and forces.\n",
    "Moreover, in this manner energy conservation and the correct behaviour under rotations of the molecule is guaranteed.\n",
    "\n",
    "Using forces in addition to energies to construct a machine learning model offers several advantages.\n",
    "Accurate force predictions are important for molecular dynamics simulations, which will be covered in the subsequent\n",
    "tutorial. Forces also encode a greater wealth of information than the energies.\n",
    "For every molecule, only one energy is present, while there are $3N_\\mathrm{atoms}$ force entries.\n",
    "This property, combined with the fact that reference forces can be computed at the same cost as energies, makes models\n",
    "trained on forces and energies very data efficient.\n",
    "\n",
    "In the following, we will show how to train such force models and how to use them in practical applications."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Preparing the data\n",
    "\n",
    "The process of preparing the data is similar to the tutorial on [QM9](tutorial_02_qm9.ipynb). We begin by importing all\n",
    "relevant packages and generating a directory for the tutorial experiments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchmetrics\n",
    "import schnetpack as spk\n",
    "import schnetpack.transform as trn\n",
    "import pytorch_lightning as pl\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "forcetut = \"./forcetut\"\n",
    "if not os.path.exists(forcetut):\n",
    "    os.makedirs(forcetut)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Next, the data needs to be loaded from a suitable dataset. \n",
    "For convenience, we use the MD17 dataset class provided in SchNetPack, which automatically downloads and builds suitable\n",
    "databases containing energies and forces for a range of small organic molecules.\n",
    "In this case, we use the ethanol molecule as an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from schnetpack.datasets import MD17\n",
    "\n",
    "ethanol_data = MD17(\n",
    "    os.path.join(forcetut, \"ethanol.db\"),\n",
    "    molecule=\"ethanol\",\n",
    "    batch_size=10,\n",
    "    num_train=1000,\n",
    "    num_val=1000,\n",
    "    transforms=[\n",
    "        trn.ASENeighborList(cutoff=5.0),\n",
    "        trn.RemoveOffsets(MD17.energy, remove_mean=True, remove_atomrefs=False),\n",
    "        trn.CastTo32(),\n",
    "    ],\n",
    "    num_workers=1,\n",
    "    pin_memory=True,  # set to false, when not using a GPU\n",
    ")\n",
    "ethanol_data.prepare_data()\n",
    "ethanol_data.setup()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The data is split into training (1000 points), validation (1000 points) and test set (remainder).\n",
    "Once again, we subtract the mean of the energies in the training data with a preprocessing transform to precondition\n",
    "our model. This only needs to be done for the energies, since the forces are obtained as derivatives and automatically\n",
    "capture the scale of the data. The subtraction of atomic reference energies is not necessary here, since only molecules\n",
    "of the same composition are used.\n",
    "\n",
    "For custom datasets, the data would have to be loaded via the SchNetPack `ASEAtomsData` and `AtomsDataModule` classes\n",
    "(see tutorial on [data preparation](tutorial_01_preparing_data.ipynb)). In this case, one needs to make sure that the\n",
    "naming of properties is kept consistent with the config. The `schnetpack.properties` module provides standard names\n",
    "for a wide range of properties. Here, we use the definitions provided with the `MD17` class.\n",
    "\n",
    "In order to train force models, forces need to be included in the reference data.\n",
    "Once the dataset has been loaded, this can be checked as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "properties = ethanol_data.dataset[0]\n",
    "print(\"Loaded properties:\\n\", *[\"{:s}\\n\".format(i) for i in properties.keys()])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "As can be seen, `energy` and `forces` are included in the properties dictionary. To have a look at the `forces` array\n",
    "and check whether it has the expected dimensions, we can call:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "print(\"Forces:\\n\", properties[MD17.forces])\n",
    "print(\"Shape:\\n\", properties[MD17.forces].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Building the model\n",
    "\n",
    "After having prepared the data in the above way, we can now build and train the force model.\n",
    "This is done in the same three steps as described in [QM9 tutorial](tutorial_02_qm9.ipynb):\n",
    "\n",
    "1. Defining input modules\n",
    "2. Building the representation\n",
    "3. Defining an output module\n",
    "\n",
    "For the representation we can use the same layers as in the previous tutorial:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "cutoff = 5.0\n",
    "n_atom_basis = 30\n",
    "\n",
    "pairwise_distance = (\n",
    "    spk.atomistic.PairwiseDistances()\n",
    ")  # calculates pairwise distances between atoms\n",
    "radial_basis = spk.nn.GaussianRBF(n_rbf=20, cutoff=cutoff)\n",
    "schnet = spk.representation.SchNet(\n",
    "    n_atom_basis=n_atom_basis,\n",
    "    n_interactions=3,\n",
    "    radial_basis=radial_basis,\n",
    "    cutoff_fn=spk.nn.CosineCutoff(cutoff),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Since we want to model forces, we need an additional output module. We will still use the Atomwise to predict the\n",
    "energy. However, since the forces should be described as the derivative of the energy, we have to indicate that the\n",
    "corresponding derviative of the model should be computed.\n",
    "\n",
    "This is done with the ``Forces`` module, which computes the negative derivative of the energy\n",
    " (specified by the supplied ``energy_key``) with respect to the atom positions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "pred_energy = spk.atomistic.Atomwise(n_in=n_atom_basis, output_key=MD17.energy)\n",
    "pred_forces = spk.atomistic.Forces(energy_key=MD17.energy, force_key=MD17.forces)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The input, representation and output modules are then assembled to the neural network potential:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "nnpot = spk.model.NeuralNetworkPotential(\n",
    "    representation=schnet,\n",
    "    input_modules=[pairwise_distance],\n",
    "    output_modules=[pred_energy, pred_forces],\n",
    "    postprocessors=[\n",
    "        trn.CastTo64(),\n",
    "        trn.AddOffsets(MD17.energy, add_mean=True, add_atomrefs=False),\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Training the model\n",
    "\n",
    "Before we can train the model, the training task has to be defined, including the model, loss functions and\n",
    "optimizers. First, the outputs of the models are connected to their respective loss functions using `ModelOutput`.\n",
    "To train the model on energies and forces, we will use a combined loss function:\n",
    "\n",
    "\\begin{equation}\n",
    "\\mathcal{L}(E_\\mathrm{ref},\\mathbf{F}_\\mathrm{ref},E_\\mathrm{pred}, \\mathbf{F}_\\mathrm{pred}) = \\frac{1}{n_\\text{train}} \\sum_{n=1}^{n_\\text{train}} \\left[  \\rho_1 \\left( E_\\mathrm{ref} - E_\\mathrm{pred} \\right)^2 + \\frac{\\rho_2}{3N_\\mathrm{atoms}} \\sum^{N_\\mathrm{atoms}}_\\alpha \\left\\| \\mathbf{F}_\\mathrm{ref}^{(\\alpha)} - \\mathbf{F}_\\mathrm{pred}^{(\\alpha)} \\right\\|^2 \\right],\n",
    "\\end{equation}\n",
    "\n",
    "where we take the predicted forces to be:\n",
    "\n",
    "\\begin{equation}\n",
    "\\mathbf{F}_\\mathrm{pred}^{(\\alpha)} = -\\frac{\\partial E_\\mathrm{pred}}{\\partial \\mathbf{R}^{(\\alpha)}}.\n",
    "\\end{equation}\n",
    "\n",
    "We have introduced the loss weights $\\rho_i$ in order to control the tradeoff between energy and force loss.\n",
    "In SchNetPack, we can implement such a weighted loss function by setting the loss weights of `ModelOutput`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "output_energy = spk.task.ModelOutput(\n",
    "    name=MD17.energy,\n",
    "    loss_fn=torch.nn.MSELoss(),\n",
    "    loss_weight=0.01,\n",
    "    metrics={\"MAE\": torchmetrics.MeanAbsoluteError()},\n",
    ")\n",
    "\n",
    "output_forces = spk.task.ModelOutput(\n",
    "    name=MD17.forces,\n",
    "    loss_fn=torch.nn.MSELoss(),\n",
    "    loss_weight=0.99,\n",
    "    metrics={\"MAE\": torchmetrics.MeanAbsoluteError()},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now, the training task can be assembled as in the last tutorial:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "task = spk.task.AtomisticTask(\n",
    "    model=nnpot,\n",
    "    outputs=[output_energy, output_forces],\n",
    "    optimizer_cls=torch.optim.AdamW,\n",
    "    optimizer_args={\"lr\": 1e-4},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Finally, we train the model using the PyTorch Lightning `Trainer` for 5 epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "logger = pl.loggers.TensorBoardLogger(save_dir=forcetut)\n",
    "callbacks = [\n",
    "    spk.train.ModelCheckpoint(\n",
    "        model_path=os.path.join(forcetut, \"best_inference_model\"),\n",
    "        save_top_k=1,\n",
    "        monitor=\"val_loss\",\n",
    "    )\n",
    "]\n",
    "\n",
    "trainer = pl.Trainer(\n",
    "    callbacks=callbacks,\n",
    "    logger=logger,\n",
    "    default_root_dir=forcetut,\n",
    "    max_epochs=5,  # for testing, we restrict the number of epochs\n",
    ")\n",
    "trainer.fit(task, datamodule=ethanol_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Training will produce several files in the `model_path` directory, which is `forcetut` in our case.\n",
    "The split is stored in `split.npz`. \n",
    "Checkpoints are written to `checkpoints` periodically, which can be used to restart training.\n",
    "A copy of the best model is stored as `best_inference_model`, which can directly be accessed using the `torch.load`\n",
    "function.\n",
    "\n",
    "You can have a look at the log using Tensorboard:\n",
    "```\n",
    "tensorboard --logdir=forcetut/lightning_logs\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "It should be noted that the model trained here is used exclusively for demonstrative purposes. Accordingly, its size\n",
    "and the training time have been reduced significantly. This puts strong constraints on the accuracy that can be\n",
    "obtained. For practical applications, one would e.g. increase the number of features, the interaction layers, the\n",
    "learning rate schedule and train until convergence (removing the `n_epochs` keyword from the `trainer`).\n",
    "To quickly get started training state-of-the-art models, have a look at the command line interface,\n",
    "which comes with a series of pre-built configurations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Using the model\n",
    "\n",
    "Since all models in SchNetPack are stored in the same way, we can use the trained force model in exactly the same manner\n",
    "as described in the [QM9 tutorial](tutorial_02_qm9.ipynb). To load the model stored in the `best_inference_model` file,\n",
    "we use the `torch.load` function. It will automatically be moved to the device it was trained on. The `AtomsConverter` can then be used to directly operate on ASE atoms objects (e.g. a molecule loaded from a file)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from ase import Atoms\n",
    "from schnetpack.utils.compatibility import load_model\n",
    "\n",
    "# set device\n",
    "device = \"cuda\"\n",
    "\n",
    "# load model\n",
    "model_path = os.path.join(forcetut, \"best_inference_model\")\n",
    "best_model = load_model(model_path, device=device)\n",
    "\n",
    "# set up converter\n",
    "converter = spk.interfaces.AtomsConverter(\n",
    "    neighbor_list=trn.ASENeighborList(cutoff=5.0), dtype=torch.float32, device=device\n",
    ")\n",
    "\n",
    "# create atoms object from dataset\n",
    "structure = ethanol_data.test_dataset[0]\n",
    "atoms = Atoms(\n",
    "    numbers=structure[spk.properties.Z], positions=structure[spk.properties.R]\n",
    ")\n",
    "\n",
    "# convert atoms to SchNetPack inputs and perform prediction\n",
    "inputs = converter(atoms)\n",
    "results = best_model(inputs)\n",
    "\n",
    "print(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Interface to ASE\n",
    "\n",
    "Having access to molecular forces also makes it possible to perform a variety of different simulations.\n",
    "The `SpkCalculator` offers a simple way to perform all computations available in the ASE package ([QM9 tutorial](tutorial_02_qm9.ipynb)).\n",
    "Below, we create an ASE calculator from the trained model and the previously generated `atoms` object\n",
    "(see [Preparing the data](#Preparing-the-data)).\n",
    "One important point is, that the MD17 dataset uses kcal/mol and kcal/mol/Å as units for energies and forces.\n",
    "For use with ASE, these need to be converted to the standard internal ASE units eV and eV/Å.\n",
    "To do so, we need to pass the units of the energies and positions used by the model to the calculator.\n",
    "The calculator will use these to derive the units of properties like forces and stress."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "calculator = spk.interfaces.SpkCalculator(\n",
    "    model=model_path,\n",
    "    neighbor_list=trn.ASENeighborList(cutoff=5.0),\n",
    "    energy_key=MD17.energy,\n",
    "    force_key=MD17.forces,\n",
    "    energy_unit=\"kcal/mol\",\n",
    "    position_unit=\"Ang\",\n",
    ")\n",
    "\n",
    "atoms.calc = calculator\n",
    "\n",
    "print(\"Prediction:\")\n",
    "print(\"energy:\", atoms.get_total_energy())\n",
    "print(\"forces:\", atoms.get_forces())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Among the simulations which can be done by using ASE and a force model are geometry optimisation,\n",
    "normal mode analysis and simple molecular dynamics simulations.\n",
    "\n",
    "The `AseInterface` of SchNetPack offers a convenient way to perform basic versions of these computations.\n",
    "Only a file specifying the geometry of the molecule and a pretrained model are needed.\n",
    "\n",
    "We will first generate a XYZ file containing an ethanol configuration:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from ase import io\n",
    "\n",
    "# Generate a directory for the ASE computations\n",
    "ase_dir = os.path.join(forcetut, \"ase_calcs\")\n",
    "\n",
    "if not os.path.exists(ase_dir):\n",
    "    os.mkdir(ase_dir)\n",
    "\n",
    "# Write a sample molecule\n",
    "molecule_path = os.path.join(ase_dir, \"ethanol.xyz\")\n",
    "io.write(molecule_path, atoms, format=\"xyz\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The `AseInterface` is initialized by passing the path to the molecule, the model and a computation directory.\n",
    "In addition, the names of energies and forces model output,\n",
    "as well as their units, need to be provided (similar to the `SpkCalculator`).\n",
    "Computation device and floating point precision can be set via the `device` and `dtype` arguments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "ethanol_ase = spk.interfaces.AseInterface(\n",
    "    molecule_path,\n",
    "    ase_dir,\n",
    "    model_file=model_path,\n",
    "    neighbor_list=trn.ASENeighborList(cutoff=5.0),\n",
    "    energy_key=MD17.energy,\n",
    "    force_key=MD17.forces,\n",
    "    energy_unit=\"kcal/mol\",\n",
    "    position_unit=\"Ang\",\n",
    "    device=\"cpu\",\n",
    "    dtype=torch.float64,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Geometry optimization\n",
    "\n",
    "For some applications it is neccessary to relax a molecule to an energy minimum.\n",
    "In order to perform this optimization of the molecular geometry, we can simply call"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "ethanol_ase.optimize(fmax=1e-2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Since we trained only a reduced model, the accuracy of energies and forces is not optimal and several steps are\n",
    "needed to optimize the geometry.\n",
    "\n",
    "### Normal mode analysis\n",
    "\n",
    "Once the geometry was optimized, normal mode frequencies can be obtained from the Hessian (matrix of second derivatives)\n",
    "of the molecule. The Hessian is a measure of the curvature of the potential energy surface and normal mode frequencies\n",
    "are useful for determining, whether an optimization has reached a minimum. Using the `AseInterface`, normal mode\n",
    "frequencies can be obtained via:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "ethanol_ase.compute_normal_modes()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Imaginary frequencies indicate, that the geometry optimisation has not yet reached a minimum.\n",
    "The `AseInterface` also creates an `normal_modes.xyz` file which can be used to visualize the vibrations with jmol.\n",
    "\n",
    "### Molecular dynamics\n",
    "\n",
    "Finally, it is also possible to basic run molecular dynamics simulations using this interface.\n",
    "To do so, we first need to prepare the system, where we specify the simulation file.\n",
    "This routine automatically initializes the velocities of the atoms to a random number corresponding to a certain average\n",
    "kinetic energy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "ethanol_ase.init_md(\"simulation\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The actual simulation is performed by calling the function `run_md` with a certain number of steps:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "ethanol_ase.run_md(1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "During simulation, energies and geometries are logged to `simulation.log` and `simulation.traj`, respectively.\n",
    "\n",
    "We can for example visualize the evolution of the systems total and potential energies as\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Load logged results\n",
    "results = np.loadtxt(os.path.join(ase_dir, \"simulation.log\"), skiprows=1)\n",
    "\n",
    "# Determine time axis\n",
    "time = results[:, 0]\n",
    "\n",
    "# Load energies\n",
    "energy_tot = results[:, 1]\n",
    "energy_pot = results[:, 2]\n",
    "energy_kin = results[:, 3]\n",
    "\n",
    "# Construct figure\n",
    "plt.figure(figsize=(14, 6))\n",
    "\n",
    "# Plot energies\n",
    "plt.subplot(2, 1, 1)\n",
    "plt.plot(time, energy_tot, label=\"Total energy\")\n",
    "plt.plot(time, energy_pot, label=\"Potential energy\")\n",
    "plt.ylabel(\"E [eV]\")\n",
    "plt.legend()\n",
    "\n",
    "plt.subplot(2, 1, 2)\n",
    "plt.plot(time, energy_kin, label=\"Kinetic energy\")\n",
    "plt.ylabel(\"E [eV]\")\n",
    "plt.xlabel(\"Time [ps]\")\n",
    "plt.legend()\n",
    "\n",
    "temperature = results[:, 4]\n",
    "print(\"Average temperature: {:10.2f} K\".format(np.mean(temperature)))\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "As can be seen, the potential and kinetic energies fluctuate, while the total energy (sum of potential and kinetic\n",
    "energy) remains approximately constant. This is a good demonstration for the energy conservation obtained by modeling\n",
    "forces as energy derivatives. Unfortunately, this also means that energy conservation is not a sufficient measure for\n",
    "the quality of the potential.\n",
    "\n",
    "However, frequently one is interested in simulations where the system is coupled to an external heat bath.\n",
    "This is the same as saying that we wish to keep the average kinetic energy of the system and hence temperature close to a\n",
    "certain value. Currently, the average temperature only depends on the random velocities drawn during the initialization\n",
    "of the dynamics. Keeping a constant temperature average can be achieved by using a so-called thermostat.\n",
    "In the `AseInterface`, simulations with a thermostat (to be precise a Langevin thermostat) can be carried out by\n",
    "providing the `temp_bath` keyword. A simulation with e.g. the target temperature of 300K is performed via:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "ethanol_ase.optimize(fmax=1e-2)  # reoptimize structure\n",
    "\n",
    "ethanol_ase.init_md(\"simulation_300K\", temp_bath=300, reset=True)\n",
    "ethanol_ase.run_md(20000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We can now once again plot total and potential energies.\n",
    "Instead of the kinetic energy, we plot the temperature (both quantities are directly related)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "skip_initial = 5000\n",
    "\n",
    "# Load logged results\n",
    "results = np.loadtxt(os.path.join(ase_dir, \"simulation_300K.log\"), skiprows=1)\n",
    "\n",
    "# Determine time axis\n",
    "time = results[skip_initial:, 0]\n",
    "# 0.02585\n",
    "# Load energies\n",
    "energy_tot = results[skip_initial:, 1]\n",
    "energy_pot = results[skip_initial:, 2]\n",
    "\n",
    "# Construct figure\n",
    "plt.figure(figsize=(14, 6))\n",
    "\n",
    "# Plot energies\n",
    "plt.subplot(2, 1, 1)\n",
    "plt.plot(time, energy_tot, label=\"Total energy\")\n",
    "plt.plot(time, energy_pot, label=\"Potential energy\")\n",
    "plt.ylabel(\"Energies [eV]\")\n",
    "plt.legend()\n",
    "\n",
    "# Plot Temperature\n",
    "temperature = results[skip_initial:, 4]\n",
    "\n",
    "# Compute average temperature\n",
    "print(\"Average temperature: {:10.2f} K\".format(np.mean(temperature)))\n",
    "\n",
    "plt.subplot(2, 1, 2)\n",
    "plt.plot(time, temperature, label=\"Simulation\")\n",
    "plt.ylabel(\"Temperature [K]\")\n",
    "plt.xlabel(\"Time [ps]\")\n",
    "plt.plot(time, np.ones_like(temperature) * 300, label=\"Target\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Since our molecule is now subjected to external influences via the thermostat the total energy is no longer conserved.\n",
    "However, the simulation temperature now fluctuates near to the requested 300K.\n",
    "This can also be seen by computing the temperature average over time, which is now close to the desired value in contrast to the previous simulation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Summary\n",
    "\n",
    "In this tutorial, we have trained a SchNet model on energies and forces using the MD17 ethanol dataset as an example. \n",
    "We have then evaluated the performance of the model and performed geometry optimisation, normal mode analysis and basic molecular dynamic simulations using the SchNetPack ASE interface.\n",
    "\n",
    "While these simulations can already be useful for practical applications, SchNetPack also comes with its own molecular dynamics package.\n",
    "This package makes it possible to run efficient simulations on GPU and also offers access to advanced techniques, such as ring polymer dynamics.\n",
    "In the next tutorial, we will cover how to perform molecular dynamics simulations directly with SchNetPack."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.8"
  },
  "nbsphinx": {
   "execute": "never"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}