Nonequilibrium Switching Free Energy Analysis

This notebook is useful for examining the results of nonequilibrium switching calculations performed with the Perses tool. It enables examination of work traces as well as trajectories.

Import required libraries


In [1]:
import numpy as np
import nglview
from bokeh.plotting import figure, output_notebook, show
from bokeh.layouts import row, column
import simtk.unit as unit
import mdtraj as md
import plotting_tools
output_notebook()


Loading BokehJS ...

Set the appropriate variables to the directories where the output files are stored


In [2]:
trajectory_directory = "/Users/grinawap/solvent_test_4"
trajectory_prefix = "cdk02" #these were set by the yaml file

In [3]:
data_reader = plotting_tools.NonequilibriumSwitchingAnalysis(trajectory_directory, trajectory_prefix)
cum_work_shape = np.shape(data_reader.cumulative_work)
n_steps_per_protocol = cum_work_shape[2]
protocol_steps = np.linspace(0, 1, n_steps_per_protocol)
n_protocols = cum_work_shape[1]

Now we can plot the forward and reverse work traces

These are the cumulative work going from lambda=0 to lambda=1, and the reverse


In [10]:
#create a convenience function to add lines to a bokeh plot
def plot_line(fig, x, y):
    fig.line(x, y)

def plot_work_values(protocol_steps, data_reader, n_protocols):
    """
    Convenience function to plot work values from nonequilibrium switching calculations.
    
    Arguments
    ---------
    protocol_steps : [n_steps] array
        Indexes the lambda value at each point in the corresponding work trajectory
    data_reader : NonequilibriumSwitchingAnalysis object
        object that contains the switching data
    n_protocols : int
        The total number of protocols in each direction (should be symmetric)
    """
    x_axis_label = "lambda value"
    y_axis_label = "Work (kT)"
    
    work_figure = figure(title="Nonequilibrium Switching Work", x_axis_label=x_axis_label,
                        y_axis_label=y_axis_label)
    
    for protocol_index in range(n_protocols):
        work_figure.line(protocol_steps, data_reader.cumulative_work[0, protocol_index, :], line_color="green")
        work_figure.line(protocol_steps, data_reader.cumulative_work[1, protocol_index, :], line_color="red")
    
    show(work_figure)

In [11]:
plot_work_values(protocol_steps, data_reader, n_protocols)


What should I be looking for in the work traces?

Is there a place in the trajectory where the work suddenly becomes anomalously high? This might indicate either a bad protocol or a bug in the code

In general, these plots are giving an illustration of how fast the system is changing as lambda is changed. The faster that change is, the more unfavorable the variance of the ultimate free energy calculation

Let's take a closer look at the trajectories

See something weird? Sometimes it helps to just visualize the trajectory of the atoms. We can do that fairly straightforwardly:


In [30]:
nonequilibrium_trajectory = data_reader.get_nonequilibrium_trajectory("forward", 0)
view = nglview.NGLWidget()
view.add_trajectory(nonequilibrium_trajectory)
view.add_ball_and_stick()

In [31]:
view



In [ ]: