diff --git a/simulation_visualisation/images/centered_system.png b/simulation_visualisation/images/centered_system.png
new file mode 100644
index 0000000..4a12491
Binary files /dev/null and b/simulation_visualisation/images/centered_system.png differ
diff --git a/simulation_visualisation/images/hybrid_system.png b/simulation_visualisation/images/hybrid_system.png
new file mode 100644
index 0000000..186e0cd
Binary files /dev/null and b/simulation_visualisation/images/hybrid_system.png differ
diff --git a/simulation_visualisation/images/hybrid_system_state_1.png b/simulation_visualisation/images/hybrid_system_state_1.png
new file mode 100644
index 0000000..d193cdf
Binary files /dev/null and b/simulation_visualisation/images/hybrid_system_state_1.png differ
diff --git a/simulation_visualisation/images/septop_full.png b/simulation_visualisation/images/septop_full.png
new file mode 100644
index 0000000..66d407a
Binary files /dev/null and b/simulation_visualisation/images/septop_full.png differ
diff --git a/simulation_visualisation/images/septop_state_a.png b/simulation_visualisation/images/septop_state_a.png
new file mode 100644
index 0000000..ff7617e
Binary files /dev/null and b/simulation_visualisation/images/septop_state_a.png differ
diff --git a/simulation_visualisation/images/statea_ligand.png b/simulation_visualisation/images/statea_ligand.png
new file mode 100644
index 0000000..b3910a2
Binary files /dev/null and b/simulation_visualisation/images/statea_ligand.png differ
diff --git a/simulation_visualisation/simulation_visualisation.ipynb b/simulation_visualisation/simulation_visualisation.ipynb
new file mode 100644
index 0000000..28fcf52
--- /dev/null
+++ b/simulation_visualisation/simulation_visualisation.ipynb
@@ -0,0 +1,364 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "id": "20a8d98a",
+ "metadata": {},
+ "source": [
+ "# Extracting and visualising a free energy simulation\n",
+ "\n",
+ "This notebook provides a step-by-step guide to extract and visualise a free energy simulation trajectory from a ``simulation.nc`` file using [openfe-analysis](https://github.com/OpenFreeEnergy/openfe_analysis), [MDAnalysis](https://github.com/MDAnalysis/mdanalysis) and [mdtraj](https://github.com/mdtraj/mdtraj). By the end, you should understand how to:\n",
+ "\n",
+ "1. Extract the trajectory of a ``replica`` or ``single lambda state`` from a ``simulation.nc`` file\n",
+ "2. For a given hybrid topology trajectory, extract the relevant atom positions for the end states using `MDAnalysis`\n",
+ "3. Write out the trajectorie(s) using `MDAnalysis`\n",
+ "4. Centre the ligand in the simulation box using `mdtraj`\n",
+ "\n",
+ "This visualisation workflow can be used with any OpenFE protocols, though the end-state extraction method can be different, as only the relative hybrid topology protocol uses the temperature factor method described here.\n",
+ "\n",
+ "## Downloading the example data\n",
+ "\n",
+ "First, download some example trajectory data. This may take a few minutes due to the size of the simulation file. Please skip this section if you have already done this!"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "id": "be6317291bbf3804",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "File ‘simulation.nc’ already there; not retrieving.\n",
+ "\n",
+ "File ‘hybrid_system.pdb’ already there; not retrieving.\n",
+ "\n"
+ ]
+ }
+ ],
+ "source": [
+ "! wget https://zenodo.org/records/15375081/files/simulation.nc -nc\n",
+ "! wget https://zenodo.org/records/15375081/files/hybrid_system.pdb -nc"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "5d066221-0b8e-4b1b-a047-c989a5957cf7",
+ "metadata": {},
+ "source": [
+ "## Extracting the trajectory with `MDAnalysis`\n",
+ "\n",
+ "The `openfe-analysis` package provides an `MDAnalysis` reader to help extract the trajectory data from the `simulation.nc` file. As the file contains multipule replicas simulated at different lambda states, we must choose which of these to load as a single trajectory. We have two options available to construct the trajectory:\n",
+ "- `state_id`: will construct a trajectory which follows a single Hamiltonian lambda state at the specified value.\n",
+ "- `recplica_id`: will construct a trajectory which follows a single replica at the specified value.\n",
+ "\n",
+ "In this example which uses a trajectory from a relative binding free energy calculation we will load the trajectory at `lambda=0` or the end state corresponding to Ligand A and visulaise the trajectory with `nglview`."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 26,
+ "id": "05ba7dc7",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "/Users/joshua/mambaforge/envs/openfe_dev/lib/python3.12/site-packages/openfe_analysis/utils/multistate.py:41: UserWarning: This is an older NetCDF file that does not yet contain information about the write frequency of positions and velocities. We will assume that positions and velocities were written out at every iteration. \n",
+ " warnings.warn(wmsg)\n"
+ ]
+ }
+ ],
+ "source": [
+ "import MDAnalysis as mda\n",
+ "import mdtraj as md\n",
+ "from openfe_analysis import FEReader\n",
+ "import nglview as nv\n",
+ "import numpy as np\n",
+ "\n",
+ "u_0 = mda.Universe(\"hybrid_system.pdb\", \"simulation.nc\", format=FEReader, state_id=0)\n",
+ "\n",
+ "v = nv.show_mdanalysis(u_0)\n",
+ "# v"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "db81b5fe-9195-488e-b7e9-b3fad9748bdb",
+ "metadata": {},
+ "source": [
+ "\n",
+ "\n",
+ " \n",
+ "\n",
+ "\n",
+ "
\n",
+ ""
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "f0e7b0e7-04b5-42c0-ab65-2202275288d4",
+ "metadata": {},
+ "source": [
+ " Note: The OpenFE relative binding free energy protocol does not save water positions by default, this can be changed via the
output_indices protocol setting.
\n",
+ "\n",
+ "\n",
+ "To view the final state at `lambda=1` we can use negative indexing if we don't know the total number of lambda states."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 27,
+ "id": "009a8ac7-0ca9-494c-bbd6-177e752b0d4f",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "/Users/joshua/mambaforge/envs/openfe_dev/lib/python3.12/site-packages/openfe_analysis/utils/multistate.py:41: UserWarning: This is an older NetCDF file that does not yet contain information about the write frequency of positions and velocities. We will assume that positions and velocities were written out at every iteration. \n",
+ " warnings.warn(wmsg)\n"
+ ]
+ }
+ ],
+ "source": [
+ "u_1 = mda.Universe(\"hybrid_system.pdb\", \"simulation.nc\", format=FEReader, state_id=-1)\n",
+ "\n",
+ "v = nv.show_mdanalysis(u_1)\n",
+ "v.center()\n",
+ "# v"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "4611bee9-b3aa-4414-ab15-54c13947f957",
+ "metadata": {},
+ "source": [
+ "\n",
+ "\n",
+ " \n",
+ "\n",
+ "\n",
+ "
\n",
+ ""
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "cbc5534f-846c-486c-9b8b-01e1bab8d880",
+ "metadata": {},
+ "source": [
+ "# Extracting the end state positions with `MDAnalysis` \n",
+ "\n",
+ "## Relative hybrid topology protocol\n",
+ "\n",
+ "The trajectory data stored in the `simulation.nc` file contains the positions of the end-state ligands in their hybrid topology format. This means only atoms that are unique to the end-states have individual positions, with conserved core atoms sharing a single set of positions. As you might have noticed in the visualisation above, this can complicate the analysis and visualisation of the protein-ligand interactions. However, we can identify the atoms relevant to the end states or core atoms using the tempature factors in the topology file:\n",
+ "\n",
+ "- `0.0`: The non-alchemical atoms (protein, solvent, etc)\n",
+ "- `0.25`: The unique atoms of state A\n",
+ "- `0.5`: The conserved core atoms present in both end states\n",
+ "- `0.75`: The unique atoms of state B\n",
+ "\n",
+ "With this information, we can easily extract the atom positions relevant to `state A` for `lambda=0`:"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 28,
+ "id": "0ac90329-4a4a-4c76-be44-447c7ef9bb8b",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# get atoms for state A\n",
+ "tempfactor = 0.25\n",
+ "\n",
+ "state = sum([u_0.atoms[u_0.atoms.tempfactors == i] for i in (0, 0.5, tempfactor)])\n",
+ "\n",
+ "v = nv.show_mdanalysis(state)\n",
+ "\n",
+ "# v"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "52c52e28-8375-4815-9b0e-7a118ab15f55",
+ "metadata": {},
+ "source": [
+ "\n",
+ "\n",
+ " \n",
+ "\n",
+ "\n",
+ "
\n",
+ ""
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "79512b7c-6759-4c99-92fb-70d66383c6b8",
+ "metadata": {},
+ "source": [
+ "## Separated Topologies\n",
+ "\n",
+ "This protocol represents both end-state ligands explicitly with unique atom positions. To extract the relevant end-state ligand coordinates, you can use `chainid` to select the state where `chainid` `A` and `B` correspond to the `A` and `B` ligands, respectively, in the `solvent` leg. In the `complex` leg, due to the way the systems are constructed, the `chainid`s are `B` and `E`, which again correspond to the `A` and `B` end-states respectively. "
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 29,
+ "id": "4209f067-5dba-4cf8-9c45-03291b371493",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# load the example septop solvent leg topology file\n",
+ "septop = mda.Universe(\"topologies/alchemical_system_septop.pdb\")\n",
+ "v = nv.show_mdanalysis(septop)\n",
+ "v.center()\n",
+ "# make the ions visible\n",
+ "v.add_representation(\"ball+stick\", selection=\"chainid D\")\n",
+ "# v"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "61338e3d-b0d4-45cb-8053-fee8753bf262",
+ "metadata": {},
+ "source": [
+ "\n",
+ "\n",
+ " \n",
+ "\n",
+ "\n",
+ "
\n",
+ ""
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 33,
+ "id": "f836e12d-97f2-474a-b926-a9da05c9eb49",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# select end-state A benzene only\n",
+ "state_a = septop.select_atoms(\"resname UNK and chainid A\")\n",
+ "v = nv.show_mdanalysis(state_a)\n",
+ "v.center()\n",
+ "# v"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "ebd5ae31-5665-4ac9-93ec-956c55e5bfff",
+ "metadata": {},
+ "source": [
+ "\n",
+ "\n",
+ " \n",
+ "\n",
+ "\n",
+ "
\n",
+ ""
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "7026738c-e4fe-461e-85ca-f726e9768a5e",
+ "metadata": {},
+ "source": [
+ "## Saving the trajectory to file with `MDAnalysis`\n",
+ "\n",
+ "We can now use `MDAnalysis` to save the trajectory of the `state A` atoms to a common file format, note that we will also need to write out a new topology file that can be used to load this trajectory:"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 58,
+ "id": "6e410679",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "/Users/joshua/mambaforge/envs/openfe_dev/lib/python3.12/site-packages/MDAnalysis/coordinates/PDB.py:1154: UserWarning: Found no information for attr: 'formalcharges' Using default value of '0'\n",
+ " warnings.warn(\"Found no information for attr: '{}'\"\n"
+ ]
+ }
+ ],
+ "source": [
+ "# write a new PDB topology file for the state A atoms only\n",
+ "state.write(\"state_a_topology.pdb\")\n",
+ "# write the trajectory to an xtc file\n",
+ "with mda.Writer('out.xtc', n_atoms=len(state.atoms)) as w:\n",
+ " for ts in u_0.trajectory:\n",
+ " w.write(u_0.atoms[state.atoms.ix])"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "75602d58-6709-4d9b-8a52-9896fdc6a0af",
+ "metadata": {},
+ "source": [
+ "## Centering the Ligand with `mdtraj`\n",
+ "\n",
+ "You may have noticed in the view above that the ligand seems to have drifted away from the protein, this is a visualisation artifact caused by the use of periodic boundary conditions and the way in which `OpenMM` tries to ensure that all particle positions are written into a single periodic box. We can fix this, however, using `mdtraj` and the [image_molecules](https://mdtraj.org/1.9.3/api/generated/mdtraj.Trajectory.html?highlight=image_molecules#mdtraj.Trajectory.image_molecules) function:"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 30,
+ "id": "a5f16795-1b9c-4198-8694-6568eaba06c7",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "traj = md.load_xtc(\"out.xtc\", top=\"state_a_topology.pdb\")\n",
+ "traj = traj.image_molecules()\n",
+ "\n",
+ "v = nv.show_mdtraj(traj)\n",
+ "\n",
+ "v.center()\n",
+ "# v"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "3347d10e-2533-4cbc-88b9-b273d9a1d8c9",
+ "metadata": {},
+ "source": [
+ "\n",
+ "\n",
+ " \n",
+ "\n",
+ "\n",
+ "
\n",
+ ""
+ ]
+ }
+ ],
+ "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.9"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/simulation_visualisation/topologies/alchemical_system_septop.pdb b/simulation_visualisation/topologies/alchemical_system_septop.pdb
new file mode 100644
index 0000000..44bbc47
--- /dev/null
+++ b/simulation_visualisation/topologies/alchemical_system_septop.pdb
@@ -0,0 +1,62 @@
+ATOM 1 C1x UNK A 1 25.978 5.327 4.779 1.00 0.00 C
+ATOM 2 C2x UNK A 1 26.395 5.074 3.499 1.00 0.00 C
+ATOM 3 C3x UNK A 1 27.340 5.860 2.902 1.00 0.00 C
+ATOM 4 C4x UNK A 1 27.837 6.921 3.569 1.00 0.00 C
+ATOM 5 C5x UNK A 1 27.420 7.196 4.856 1.00 0.00 C
+ATOM 6 C6x UNK A 1 26.498 6.379 5.469 1.00 0.00 C
+ATOM 7 H1x UNK A 1 25.230 4.686 5.245 1.00 0.00 H
+ATOM 8 H2x UNK A 1 25.968 4.235 2.950 1.00 0.00 H
+ATOM 9 H3x UNK A 1 27.689 5.631 1.895 1.00 0.00 H
+ATOM 10 H4x UNK A 1 28.573 7.566 3.089 1.00 0.00 H
+ATOM 11 H5x UNK A 1 27.821 8.060 5.386 1.00 0.00 H
+ATOM 12 H6x UNK A 1 26.187 6.572 6.496 1.00 0.00 H
+TER 13 UNK A 1
+ATOM 14 H1x UNK B 1 37.049 8.204 4.787 1.00 0.00 H
+ATOM 15 H2x UNK B 1 36.339 7.604 6.304 1.00 0.00 H
+ATOM 16 C1x UNK B 1 36.128 7.877 5.270 1.00 0.00 C
+ATOM 17 H3x UNK B 1 35.400 8.687 5.248 1.00 0.00 H
+ATOM 18 C2x UNK B 1 34.120 4.787 4.444 1.00 0.00 C
+ATOM 19 C3x UNK B 1 34.537 4.534 3.164 1.00 0.00 C
+ATOM 20 C4x UNK B 1 35.482 5.320 2.567 1.00 0.00 C
+ATOM 21 C5x UNK B 1 35.979 6.381 3.234 1.00 0.00 C
+ATOM 22 C6x UNK B 1 35.562 6.656 4.521 1.00 0.00 C
+ATOM 23 C7x UNK B 1 34.640 5.839 5.134 1.00 0.00 C
+ATOM 24 H4x UNK B 1 33.372 4.146 4.910 1.00 0.00 H
+ATOM 25 H5x UNK B 1 34.110 3.696 2.614 1.00 0.00 H
+ATOM 26 H6x UNK B 1 35.831 5.092 1.560 1.00 0.00 H
+ATOM 27 H7x UNK B 1 36.715 7.026 2.754 1.00 0.00 H
+ATOM 28 H8x UNK B 1 34.329 6.032 6.160 1.00 0.00 H
+TER 29 UNK B 1
+ATOM 30 Na NA D 1 17.857 -4.057 -2.352 1.00 0.00 Na
+ATOM 31 Na NA D 2 41.388 6.331 12.660 1.00 0.00 Na
+ATOM 32 Cl CL D 3 37.257 -6.694 -1.840 1.00 0.00 Cl
+ATOM 33 Cl CL D 4 30.875 -3.775 0.157 1.00 0.00 Cl
+TER 34 CL D 4
+CONECT 1 2 6 7
+CONECT 2 1 3 8
+CONECT 3 2 4 9
+CONECT 4 3 5 10
+CONECT 5 4 6 11
+CONECT 6 1 5 12
+CONECT 7 1
+CONECT 8 2
+CONECT 9 3
+CONECT 10 4
+CONECT 11 5
+CONECT 12 6
+CONECT 14 16
+CONECT 15 16
+CONECT 16 14 15 17 22
+CONECT 17 16
+CONECT 18 19 23 24
+CONECT 19 18 20 25
+CONECT 20 19 21 26
+CONECT 21 20 22 27
+CONECT 22 16 21 23
+CONECT 23 18 22 28
+CONECT 24 18
+CONECT 25 19
+CONECT 26 20
+CONECT 27 21
+CONECT 28 23
+END