{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# `torch-sim`\n", "\n", "[torch-sim](https://radical-ai.github.io/torch-sim/index.html) is a tool for running (batched) molecular dynamics and structure optimisations in a GPU-accelerated manner.\n", "\n", "You can use **any** model from `graph-pes` with `torch-sim`! Doing this is extremely straightforward, as this notebook hopes to demonstrate. \n", "One caveat is that `torch-sim` expects all models to return energies and stresses in units of `eV` and forces in units of `eV/Å`. If you have a model that was trained in a different unit system, you will need to wrap it in a [graph_pes.models.UnitConverter](https://jla-gardner.github.io/graph-pes/models/root.html#unit-conversion) object so that it can be used in `torch-sim`.\n", "\n", "To get started, you must first ensure you have the `torch-sim-atomistic` package installed:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# this notebook ran successfully with this torch-sim version. \n", "# future versions may require changes as the package evolves.\n", "!pip install graph-pes torch-sim-atomistic==0.2.0" ] }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "Using device: cpu" ] }, "execution_count": 1, "metadata": {}, "output_type": "execute_result" } ], "source": [ "from pathlib import Path\n", "\n", "import ase\n", "import ase.io\n", "import load_atoms\n", "import torch\n", "import torch_sim\n", "from ase.build import bulk\n", "\n", "from graph_pes import models\n", "\n", "%config InlineBackend.figure_format = 'retina'\n", "torch.manual_seed(42) # for reproducibility\n", "\n", "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", "print(f\"Using device: {DEVICE}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## NVT MD\n", "\n", "`torch-sim` makes it easy to run NVT molecular dynamics using a Langevin thermostat. \n", "\n", "To keep things simple, we use a simple `LennardJones` model to run NVT for a single structure.\n", "It's worth repeating that this proceduce will be identical for any other model in `graph-pes`, be that a model you've\n", "[trained from scratch yourself](https://jla-gardner.github.io/graph-pes/quickstart/quickstart.html), a [pre-trained foundation model](https://jla-gardner.github.io/graph-pes/interfaces/mace.html) or a foundation model you have [fine-tuned on your own data](https://jla-gardner.github.io/graph-pes/quickstart/fine-tuning.html)." ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", "
\n", " \n", "