namd-demos/rmsd-rdf/rmsd-rdf.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## On-the fly Analysis\n",
    "\n",
    "In this demo, we will be doing on-the-fly analysis of the data recived by the IMDClient from the MD engine. The two analysis we will be doing will be RDF and RMSF respectively."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Libray Imports\n",
    "\n",
    "We import requisite libraries for the client, the analysis and plotting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import required libraries\n",
    "\n",
    "# client-related libraries\n",
    "from imdclient.IMD import IMDReader\n",
    "import MDAnalysis as mda\n",
    "\n",
    "# MDAnalysis libraries for analysis\n",
    "from MDAnalysis.analysis.rdf import InterRDF\n",
    "from MDAnalysis.analysis.rms import RMSF\n",
    "\n",
    "# plotting and general libraries\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# parallel processing libraries\n",
    "import concurrent.futures\n",
    "import threading"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Initialize the client IMDReader object and analysis objects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NAMD_TOPOL = \"../exec-files/T3_MNN.psf\"\n",
    "\n",
    "# Enable widget mode for live interactive plots in Jupyter\n",
    "%matplotlib widget\n",
    "\n",
    "# Function to initialize the MDAnalysis universe and analyzers\n",
    "def initialize_analysis(topol_path, stream_path):\n",
    "    \"\"\"\n",
    "    Initialize MDAnalysis Universe and set up RDF and RMSF analyzers.\n",
    "\n",
    "    Args:\n",
    "        universe_path (str): Path to the topology file.\n",
    "        trajectory_path (str): Path or identifier for the trajectory file.\n",
    "\n",
    "    Returns:\n",
    "        tuple: MDAnalysis universe, RDF analyzer, RMSF analyzer.\n",
    "    \"\"\"\n",
    "    u = mda.Universe(topol_path, stream_path)\n",
    "    protein = u.select_atoms(\"protein\")\n",
    "    calphas = protein.select_atoms(\"protein and name CA\")\n",
    "    # ionText = \"name POT CLA CAL SOD LIT CES IOD BRO H11 H21 H31 H41\"\n",
    "    # ions = u.select_atoms(ionText)  # Select ions based on their names\n",
    "    g = calphas\n",
    "    # g = u.atoms\n",
    "\n",
    "    rdf_cutoff = 16 # Cutoff for RDF calculation\n",
    "\n",
    "    # Initialize RDF and RMSF\n",
    "    rdf = InterRDF(g, g, range=(0.0, rdf_cutoff))\n",
    "    rmsf = RMSF(g)\n",
    "    rdf._ts = u.trajectory.ts\n",
    "    rmsf._ts = u.trajectory.ts\n",
    "\n",
    "    # Prepare analyzers\n",
    "    rdf._prepare()\n",
    "    rmsf._prepare()\n",
    "\n",
    "    return u, rdf, rmsf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Initialize the plot for live-visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to initialize the plot for live updates\n",
    "def initialize_live_plot():\n",
    "    \"\"\"\n",
    "    Initialize a Matplotlib figure for live plotting.\n",
    "\n",
    "    Returns:\n",
    "        tuple: Figure, axes, line objects for RDF and RMSF.\n",
    "    \"\"\"\n",
    "    # Create subplots\n",
    "    fig, axes = plt.subplots(2, 1, figsize=(5, 6))\n",
    "\n",
    "    # RDF subplot\n",
    "    rdf_line, = axes[0].plot([], [], label=\"RDF\", color=\"blue\")\n",
    "    axes[0].set_title(\"Radial Distribution Function (RDF)\")\n",
    "    axes[0].set_xlabel(\"Distance (Å)\")\n",
    "    axes[0].set_ylabel(\"RDF\")\n",
    "    axes[0].set_ylim(0, 15)\n",
    "    axes[0].legend()\n",
    "\n",
    "    # RMSF subplot\n",
    "    rmsf_line, = axes[1].plot([], [], label=\"RMSF\", color=\"green\")\n",
    "    axes[1].set_title(\"Root Mean Square Fluctuation (RMSF)\")\n",
    "    axes[1].set_xlabel(\"Residue Index\")\n",
    "    axes[1].set_ylabel(\"RMSF (Å)\")\n",
    "    axes[1].legend()\n",
    "\n",
    "    plt.tight_layout()\n",
    "    return fig, axes, rdf_line, rmsf_line"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Perfrom the analysis on-the-fly and upate the plot\n",
    "\n",
    "The two seperate analysis are perfromed on seperate threads and synced before the plot is updated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to perform parallel calculations with live plotting\n",
    "def perform_live_plotting(u, rdf, rmsf, fig, rdf_line, rmsf_line, update_frequency=10):\n",
    "    \"\"\"\n",
    "    Perform RDF and RMSF calculations with live plot updates.\n",
    "\n",
    "    Args:\n",
    "        u (Universe): MDAnalysis universe.\n",
    "        rdf (InterRDF): RDF analyzer.\n",
    "        rmsf (RMSF): RMSF analyzer.\n",
    "        fig (Figure): Matplotlib figure object.\n",
    "        rdf_line (Line2D): Line object for RDF.\n",
    "        rmsf_line (Line2D): Line object for RMSF.\n",
    "        update_frequency (int): Frequency of updates to the live plot.\n",
    "    \"\"\"\n",
    "    executor = concurrent.futures.ThreadPoolExecutor()\n",
    "    count = 0\n",
    "\n",
    "    # Initialize data containers\n",
    "    rdf_bins = rdf.bins\n",
    "    # rdf_data = np.zeros_like(rdf_bins)\n",
    "    # rmsf_data = np.zeros(rmsf.n_atoms)\n",
    "\n",
    "    for ts in u.trajectory:\n",
    "        # Submit parallel tasks for RDF and RMSF calculations\n",
    "        rdf_future = executor.submit(rdf._single_frame)\n",
    "\n",
    "        rmsf._frame_index = count\n",
    "        rmsf_future = executor.submit(rmsf._single_frame)\n",
    "\n",
    "        # Wait for threads to complete\n",
    "        rdf_future.result()\n",
    "        rmsf_future.result()\n",
    "\n",
    "        count += 1\n",
    "\n",
    "        rdf.n_frames = count\n",
    "        rdf._conclude()\n",
    "        rmsf._conclude()\n",
    "\n",
    "        # Update results\n",
    "        rdf_data = np.array(rdf.results[\"rdf\"])\n",
    "        rmsf_data = np.array(rmsf.results[\"rmsf\"])\n",
    "        # rdf_data = np.array(rdf.results.rdf)\n",
    "        # rmsf_data = np.array(rmsf.rmsf)\n",
    "\n",
    "        # Live update of the plot\n",
    "        if count % update_frequency == 0:\n",
    "            rdf_line.set_data(rdf_bins, rdf_data)\n",
    "            rmsf_line.set_data(np.arange(len(rmsf_data)), rmsf_data)\n",
    "\n",
    "            # Adjust axis limits\n",
    "            rdf_line.axes.relim()\n",
    "            rdf_line.axes.autoscale_view()\n",
    "            rmsf_line.axes.relim()\n",
    "            rmsf_line.axes.autoscale_view()\n",
    "\n",
    "            # Efficient redraw for Jupyter\n",
    "            fig.canvas.draw_idle()\n",
    "\n",
    "    # Finalize analyzers\n",
    "    rdf.n_frames = count\n",
    "    rdf._conclude()\n",
    "    rmsf._conclude()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Main program\n",
    "\n",
    "Analysis-plotting is done in a seperate thread so that the main program can handle other JUPyter background processes to display the plot. This is done using the threading library and an optional addition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Main script to execute the analysis and live plotting\n",
    "if __name__ == \"__main__\":\n",
    "    # Set paths for universe and trajectory\n",
    "    topol_path, stream_path = NAMD_TOPOL, \"imd://localhost:8888\"\n",
    "\n",
    "    # Initialize universe and analyzers\n",
    "    u, rdf, rmsf = initialize_analysis(topol_path, stream_path)\n",
    "\n",
    "    # Initialize live plotting\n",
    "    fig, axes, rdf_line, rmsf_line = initialize_live_plot()\n",
    "\n",
    "    # Perform calculations with live plotting\n",
    "    data_thread = threading.Thread(target=perform_live_plotting, args=(u, rdf, rmsf, fig, rdf_line, rmsf_line, 1))\n",
    "    data_thread.daemon = True  # Thread exits when main program ends\n",
    "    data_thread.start()\n",
    "    # perform_live_plotting(u, rdf, rmsf, fig, rdf_line, rmsf_line, update_frequency=1)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "imdclient-demos",
   "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}