diff --git a/CHANGES b/CHANGES index b95aa851..c64a912b 100644 --- a/CHANGES +++ b/CHANGES @@ -21,6 +21,7 @@ The rules for this file: Enhancements - Allow the overlap matrix of the MBAR estimator to be plotted. (PR #107) + - Allow the dhdl of the TI estimator to be plotted. (PR #110) Deprecations diff --git a/docs/images/dhdl_TI.png b/docs/images/dhdl_TI.png new file mode 100644 index 00000000..1e85be39 Binary files /dev/null and b/docs/images/dhdl_TI.png differ diff --git a/docs/visualisation.rst b/docs/visualisation.rst index e7a03336..8d67d3b3 100644 --- a/docs/visualisation.rst +++ b/docs/visualisation.rst @@ -4,6 +4,14 @@ It is quite often that the user want to visualise the results to gain confidence on the computed free energy. **alchemlyb** provides various visualisation tools to help user to judge the estimate. +.. currentmodule:: alchemlyb.visualisation + +.. autosummary:: + :toctree: visualisation + + plot_mbar_overlap_matrix + plot_ti_dhdl + .. _plot_overlap_matrix: Overlap Matrix of the MBAR @@ -34,6 +42,38 @@ Will give a plot looks like this .. image:: images/O_MBAR.png +.. _plot_TI_dhdl: + +dhdl Plot of the TI +------------------- +In order for the :class:`~alchemlyb.estimators.TI` estimator to work reliably, +the change in the dhdl between lambda state 0 and lambda state 1 should be +adequately sampled. The function :func:`~alchemlyb.visualisation.plot_ti_dhdl` +can be used to assess the change of the dhdl across the lambda states. + +More than one :class:`~alchemlyb.estimators.TI` estimators can be plotted +together as well. :: + + >>> import pandas as pd + >>> from alchemtest.gmx import load_benzene + >>> from alchemlyb.parsing.gmx import extract_dHdl + >>> from alchemlyb.estimators import TI + + >>> bz = load_benzene().data + >>> dHdl_coul = pd.concat([extract_dHdl(xvg, T=300) for xvg in bz['Coulomb']]) + >>> ti_coul = TI().fit(dHdl_coul) + >>> dHdl_vdw = pd.concat([extract_dHdl(xvg, T=300) for xvg in bz['VDW']]) + >>> ti_vdw = TI().fit(dHdl_vdw) + + >>> from alchemlyb.visualisation import plot_ti_dhdl + >>> ax = plot_ti_dhdl([ti_coul, ti_vdw], labels=['Coul', 'VDW'], colors=['r', 'g']) + >>> ax.figure.savefig('dhdl_TI.pdf') + + +Will give a plot looks like this + +.. image:: images/dhdl_TI.png + .. [Klimovich2015] Klimovich, P.V., Shirts, M.R. & Mobley, D.L. Guidelines for the analysis of free energy calculations. J Comput Aided Mol Des 29, 397–411 (2015). https://doi.org/10.1007/s10822-015-9840-9 diff --git a/docs/visualisation/alchemlyb.visualisation.plot_mbar_overlap_matrix.rst b/docs/visualisation/alchemlyb.visualisation.plot_mbar_overlap_matrix.rst index 5323d1a9..a4de37da 100644 --- a/docs/visualisation/alchemlyb.visualisation.plot_mbar_overlap_matrix.rst +++ b/docs/visualisation/alchemlyb.visualisation.plot_mbar_overlap_matrix.rst @@ -1,4 +1,4 @@ -.. _plot_overlap: +.. _visualisation_plot_mbar_overlap_matrix: Plot Overlap Matrix from MBAR ============================= diff --git a/docs/visualisation/alchemlyb.visualisation.plot_ti_dhdl.rst b/docs/visualisation/alchemlyb.visualisation.plot_ti_dhdl.rst new file mode 100644 index 00000000..d4247515 --- /dev/null +++ b/docs/visualisation/alchemlyb.visualisation.plot_ti_dhdl.rst @@ -0,0 +1,22 @@ +.. _visualisation_plot_mbar_plot_ti_dhdl: + +Plot dhdl from TI +================= + +The function :func:`~alchemlyb.visualisation.plot_ti_dhdl` allows +the user to plot the dhdl from :class:`~alchemlyb.estimators.TI` estimator. +Several :class:`~alchemlyb.estimators.TI` estimators could be passed to the +function to give a concerted picture of the whole alchemical transformation. +When custom labels are desirable, the user could pass a list of strings to the +*labels* for labelling each alchemical transformation differently. The color of +each alchemical transformation could also be set by passing a list of color +string to the *colors*. The unit in the y axis could be labelled to other units +by setting *units*, which by default is kcal/mol. The user can pass +:class:`matplotlib.axes.Axes` into the function to have the dhdl drawn on a +specific axes. + +Please check :ref:`How to plot TI dhdl ` for usage. + +API Reference +------------- +.. autofunction:: alchemlyb.visualisation.plot_ti_dhdl \ No newline at end of file diff --git a/src/alchemlyb/estimators/mbar_.py b/src/alchemlyb/estimators/mbar_.py index dd8127ae..975d177b 100644 --- a/src/alchemlyb/estimators/mbar_.py +++ b/src/alchemlyb/estimators/mbar_.py @@ -45,9 +45,6 @@ class MBAR(BaseEstimator): states_ : list Lambda states for which free energy differences were obtained. - overlap_matrix: numpy.matrix - The overlap matrix. - """ def __init__(self, maximum_iterations=10000, relative_tolerance=1.0e-7, diff --git a/src/alchemlyb/estimators/ti_.py b/src/alchemlyb/estimators/ti_.py index 7f25dbe3..383341c9 100644 --- a/src/alchemlyb/estimators/ti_.py +++ b/src/alchemlyb/estimators/ti_.py @@ -26,6 +26,9 @@ class TI(BaseEstimator): states_ : list Lambda states for which free energy differences were obtained. + dhdl : DataFrame + The estimated dhdl of each state. + """ def __init__(self, verbose=False): @@ -92,6 +95,7 @@ def fit(self, dHdl): self.delta_f_ = pd.DataFrame(adelta - adelta.T, columns=means.index.values, index=means.index.values) + self.dhdl = means # yield standard deviation d_delta_f_ between each state self.d_delta_f_ = pd.DataFrame(np.sqrt(ad_delta + ad_delta.T), diff --git a/src/alchemlyb/tests/test_visualisation.py b/src/alchemlyb/tests/test_visualisation.py index d482caad..06808ea6 100644 --- a/src/alchemlyb/tests/test_visualisation.py +++ b/src/alchemlyb/tests/test_visualisation.py @@ -1,10 +1,13 @@ import matplotlib +import matplotlib.pyplot as plt import pandas as pd +import numpy as np from alchemtest.gmx import load_benzene -from alchemlyb.parsing.gmx import extract_u_nk -from alchemlyb.estimators import MBAR +from alchemlyb.parsing.gmx import extract_u_nk, extract_dHdl +from alchemlyb.estimators import MBAR, TI from alchemlyb.visualisation.mbar_matrix import plot_mbar_overlap_matrix +from alchemlyb.visualisation.ti_dhdl import plot_ti_dhdl def test_plot_mbar_omatrix(): '''Just test if the plot runs''' @@ -25,3 +28,29 @@ def test_plot_mbar_omatrix(): assert isinstance(plot_mbar_overlap_matrix(overlap_maxtrix), matplotlib.axes.Axes) +def test_plot_ti_dhdl(): + '''Just test if the plot runs''' + bz = load_benzene().data + dHdl_coul = pd.concat([extract_dHdl(xvg, T=300) for xvg in bz['Coulomb']]) + ti_coul = TI() + ti_coul.fit(dHdl_coul) + assert isinstance(plot_ti_dhdl(ti_coul), + matplotlib.axes.Axes) + fig, ax = plt.subplots(figsize=(8, 6)) + assert isinstance(plot_ti_dhdl(ti_coul, ax=ax), + matplotlib.axes.Axes) + assert isinstance(plot_ti_dhdl(ti_coul, labels=['Coul']), + matplotlib.axes.Axes) + assert isinstance(plot_ti_dhdl(ti_coul, labels=['Coul'], colors=['r']), + matplotlib.axes.Axes) + dHdl_vdw = pd.concat([extract_dHdl(xvg, T=300) for xvg in bz['VDW']]) + ti_vdw = TI().fit(dHdl_vdw) + assert isinstance(plot_ti_dhdl([ti_coul, ti_vdw]), + matplotlib.axes.Axes) + ti_coul.dhdl = pd.DataFrame.from_dict( + {'fep': range(100)}, + orient='index', + columns=np.arange(100)/100).T + assert isinstance(plot_ti_dhdl(ti_coul), + matplotlib.axes.Axes) + diff --git a/src/alchemlyb/visualisation/__init__.py b/src/alchemlyb/visualisation/__init__.py index 4fc73051..5cf2c3f9 100644 --- a/src/alchemlyb/visualisation/__init__.py +++ b/src/alchemlyb/visualisation/__init__.py @@ -1 +1,2 @@ -from .mbar_matrix import plot_mbar_overlap_matrix \ No newline at end of file +from .mbar_matrix import plot_mbar_overlap_matrix +from .ti_dhdl import plot_ti_dhdl \ No newline at end of file diff --git a/src/alchemlyb/visualisation/ti_dhdl.py b/src/alchemlyb/visualisation/ti_dhdl.py new file mode 100644 index 00000000..7b53217b --- /dev/null +++ b/src/alchemlyb/visualisation/ti_dhdl.py @@ -0,0 +1,175 @@ +"""Functions for Plotting the dhdl for the TI estimator. + +To assess the quality of the TI estimator, the dhdl from lambda state 0 +to lambda state 1 can plotted to assess if the change in dhdl is sampled +thoroughly. + +The code for producing the dhdl plot is modified based on +: `Alchemical Analysis `_. + +""" +from __future__ import division + +import matplotlib.pyplot as plt +from matplotlib.font_manager import FontProperties as FP +import numpy as np + +def plot_ti_dhdl(dhdl_data, labels=None, colors=None, units='kcal/mol', ax=None): + '''Plot the dhdl of TI. + + Parameters + ---------- + dhdl_data : :class:`~alchemlyb.estimators.TI` or list + One or more :class:`~alchemlyb.estimators.TI` estimator, where the + dhdl value will be taken from. + labels : List + list of labels for labelling all the alchemical transformations. + colors : List + list of colors for plotting all the alchemical transformations. + Default: ['r', 'g', '#7F38EC', '#9F000F', 'b', 'y'] + units : str + The unit of the estimate. Default: 'kcal/mol' + ax : matplotlib.axes.Axes + Matplotlib axes object where the plot will be drawn on. If ax=None, + a new axes will be generated. + + Returns + ------- + matplotlib.axes.Axes + An axes with the TI dhdl drawn. + + Notes + ----- + The code is taken and modified from + : `Alchemical Analysis `_ + + ''' + # Make it into a list + try: + len(dhdl_data) + except TypeError: + dhdl_data = [dhdl_data, ] + + if ax is None: + fig, ax = plt.subplots(figsize=(8, 6)) + + ax.spines['bottom'].set_position('zero') + ax.spines['top'].set_color('none') + ax.spines['right'].set_color('none') + ax.xaxis.set_ticks_position('bottom') + ax.yaxis.set_ticks_position('left') + + for k, spine in ax.spines.items(): + spine.set_zorder(12.2) + + # Make level names + if labels is None: + lv_names2 = [] + for dhdl in dhdl_data: + # Assume that the dhdl has only one columns + lv_names2.append(dhdl.dhdl.columns.values[0].capitalize()) + else: + if len(labels) == len(dhdl_data): + lv_names2 = labels + else: # pragma: no cover + raise ValueError( + 'Length of labels ({}) should be the same as the number of data ({})'.format( + len(labels), len(dhdl_data))) + + if colors is None: + colors = ['r', 'g', '#7F38EC', '#9F000F', 'b', 'y'] + else: + if len(colors) >= len(dhdl_data): + pass + else: # pragma: no cover + raise ValueError( + 'Number of colors ({}) should be larger than the number of data ({})'.format( + len(labels), len(dhdl_data))) + + # Get the real data out + xs, ndx, dx = [0], 0, 0.001 + min_y, max_y = 0, 0 + for dhdl in dhdl_data: + x = dhdl.dhdl.index.values + y = dhdl.dhdl.values.ravel() + + min_y = min(y.min(), min_y) + max_y = max(y.max(), max_y) + + for i in range(len(x) - 1): + if i % 2 == 0: + ax.fill_between(x[i:i + 2] + ndx, 0, y[i:i + 2], + color=colors[ndx], alpha=1.0) + else: + ax.fill_between(x[i:i + 2] + ndx, 0, y[i:i + 2], + color=colors[ndx], alpha=0.5) + + xlegend = [-100 * wnum for wnum in range(len(lv_names2))] + ax.plot(xlegend, [0 * wnum for wnum in xlegend], ls='-', + color=colors[ndx], + label=lv_names2[ndx]) + xs += (x + ndx).tolist()[1:] + ndx += 1 + + # Make sure the tick labels are not overcrowded. + xs = np.array(xs) + dl_mat = np.array([xs - i for i in xs]) + ri = range(len(xs)) + + def getInd(r=ri, z=[0]): + primo = r[0] + min_dl = ndx * 0.02 * 2 ** (primo > 10) + if dl_mat[primo].max() < min_dl: + return z + for i in r: # pragma: no cover + for j in range(len(xs)): + if dl_mat[i, j] > min_dl: + z.append(j) + return getInd(ri[j:], z) + + xt = [] + for i in range(len(xs)): + if i in getInd(): + xt.append(i) + else: + xt.append('') + + plt.xticks(xs[1:], xt[1:], fontsize=10) + ax.yaxis.label.set_size(10) + + # Remove the abscissa ticks and set up the axes limits. + for tick in ax.get_xticklines(): + tick.set_visible(False) + ax.set_xlim(0, ndx) + min_y *= 1.01 + max_y *= 1.01 + + # Modified so that the x label won't conflict with the lambda label + min_y -= (max_y-min_y)*0.1 + + ax.set_ylim(min_y, max_y) + + for i, j in zip(xs[1:], xt[1:]): + ax.annotate( + ('%.2f' % (i - 1.0 if i > 1.0 else i) if not j == '' else ''), + xy=(i, 0), xytext=(i, 0.01), size=10, rotation=90, + textcoords=('data', 'axes fraction'), va='bottom', ha='center', + color='#151B54') + if ndx > 1: + lenticks = len(ax.get_ymajorticklabels()) - 1 + if min_y < 0: lenticks -= 1 + if lenticks < 5: # pragma: no cover + from matplotlib.ticker import AutoMinorLocator as AML + ax.yaxis.set_minor_locator(AML()) + ax.grid(which='both', color='w', lw=0.25, axis='y', zorder=12) + ax.set_ylabel( + r'$\mathrm{\langle{\frac{ \partial U } { \partial \lambda }}\rangle_{\lambda}\/%s}$' % units, + fontsize=20, color='#151B54') + ax.annotate('$\mathit{\lambda}$', xy=(0, 0), xytext=(0.5, -0.05), size=18, + textcoords='axes fraction', va='top', ha='center', + color='#151B54') + lege = ax.legend(prop=FP(size=14), frameon=False, loc=1) + for l in lege.legendHandles: + l.set_linewidth(10) + return ax +