<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 MSM of Mueller potential
In this example we work with the Muller potential, introduced by K. Muller and L. D. Brown in Theo. Chim. Acta (Berl.) 53, 75-93 (1979)
Step1:  Data generation
First of all we import the mueller module, which contains a class (borrowed from Robert Mc Gibbon) that runs the simulations using OpenMM.
Step2:  First we set a number of parameters for the run.
Step3:  Data analysis
Download brownian trajectory from here, and read the latter by using trajectory from COSSIO generator.
Step4:  Trajectory analysis and assignment
We discretize the trajectory by a 2D binning, so the mueller potential is divided in microstates (plot 1). In consequence, both x and y coordinates merge onto a s coordinate (plot 2).
Step5:  MSM
Step6:  MSM
Step7:  The first mode captured by $\lambda_1$ is significantly slower than the others. That mode, which is described by the right eigenvector $\psi^R_1$ as the transition between the two main minima. Next plots correspond to the eigenvector and its projection on the mueller potential.
Step8:  HDBSCAN clustering
Since the 2D binning used above in order to assign discrete states does not take into account the potential form, in the following a density-based clustering is employed instead. Namely, the HDBSCAN algorithm is suitable to deal with any shape clusters with different densities. The latter must be installed from here before going on.
Discretization
Clusters are built from trajectory data. Then, we refine noise and analyze if the resulting condensed tree plot makes sense for our system.
Step9:  MSM validation
We compare the relaxation times obtained from the transition matrix with those corresponding to the 2D binning shown above. Obviously, since there are three clusters in HDBSCAN, we obtain two possible transitions.
Step10:  RATE matrix analysis
The simplicity of this discretization with only three states allows us to employ the Maximum-Likelihood Propagator-Based (MLPB) method in order to compute rate matrix and corresponding eigenvectors.
Step11:  COUNT and TRANSITION matrices
Count and transition matrices give us an intuition of the dynamics for this simple system. Remind the former is strongly dependent on the lag time.
<ASSISTANT_TASK:>
Python Code:
%matplotlib inline
%load_ext autoreload
%autoreload 2
import time
import numpy as np
import h5py
import scipy.stats
import matplotlib.pyplot as plt
import seaborn as sns
sns.set(style="ticks", color_codes=True, font_scale=1.5)
sns.set_style({"xtick.direction": "in", "ytick.direction": "in"})
from collections import OrderedDict
Explanation: MSM of Mueller potential
In this example we work with the Muller potential, introduced by K. Muller and L. D. Brown in Theo. Chim. Acta (Berl.) 53, 75-93 (1979):
\begin{equation}
    V(x,y) = \sum_{j=1}^{4}A_{j} \cdot \exp(a_{j}(x-X_{j}^{0})^{2}+b_{j}(x-X_{j}^{0})(y-Y_{j}^{0})+c_{j}(y-Y_{j}^{0})^{2})
\end{equation}
where $a$=(-1,-1,-6.5,0.7); $b$=(0, 0, 11, 0.6); $c$=(-10,-10,-6.5, 0.7); $A$=(-200,-100,-170, 15); $X^{0}$=(1, 0,-0.5,-1); $Y^{0}$=(0, 0.5, 1.5, 1)
We first use OpenMM to generate the data and then analyze it user MasterMSM.
End of explanation
import mueller
import simtk.openmm as mm
Explanation: Data generation
First of all we import the mueller module, which contains a class (borrowed from Robert Mc Gibbon) that runs the simulations using OpenMM.
End of explanation
from simtk.unit import kelvin, picosecond, femtosecond, nanometer, dalton
# each particle is totally independent, propagating under the same potential
mass = 1.0*dalton
temperature = 750 * kelvin
friction = 50/picosecond
timestep = 10.0*femtosecond
# Choose starting conformations uniform on the grid between (-1.5, -0.2) and (1.2, 2)
nParticles = 4
startingPositions = (np.random.rand(nParticles, 3)*np.array([2.7, 1.8, 1])) \
            + np.array([-1.5, -0.2, 0])
    
system = mm.System()
mullerforce = mueller.MullerForce()
for i in range(nParticles):
    system.addParticle(mass)
    mullerforce.addParticle(i, [])
system.addForce(mullerforce)
integrator = mm.LangevinIntegrator(temperature, friction, timestep)
context = mm.Context(system, integrator)
context.setPositions(startingPositions)
context.setVelocitiesToTemperature(temperature)
import cProfile, pstats, io
from pstats import SortKey
# Run the simulation
traj = []
for i in range(int(1e7)):
    coords = context.getState(getPositions=True).getPositions(asNumpy=True).value_in_unit(nanometer)
    integrator.step(50)
    traj.append(coords)
for i in range(nParticles):
    data = np.vstack([x[i] for x in traj])
    h5file = "data/mueller_test_n%i.h5"%i
    with h5py.File(h5file, "w") as hf:
        hf.create_dataset("data", data=data)
Explanation: First we set a number of parameters for the run.
End of explanation
import h5py
i = 0
h5file = "data/mueller_test_n%g.h5"%i
f = h5py.File(h5file, 'r')
data = np.array(f['data'])
f.close()
fig, ax = plt.subplots(figsize=(4,4))
mullerforce.plot(ax=ax, cmap='binary_r', levels=np.arange(-150,50,10))
for i in range(nParticles):
    h5file = "data/mueller_test_n%g.h5"%i
    f = h5py.File(h5file, 'r')
    data = np.array(f['data'])
    f.close()
    ax.plot(data[:-1:10,0], data[:-1:10,1], lw=0.1)
ax.set_xlabel('x')
ax.set_ylabel('y')
plt.tight_layout()
Explanation: Data analysis
Download brownian trajectory from here, and read the latter by using trajectory from COSSIO generator.
End of explanation
def discretize(data):
    # Discretize trajectory and plot onto the mueller potential
    discret = np.zeros(data.shape[0])
    statistic, x_edge, y_edge, binnumber = \
            scipy.stats.binned_statistic_2d(data[:,0], data[:,1],\
            discret, 'count', \
            bins = [np.linspace(-1.5, 1, 6), np.linspace(-0.5, 2, 6)], \
                        expand_binnumbers=False)
    return binnumber
data.shape
binned_trajs = []
fig, ax = plt.subplots(3,1, figsize=(8,4), sharex=True)
for i in range(nParticles):
    h5file = "data/mueller_test_n%g.h5"%i
    f = h5py.File(h5file, 'r')
    data = np.array(f['data'])
    f.close()
    print (data.shape)
    binnumber = discretize(data)
    binned_trajs.append(list(binnumber))
    ax[0].plot(data[:,0], lw=0.2)
    ax[1].plot(data[:,1], lw=0.2)
    ax[2].plot(binnumber, lw=0.2)
#ax[0].set_xlim(0, 500000)
ax[0].set_ylabel('x')
ax[0].set_ylabel('y')
plt.tight_layout()
data[:,1]
len(binnumber), len(data)
binnumber
Explanation: Trajectory analysis and assignment
We discretize the trajectory by a 2D binning, so the mueller potential is divided in microstates (plot 1). In consequence, both x and y coordinates merge onto a s coordinate (plot 2).
End of explanation
# MASTER EQUATION MODEL
from mastermsm.trajectory import traj
trajectories = [traj.TimeSeries(distraj=binned_trajs[i], dt=1) \
           for i in range(nParticles)]
[trajectories[i].find_keys() for i in range(nParticles)]
from mastermsm.msm import msm
msm_2D = msm.SuperMSM(trajectories,sym=False)
msm_2D.convergence_test(time=[1, 2, 5, 10, 20, 50], error=False)
fig, ax = plt.subplots()
for i in range(0,6):
    tau_vs_lagt = np.array([[x, msm_2D.msms[x].tauT[i]] \
               for x in sorted(msm_2D.msms.keys())])
    ax.plot(tau_vs_lagt[:,0], tau_vs_lagt[:,1], marker='o', ms=7)
ax.set_xlabel(r'$\Delta$t ', fontsize=16)
ax.set_ylabel(r'$\tau_i$', fontsize=16)
ax.set_xscale('log')
ax.set_yscale('log')
ax.fill_between(10**np.arange(-1,4,0.2), 1e-1, 10**np.arange(-1,4,0.2), facecolor='lightgray', alpha=0.5)
ax.set_ylim(1e-1,1000)
ax.set_xlim(0.7,1400)
plt.tight_layout()
Explanation: MSM: convergence test
First, we build an instance of TimeSeries class to find the strongly connected states corresponding to the assignment shown above. Then, we use the SuperMSM class to study the dependence of the relaxation times, $\tau_i$, on the lag time $\Delta t$ adopted to build the transition matrix.
End of explanation
lagt=10
msm_2D.do_msm(lagt)
msm_2D.msms[lagt].do_trans(evecs=True)
# VALIDATION - AUTOCORRELATION FUNCTION (acf)
acf = msm_2D.msms[lagt].acf_mode()
time = np.arange(len(acf[1]))*msm_2D.data[0].dt
fig, ax = plt.subplots()
ax.plot(time, acf[1], 'o', label='ACF(left-eig)')
ax.plot(time,np.exp(-time*1./msm_2D.msms[lagt].tauT[0]), \
            label='exp(-t/%5.3f)' % msm_2D.msms[lagt].tauT[0])
ax.set_ylim(0,1)
ax.set_xlim(0,1300)
ax.set_xlabel('Time')
ax.set_ylabel('C$_{11}$(t)')
ax.legend()
plt.show
# Spectrum of relaxation times
fig, ax = plt.subplots()
ax.plot(range(1,len(msm_2D.msms[lagt].tauT)+1), \
            np.log(msm_2D.msms[lagt].tauT), 'x-', markersize=12, linewidth=4)
ax.set_xlabel('Eigenvalue')
ax.set_ylabel(r'log $\tau_i$')
plt.tight_layout()
plt.show
Explanation: MSM: validation
Once a proper lag time is chosen, we generate a MSM and plot the autocorrelation corresponding to the left eigenvectors (Buchete et al. JPCB 2008), as well as the spectrum of (logarithm of) relaxation times.
End of explanation
fig, ax = plt.subplots(figsize=(10,4))
ax.plot(msm_2D.msms[lagt].rvecsT[:,1])
ax.fill_between(range(len(msm_2D.msms[lagt].rvecsT[:,1])), 0, \
                msm_2D.msms[lagt].rvecsT[:,1], \
                where=msm_2D.msms[lagt].rvecsT[:,1]>0,\
                facecolor='c', interpolate=True,alpha=.4)
ax.fill_between(range(len(msm_2D.msms[lagt].rvecsT[:,1])), 0, \
                msm_2D.msms[lagt].rvecsT[:,1], \
                where=msm_2D.msms[lagt].rvecsT[:,1]<0,\
                facecolor='g', interpolate=True,alpha=.4)
ax.set_ylabel("$\Psi^R_1$")
plt.show
# 2D projection
binnumber2 = np.zeros(len(data[:,1]))
for i in range(len(binnumber)-1):
    j = binnumber[i]
    for k in range(len(msm_2D.msms[lagt].keep_keys)-1):
        if j==msm_2D.msms[lagt].keep_keys[k]:
            binnumber2[i] = np.log(msm_2D.msms[lagt].peqT[k])
            break
y = np.linspace(-0.5,2.,200)
x = np.linspace(-2.,1.1,200)
M = np.zeros((len(x), len(y)))
lines = []
for i, x_i in enumerate(x):
    for j, y_i in enumerate(y):
        aux = fmueller(x_i,y_i)
        if aux<100: M[i, j] = aux
X, Y = np.meshgrid(x,y)
fig, ax = plt.subplots()
levels = np.arange(-130.0, 40.0, 20)
CS = ax.contour(X, Y, M.T, levels)
ax.clabel(CS, inline=0.9, fontsize=6)
cmaps = OrderedDict()
cmaps['Diverging'] = ['PiYG', 'PRGn', 'BrBG', 'PuOr', 'RdGy', 'RdBu',
            'RdYlBu', 'RdYlGn', 'Spectral', 'coolwarm', 'bwr', 'seismic']
plt.scatter(data[:,1],data[:,2],c=binnumber2,cmap=plt.get_cmap('seismic'))
plt.colorbar()
plt.title('log of equilibrium distribution')
plt.show
Explanation: The first mode captured by $\lambda_1$ is significantly slower than the others. That mode, which is described by the right eigenvector $\psi^R_1$ as the transition between the two main minima. Next plots correspond to the eigenvector and its projection on the mueller potential.
End of explanation
import hdbscan
X = data[:,[1,2]]
hb = hdbscan.HDBSCAN(min_cluster_size=200).fit(X)
labels = hb.labels_
for i, x_i in enumerate(labels):
    if hb.probabilities_[i] < 0.1:
        labels[i] = -1
n_clusters_ = len(set(labels)) - (1 if -1 in labels else 0)
colors = ['royalblue', 'maroon', 'forestgreen', 'mediumorchid', \
        'tan', 'deeppink', 'olive', 'goldenrod', 'lightcyan', 'lightgray']
vectorizer = np.vectorize(lambda x: colors[x % len(colors)])
plt.scatter(X[:,0],X[:,1], c=vectorizer(labels))
plt.title('Estimated number of clusters: %d' % n_clusters_)
plt.show
hb.condensed_tree_.plot()
plt.title('tree')
plt.show
Explanation: HDBSCAN clustering
Since the 2D binning used above in order to assign discrete states does not take into account the potential form, in the following a density-based clustering is employed instead. Namely, the HDBSCAN algorithm is suitable to deal with any shape clusters with different densities. The latter must be installed from here before going on.
Discretization
Clusters are built from trajectory data. Then, we refine noise and analyze if the resulting condensed tree plot makes sense for our system.
End of explanation
# remove noise (label=-1) to obtain a trajectory for MasterMSM
last = labels[0]
for i, x_i in enumerate(labels):
    if x_i == -1:
        labels[i] = last
    else:
        last = x_i
labels[0] = labels[1]
cluster = traj.TimeSeries(distraj=list(labels), dt=1)
cluster.find_keys()
cluster.keys.sort()
msm_cluster = msm.SuperMSM([cluster],sym=True)
lagt=5 # lag time
msm_cluster.do_msm(lagt)
msm_cluster.msms[lagt].do_trans(evecs=True)
fig, ax = plt.subplots()
ax.plot(range(1,len(msm_cluster.msms[lagt].tauT)+1), \
        np.log(msm_cluster.msms[lagt].tauT), 'o-', label='HDBSCAN')
# Note for msm_2D we chose lagt=10
ax.plot(range(1,len(msm_cluster.msms[lagt].tauT)+1), \
        np.log(msm_2D.msms[10].tauT[:len(msm_cluster.msms[lagt].tauT)]), \
        'x-', label='2D binning')
ax.set_xlabel('Eigenvalue')
ax.set_ylabel(r'log $\tau_i$')
plt.tight_layout()
plt.legend()
plt.title('log_mueller_relax_times.png')
plt.show
Explanation: MSM validation
We compare the relaxation times obtained from the transition matrix with those corresponding to the 2D binning shown above. Obviously, since there are three clusters in HDBSCAN, we obtain two possible transitions.
End of explanation
msm_cluster.do_lbrate()
msm_cluster.msms[lagt].do_rate(method='MLPB', evecs=True, init=msm_cluster.lbrate)
fig, ax = plt.subplots(2,1, figsize=(6,4), sharex=False)
plt.subplots_adjust(wspace=0.9, hspace=0.9)
ax[0].plot(msm_cluster.msms[lagt].rvecsK[:,1])
ax[0].fill_between(range(len(msm_cluster.msms[lagt].rvecsK[:,1])), \
        0, msm_cluster.msms[lagt].rvecsK[:,1], alpha=.4)
ax[1].plot(msm_cluster.msms[lagt].lvecsK[:,1])
ax[1].axhline(0,0,25, c='k', ls='--', lw=1)
ax[1].fill_between(range(len(msm_cluster.msms[lagt].lvecsK[:,1])), \
        msm_cluster.msms[lagt].lvecsK[:,1], 0, alpha=.4)
ax[0].set_ylabel("$\Psi^R_1$")
ax[1].set_ylabel("$\Psi^L_1$")
plt.tight_layout(h_pad=0)
ax[0].set_title('1st_eigenvectors_R_L.png')
plt.show
Explanation: RATE matrix analysis
The simplicity of this discretization with only three states allows us to employ the Maximum-Likelihood Propagator-Based (MLPB) method in order to compute rate matrix and corresponding eigenvectors.
End of explanation
plt.figure()
plt.imshow(np.log10(msm_cluster.msms[lagt].count), \
    interpolation='none', cmap='viridis_r', origin='lower')
plt.ylabel('$\it{j}$')
plt.xlabel('$\it{i}$')
plt.colorbar()
plt.title('Count matrix (log), $\mathbf{N}$')
plt.show
plt.imshow(np.log10(msm_cluster.msms[lagt].trans), interpolation='none', \
    cmap='viridis_r', origin='lower')
plt.ylabel('$\it{j}$')
plt.xlabel('$\it{i}$')
plt.colorbar()
plt.title('Transition matrix (log), $\mathbf{T}$')
plt.show
Explanation: COUNT and TRANSITION matrices
Count and transition matrices give us an intuition of the dynamics for this simple system. Remind the former is strongly dependent on the lag time.
End of explanation
<END_TASK>