add info on extracting end states, use mdtraj to image the molecules

Author: jthorton

simulation_visualisation/simulation_visualisation.ipynb

@@ -7,11 +7,12 @@
    "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) and [MDAnalysis](https://github.com/MDAnalysis/mdanalysis). By the end, you should understand how to:\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. Centre the ligand in the simulation box using `MDAnalysis`\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",
     "## Downloading the example data\n",
     "\n",
@@ -28,27 +29,24 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "--2025-09-05 10:31:50--  https://zenodo.org/records/15375081/files/simulation.nc\n",
-      "Resolving zenodo.org (zenodo.org)... 2001:1458:d00:25::100:372, 2001:1458:d00:24::100:f6, 2001:1458:d00:61::100:2f3, ...\n",
+      "--2025-09-22 13:47:47--  https://zenodo.org/records/15375081/files/simulation.nc\n",
+      "Resolving zenodo.org (zenodo.org)... 2001:1458:d00:25::100:372, 2001:1458:d00:61::100:2f3, 2001:1458:d00:24::100:f6, ...\n",
       "Connecting to zenodo.org (zenodo.org)|2001:1458:d00:25::100:372|:443... connected.\n",
       "HTTP request sent, awaiting response... 200 OK\n",
       "Length: 516886878 (493M) [application/octet-stream]\n",
-      "Saving to: ‘simulation.nc.1’\n",
+      "Saving to: ‘simulation.nc.2’\n",
       "\n",
-      "simulation.nc.1     100%[===================>] 492.94M  12.2MB/s    in 62s     \n",
-      "\n",
-      "2025-09-05 10:32:53 (7.96 MB/s) - ‘simulation.nc.1’ saved [516886878/516886878]\n",
-      "\n",
-      "--2025-09-05 10:32:53--  https://zenodo.org/records/15375081/files/hybrid_system.pdb\n",
-      "Resolving zenodo.org (zenodo.org)... 2001:1458:d00:24::100:f6, 2001:1458:d00:61::100:2f3, 2001:1458:d00:25::100:372, ...\n",
-      "Connecting to zenodo.org (zenodo.org)|2001:1458:d00:24::100:f6|:443... connected.\n",
+      "simulation.nc.2       7%[>                   ]  37.40M  6.76MB/s    eta 88s    ^C\n",
+      "--2025-09-22 13:47:55--  https://zenodo.org/records/15375081/files/hybrid_system.pdb\n",
+      "Resolving zenodo.org (zenodo.org)... 2001:1458:d00:61::100:2f3, 2001:1458:d00:24::100:f6, 2001:1458:d00:25::100:372, ...\n",
+      "Connecting to zenodo.org (zenodo.org)|2001:1458:d00:61::100:2f3|:443... connected.\n",
       "HTTP request sent, awaiting response... 200 OK\n",
       "Length: 388547 (379K) [application/octet-stream]\n",
-      "Saving to: ‘hybrid_system.pdb.1’\n",
+      "Saving to: ‘hybrid_system.pdb.2’\n",
       "\n",
-      "hybrid_system.pdb.1 100%[===================>] 379.44K  1.65MB/s    in 0.2s    \n",
+      "hybrid_system.pdb.2 100%[===================>] 379.44K  1.29MB/s    in 0.3s    \n",
       "\n",
-      "2025-09-05 10:32:54 (1.65 MB/s) - ‘hybrid_system.pdb.1’ saved [388547/388547]\n",
+      "2025-09-22 13:47:56 (1.29 MB/s) - ‘hybrid_system.pdb.2’ saved [388547/388547]\n",
       "\n"
      ]
     }
@@ -74,22 +72,22 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 28,
    "id": "05ba7dc7",
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "/Users/joshua/mambaforge/envs/openfe_fast_set/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",
+      "/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"
      ]
     },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "ca7edee8791e482fa36a171ee3edb895",
+       "model_id": "e5a30cb71c014c3c9831cf2f3723ea72",
        "version_major": 2,
        "version_minor": 0
       },
@@ -103,8 +101,10 @@
    ],
    "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",
@@ -125,23 +125,22 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 29,
    "id": "009a8ac7-0ca9-494c-bbd6-177e752b0d4f",
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "/Users/joshua/mambaforge/envs/openfe_fast_set/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",
-      "DEPRECATED: Please use 'center' method\n"
+      "/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"
      ]
     },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "18e16cf1a1144b228e6ab22a06d42002",
+       "model_id": "e9e8ec05be4c4075ad1b13bc9be88c6c",
        "version_major": 2,
        "version_minor": 0
       },
@@ -157,30 +156,45 @@
     "u_1 = mda.Universe(\"hybrid_system.pdb\", \"simulation.nc\", format=FEReader, state_id=-1)\n",
     "\n",
     "w = nv.show_mdanalysis(u_1)\n",
-    "w.center_view()\n",
+    "w.center()\n",
     "w"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "75602d58-6709-4d9b-8a52-9896fdc6a0af",
+   "id": "cbc5534f-846c-486c-9b8b-01e1bab8d880",
    "metadata": {},
    "source": [
-    "## Centring the Ligand\n",
+    "# Extracting the end state positions with `MDAnalysis` \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 transformations in `MDAnalysis` to centre the ligand and wrap the protein around it like so:"
+    "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 beta 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": 8,
-   "id": "a5f16795-1b9c-4198-8694-6568eaba06c7",
+   "execution_count": 30,
+   "id": "0ac90329-4a4a-4c76-be44-447c7ef9bb8b",
    "metadata": {},
    "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/joshua/mambaforge/envs/openfe_dev/lib/python3.12/site-packages/MDAnalysis/core/topologyattrs.py:329: DeprecationWarning: The bfactor topology attribute is only provided as an alias to the tempfactor attribute. It will be removed in 3.0. Please use the tempfactor attribute instead.\n",
+      "  warnings.warn(BFACTOR_WARNING, DeprecationWarning)\n"
+     ]
+    },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "2f9fe1601d5a4c22954bdc2dd69ad482",
+       "model_id": "8d08590641d041e9bda6b83b21c8d95e",
        "version_major": 2,
        "version_minor": 0
       },
@@ -193,18 +207,13 @@
     }
    ],
    "source": [
-    "import MDAnalysis.transformations as trans\n",
+    "# get atoms for state A\n",
+    "bfactor = 0.25\n",
+    "state_atoms = np.array([atom.ix for atom in u_0.atoms if atom.bfactor in (bfactor, 0.5, 0.0)])\n",
+    "state = u_0.atoms[state_atoms]\n",
     "\n",
-    "ligand = u_0.select_atoms(\"resname UNK\")\n",
-    "protein = u_0.select_atoms(\"protein\")\n",
-    "transforms = [\n",
-    "    trans.center_in_box(ligand, wrap=True),\n",
-    "    trans.wrap(protein),\n",
-    "    trans.nojump.NoJump()\n",
-    "]\n",
-    "u_0.trajectory.add_transformations(*transforms)\n",
+    "w = nv.show_mdanalysis(state)\n",
     "\n",
-    "w = nv.show_mdanalysis(u_0)\n",
     "w"
    ]
   },
@@ -213,21 +222,73 @@
    "id": "7026738c-e4fe-461e-85ca-f726e9768a5e",
    "metadata": {},
    "source": [
-    "## Saving the trajectory to file\n",
+    "## Saving the trajectory to file with `MDAnalysis`\n",
     "\n",
-    "We can now use `MDAnalysis` to save our trajectory to a common file format:"
+    "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": 9,
+   "execution_count": 31,
    "id": "6e410679",
    "metadata": {},
-   "outputs": [],
+   "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": [
-    "with mda.Writer('out.xtc', n_atoms=len(u_0.atoms)) as w:\n",
+    "# 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)"
+    "        w.write(u_0.atoms[state_atoms])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "75602d58-6709-4d9b-8a52-9896fdc6a0af",
+   "metadata": {},
+   "source": [
+    "## Centring 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": 33,
+   "id": "a5f16795-1b9c-4198-8694-6568eaba06c7",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "c967de9e5a9a4ac688f2e74e25ede21f",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "NGLWidget(max_frame=500)"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "traj = md.load_xtc(\"out.xtc\", top=\"state_a_topology.pdb\")\n",
+    "traj = traj.image_molecules()\n",
+    "\n",
+    "w = nv.show_mdtraj(traj)\n",
+    "\n",
+    "w"
    ]
   }
  ],