Guillaume Chevrot

“Anyone who stops learning is old, anyone who keeps learning stays young."

A interactive 2D map in the IPython notebook

In this notebook, I plot an interactive 2D map from a molecular dynamics trajectory in the IPython notebook in just a few lines of code. For this purpose, we will take advantage of the powerful package MMTK.

In [1]:
from __future__ import print_function, division

from MMTK import *
from MMTK.Trajectory import Trajectory, TrajectoryOutput, SnapshotGenerator
from MMTK.Proteins import Protein, PeptideChain

import numpy as np
In [2]:
%install_ext https://raw.githubusercontent.com/rasbt/watermark/master/watermark.py

Installed watermark.py. To use it, type:
  %load_ext watermark
In [3]:
%load_ext watermark
In [4]:
%watermark -v -m -p numpy,MMTK

CPython 2.7.9
IPython 2.3.1

numpy 1.8.2
MMTK 2.7.9

compiler   : GCC 4.4.7 20120313 (Red Hat 4.4.7-1)
system     : Linux
release    : 3.11.0-26-generic
machine    : x86_64
processor  : x86_64
CPU cores  : 24
interpreter: 64bit

NetCDF file format - MMTK trajectory (The trajectory can be easily translated from outcomes of popular Molecular dynamics software like Gromacs.)

In [6]:
filename = 'traj_prot_nojump.nc'

Reading the trajectory

In [7]:
trajectory = Trajectory(None, filename)
universe = trajectory.universe
proteins = universe.objectList(Protein)
chain = proteins[0][0]
In [8]:
# modified fonction, so it can be called with multiprocessing
def calpha_2dmap(trajectory = trajectory, dt = 10000, t = range(0, len(trajectory))):
    """
    Calculate the distance between carbon alpha in the first protein of trajectory.
    
    Parameters
    ----------
    trajectory: MMTK.Trajectory
    
    dt : int
        dt in ns. The distances will be calculated every dt ns.
    t : int 
        range of steps to take into account
    """
    dist = []
    universe = trajectory.universe
    proteins = universe.objectList(Protein)
    chain = proteins[0][0]
    traj = trajectory[t]
    for n, step in enumerate(traj):
        if int(step['time']) % dt == 0:
            universe.setConfiguration(step['configuration'])
            for i in np.arange(len(chain)-1):
                for j in np.arange(len(chain)-1):
                    dist.append(universe.distance(chain[i].peptide.C_alpha,
                                                  chain[j].peptide.C_alpha))
    return(dist)
In [9]:
dist = calpha_2dmap()
dist = np.array(dist)

Reshaping the array

We reshape the dist array, so the first dimension corresponds to the time.

In [10]:
dist = dist.reshape(int(len(dist) / ((len(chain)-1) * (len(chain)-1))), len(chain)-1, len(chain)-1)
In [11]:
dist.shape # time, n_calpha, n_calpha
Out[11]:
(31, 339, 339)

So the first dimension corresponds to the time and the 2 other dimensions correspond to the distances between the Carbon-$\alpha$ atoms of the 339 residues of the protein.

In [12]:
dist_mean = dist[-10:].mean(axis=0) # mean over the last 10 steps
In [13]:
dist_mean.shape
Out[13]:
(339, 339)

Static 2D map

We can naturally plot a static 2D map ...

In [14]:
fig = plt.figure()
plt.imshow(dist_mean, plt.get_cmap('YlOrRd'))
plt.colorbar()
Out[14]:
<matplotlib.colorbar.Colorbar instance at 0x7f665845a128>

Interactive 2D map

... but why wouldn't we plot an interactive one?

In [15]:
import IPython.html.widgets as widgets
from IPython.html.widgets import interact, interactive, fixed
In [16]:
len(dist[0])
Out[16]:
339
In [17]:
@interact(
          t=(0,len(dist)-1),
          calpha_x_min = (0, len(dist[0])),
          calpha_x_max = (0, len(dist[0])),
          calpha_y_min = (0, len(dist[0])),
          calpha_y_max = (0, len(dist[0]))
          )

def interactive_map(t=0, calpha_x_min = 0, calpha_x_max = len(dist[0]), calpha_y_min = 0, calpha_y_max = len(dist[0])):
    fig = plt.figure(figsize=(7,7))
    ax = fig.add_subplot(1, 1, 1)
    ax.set_xlim([calpha_x_min, calpha_x_max])
    ax.set_ylim([calpha_y_min, calpha_y_max])
    plt.imshow(dist[t], plt.get_cmap('YlOrRd'))
    cbar = plt.colorbar()
    cbar.ax.tick_params(labelsize=14)
    plt.clim(0.0,7.0)
    plt.title('2D map')

<matplotlib.figure.Figure at 0x7f664b315910>

NOTE: the corresponding plot is encapsulated in the first video at the end of this page.

So now we can have our interactive 2D map. You can observe the 2D map in function of time or zoom to a specific region. If you are insterested in some specific distances, you can directly narrow a specific zone like that:

In [18]:
@interact(
          t=(0,len(dist)-1),
          calpha_x_min = (211, 227),
          calpha_x_max = (211, 227),
          calpha_y_min = (110, 143),
          calpha_y_max = (110, 143))

def interactive_map(t=len(dist)-1, calpha_x_min = 0, calpha_x_max = len(dist[0]), calpha_y_min = 0, calpha_y_max = len(dist[0])):
    fig = plt.figure(figsize=(7,7))
    ax = fig.add_subplot(1, 1, 1)
    ax.set_xlim([calpha_x_min, calpha_x_max])
    ax.set_ylim([calpha_y_min, calpha_y_max])
    ax.tick_params(labelsize=14)
    plt.imshow(dist[t], plt.get_cmap('YlOrRd'), aspect = 'auto')
    cbar = plt.colorbar()
    cbar.ax.tick_params(labelsize=14)
    plt.clim(0.0,6.0)
    plt.title('2D map')
    plt.tight_layout()
    plt.savefig('2d.jpg', dpi=1000)

<matplotlib.figure.Figure at 0x7f664b52ba10>

NOTE: the corresponding plot is encapsulated in the second video.

NOTE - if you cannot watch the videos on this page, you will find them at these URLs: