dhdl plot for the TI (#110)

* partially addresses #73
* add dhdl plot from alchemical-analysis (+ improvements)
* add tests and use `# pragma: no cover` for parts of plotting code that are difficult to test
* update CHANGES

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
 

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

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
 =============================

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 <plot_TI_dhdl>` for usage.
+
+API Reference
+-------------
+.. autofunction:: alchemlyb.visualisation.plot_ti_dhdl

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,

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),

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)
+

src/alchemlyb/visualisation/__init__.py

@@ -1 +1,2 @@
-from .mbar_matrix import plot_mbar_overlap_matrix
+from .mbar_matrix import plot_mbar_overlap_matrix
+from .ti_dhdl import plot_ti_dhdl

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 <https://github.com/MobleyLab/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 <https://github.com/MobleyLab/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
+