westpa 2022.12__cp312-cp312-macosx_10_13_x86_64.whl

westpa/cli/tools/plothist.py

@@ -0,0 +1,706 @@
+import logging
+import os
+import re
+
+import h5py
+import numpy as np
+
+import matplotlib
+from matplotlib import pyplot
+from matplotlib.image import NonUniformImage
+
+from westpa.tools import WESTMasterCommand, WESTSubcommand
+from westpa.core import h5io, textio
+from westpa.fasthist import normhistnd
+from westpa.core.extloader import get_object
+
+log = logging.getLogger('plothist')
+
+# Suppress divide-by-zero in log
+np.seterr(divide='ignore', invalid='ignore')
+
+
+def sum_except_along(array, axes):
+    '''Reduce the given array by addition over all axes except those listed in the scalar or
+    iterable ``axes``'''
+
+    try:
+        iter(axes)
+    except TypeError:
+        axes = [axes]
+
+    kept = set(axes)
+    summed = list(set(range(array.ndim)) - kept)
+
+    # Reorder axes so that the kept axes are first, and in the order they
+    # were given
+    array = np.transpose(array, list(axes) + summed).copy()
+
+    # Now, the last len(summed) axes are summed over
+    for _ in range(len(summed)):
+        array = np.add.reduce(array, axis=-1)
+
+    return array
+
+
+class PlotHistBase(WESTSubcommand):
+    def __init__(self, parent):
+        super().__init__(parent)
+
+        self.input_arg_group = None
+        self.output_arg_group = None
+
+        self.input_h5 = None
+        self.opmode = None
+        self.plotscale = None
+        self.enerzero = None
+        self.plotrange = None
+        self.plottitle = None
+        self.postprocess_function = None
+        self.plot_contour = None
+
+        # Iteration range for average/evolution
+        self.avail_iter_start = None
+        self.avail_iter_stop = None
+        self.avail_iter_step = None
+        self.iter_start = None
+        self.iter_stop = None
+        self.iter_step = None
+
+        # Iteration for single point
+        self.n_iter = None
+
+        # An array of dicts describing what dimensions to work with and
+        # what their ranges should be for the plots.
+        self.dimensions = []
+
+        self.plot_output_filename = None
+        self.text_output_filename = None
+        self.hdf5_output_filename = None
+
+    def add_args(self, parser):
+        igroup = self.input_arg_group = parser.add_argument_group('input options')
+        igroup.add_argument('input', help='HDF5 file containing histogram data')
+        igroup.add_argument(
+            'firstdim',
+            nargs='?',
+            metavar='DIMENSION',
+            help='''Plot for the given DIMENSION, specified as INT[:[LB,UB]:LABEL], where
+                            INT is a zero-based integer identifying the dimension in the histogram,
+                            LB and UB are lower and upper bounds for plotting, and LABEL is the label for
+                            the plot axis. (Default: dimension 0, full range.)''',
+        )
+
+        ogroup = self.output_arg_group = parser.add_argument_group('output options')
+        ogroup.add_argument(
+            '-o',
+            '--output',
+            '--plot-output',
+            dest='plot_output',
+            default='hist.pdf',
+            metavar='PLOT_OUTPUT',
+            help='''Store plot as PLOT_OUTPUT. This may be set to an empty string
+                            (e.g. --plot-output='') to suppress plotting entirely. The output
+                            format is determined by filename extension (and thus defaults to PDF).
+                            Default: "%(default)s".''',
+        )
+        ogroup.add_argument('--hdf5-output', help='''Store plot data in the HDF5 file HDF5_OUTPUT.''')
+        ogroup.add_argument(
+            '--plot-contour',
+            dest='plot_contour',
+            action='store_const',
+            const=True,
+            default=False,
+            help='''Determines whether or not to superimpose a contour plot over the heatmap for 2D objects.''',
+        )
+
+        pgroup = parser.add_argument_group('plot options')
+        pmgroup = pgroup.add_mutually_exclusive_group()
+        pgroup.add_argument('--title', dest='title', help='Include TITLE as the top-of-graph title')
+        pmgroup.add_argument(
+            '--linear', dest='plotscale', action='store_const', const='linear', help='Plot the histogram on a linear scale.'
+        )
+        pmgroup.add_argument(
+            '--energy',
+            dest='plotscale',
+            action='store_const',
+            const='energy',
+            help='Plot the histogram on an inverted natural log scale, corresponding to (free) energy (default).',
+        )
+        pmgroup.add_argument(
+            '--zero-energy',
+            dest='enerzero',
+            metavar='E',
+            default='min',
+            help='Set the zero of energy to E, which may be a scalar, "min" or "max"',
+        )
+        pmgroup.add_argument(
+            '--log10', dest='plotscale', action='store_const', const='log10', help='Plot the histogram on a base-10 log scale.'
+        )
+        pgroup.add_argument(
+            '--range',
+            help='''Plot histogram ordinates over the given RANGE, specified as "LB,UB",
+                            where LB and UB are the lower and upper bounds, respectively. For 1-D plots,
+                            this is the Y axis. For 2-D plots, this is the colorbar axis.
+                            (Default: full range.)''',
+        )
+        pgroup.add_argument(
+            '--postprocess-function',
+            help='''Names a function (as in module.function) that will be called just prior
+                                to saving the plot. The function will be called as ``postprocess(hist, midpoints, binbounds)``
+                                where ``hist`` is the histogram that was plotted, ``midpoints`` is the bin midpoints for
+                                each dimension, and ``binbounds`` is the bin boundaries for each dimension for 2-D plots,
+                                or None otherwise. The plot must be modified in place using the pyplot stateful interface.''',
+        )
+
+        parser.set_defaults(plotscale='energy')
+
+    def process_args(self, args):
+        self.plotscale = args.plotscale
+        self.input_h5 = h5py.File(args.input, 'r')
+        self.plot_output_filename = args.plot_output
+        self.hdf5_output_filename = args.hdf5_output
+        self.plot_contour = args.plot_contour
+
+        if args.title:
+            self.plottitle = args.title
+
+        if args.range:
+            self.plotrange = self.parse_range(args.range)
+
+        if args.firstdim:
+            self.dimensions.append(self.parse_dimspec(args.firstdim))
+
+        if not args.firstdim:
+            self.dimensions.append({'idim': 0, 'label': 'dimension 0'})
+
+        if args.enerzero:
+            lenerzero = args.enerzero.lower()
+            if lenerzero not in ('min', 'max'):
+                try:
+                    self.enerzero = float(args.enerzero)
+                except ValueError:
+                    raise ValueError('invalid energy zero point {!r}'.format(args.enerzero))
+            else:
+                self.enerzero = lenerzero
+        else:
+            self.enerzero = 'min'
+
+        self.avail_iter_start, self.avail_iter_stop = h5io.get_iter_range(self.input_h5['histograms'])
+        try:
+            self.avail_iter_step = h5io.get_iter_step(self.input_h5['histograms'])
+        except KeyError:
+            self.avail_iter_step = 1
+        log.info(
+            'HDF5 file {!r} contains data for iterations {} -- {} with a step of {}'.format(
+                args.input, self.avail_iter_start, self.avail_iter_stop, self.avail_iter_step
+            )
+        )
+        if args.postprocess_function:
+            self.postprocess_function = get_object(args.postprocess_function, path=['.'])
+
+    def parse_dimspec(self, dimspec):
+        dimdata = {}
+        match = re.match(r'([0-9]+)(?::(?:([^,]+),([^:,]+))?(?::(.*))?)?', dimspec)
+        if not match:
+            raise ValueError('invalid dimension specification {!r}'.format(dimspec))
+
+        (idim_txt, lb_txt, ub_txt, label) = match.groups()
+        try:
+            dimdata['idim'] = int(idim_txt)
+            if lb_txt:
+                dimdata['lb'] = float(lb_txt)
+            if ub_txt:
+                dimdata['ub'] = float(ub_txt)
+            if label:
+                dimdata['label'] = label
+            else:
+                dimdata['label'] = 'dimension {}'.format(dimdata['idim'])
+        except ValueError as e:
+            raise ValueError('invalid dimension specification {!r}: {!r}'.format(dimspec, e))
+        return dimdata
+
+    def parse_range(self, rangespec):
+        try:
+            (lbt, ubt) = rangespec.split(',')
+            return float(lbt), float(ubt)
+        except (ValueError, TypeError) as e:
+            raise ValueError('invalid range specification {!r}: {!r}'.format(rangespec, e))
+
+    def _ener_zero(self, hist):
+        hist = -np.log(hist)
+        if self.enerzero == 'min':
+            np.subtract(hist, hist.min(), out=hist, casting="unsafe")
+        elif self.enerzero == 'max':
+            np.subtract(hist, hist.max(), out=hist, casting="unsafe")
+        else:
+            np.subtract(hist, self.enerzero, out=hist, casting="unsafe")
+        return hist
+
+
+class PlotSupports2D(PlotHistBase):
+    def __init__(self, parent):
+        super().__init__(parent)
+
+    def add_args(self, parser):
+        self.input_arg_group.add_argument(
+            'seconddim',
+            nargs='?',
+            metavar='ADDTLDIM',
+            help='''For instantaneous/average plots, plot along the given additional
+                                          dimension, producing a color map.''',
+        )
+        self.output_arg_group.add_argument(
+            '--text-output',
+            help='''Store plot data in a text format at TEXT_OUTPUT. This option is
+                                           only valid for 1-D histograms. (Default: no text output.)''',
+        )
+
+    def process_args(self, args):
+        self.text_output_filename = args.text_output
+        if args.seconddim is not None:
+            self.dimensions.append(self.parse_dimspec(args.seconddim))
+
+    def _do_1d_output(self, hist, idim, midpoints):
+        enehist = self._ener_zero(hist)
+        log10hist = np.log10(hist)
+
+        if self.hdf5_output_filename:
+            with h5py.File(self.hdf5_output_filename, 'w') as output_h5:
+                h5io.stamp_creator_data(output_h5)
+                output_h5.attrs['source_data'] = os.path.abspath(self.input_h5.filename)
+                output_h5.attrs['source_dimension'] = idim
+                output_h5['midpoints'] = midpoints
+                output_h5['histogram'] = hist
+
+        if self.text_output_filename:
+            with textio.NumericTextOutputFormatter(self.text_output_filename) as output_file:
+                output_file.write_header('source data: {} dimension {}'.format(os.path.abspath(self.input_h5.filename), idim))
+                output_file.write_header('column 0: midpoint of bin')
+                output_file.write_header('column 1: probability in bin')
+                output_file.write_header('column 2: -ln P')
+                output_file.write_header('column 3: log10 P')
+                np.savetxt(output_file, np.column_stack([midpoints, hist, enehist, log10hist]))
+
+        if self.plot_output_filename:
+            if self.plotscale == 'energy':
+                plothist = enehist
+                label = r'$-\ln\,P(x)$'
+            elif self.plotscale == 'log10':
+                plothist = log10hist
+                label = r'$\log_{10}\ P(x)$'
+            else:
+                plothist = hist
+                label = r'$P(x)$'
+            pyplot.figure()
+            pyplot.plot(midpoints, plothist)
+            pyplot.xlim(self.dimensions[0].get('lb'), self.dimensions[0].get('ub'))
+            if self.plotrange:
+                pyplot.ylim(*self.plotrange)
+            pyplot.xlabel(self.dimensions[0]['label'])
+            pyplot.ylabel(label)
+            if self.plottitle:
+                pyplot.title(self.plottitle)
+            if self.postprocess_function:
+                self.postprocess_function(plothist, midpoints, None)
+            pyplot.savefig(self.plot_output_filename)
+
+    def _do_2d_output(self, hist, idims, midpoints, binbounds):
+        enehist = self._ener_zero(hist)
+        log10hist = np.log10(hist)
+
+        if self.hdf5_output_filename:
+            with h5py.File(self.hdf5_output_filename, 'w') as output_h5:
+                h5io.stamp_creator_data(output_h5)
+                output_h5.attrs['source_data'] = os.path.abspath(self.input_h5.filename)
+                output_h5.attrs['source_dimensions'] = np.array(idims, np.min_scalar_type(max(idims)))
+                output_h5.attrs['source_dimension_labels'] = np.array([dim['label'] for dim in self.dimensions])
+                for idim in idims:
+                    output_h5['midpoints_{}'.format(idim)] = midpoints[idim]
+                output_h5['histogram'] = hist
+
+        if self.plot_output_filename:
+            if self.plotscale == 'energy':
+                plothist = enehist
+                label = r'$-\ln\,P(x)$'
+            elif self.plotscale == 'log10':
+                plothist = log10hist
+                label = r'$\log_{10}\ P(\vec{x})$'
+            else:
+                plothist = hist
+                plothist[~np.isfinite(plothist)] = np.nan
+                label = r'$P(\vec{x})$'
+
+            try:
+                vmin, vmax = self.plotrange
+            except TypeError:
+                vmin, vmax = None, None
+
+            pyplot.figure()
+            # Transpose input so that axis 0 is displayed as x and axis 1 is displayed as y
+            #            pyplot.imshow(plothist.T, interpolation='nearest', aspect='auto',
+            #                          extent=(midpoints[0][0], midpoints[0][-1], midpoints[1][0], midpoints[1][-1]),
+            #                          origin='lower', vmin=vmin, vmax=vmax)
+
+            # The following reproduces the former calls to imshow and colorbar
+            norm = matplotlib.colors.Normalize(vmin=vmin, vmax=vmax)
+            ax = pyplot.gca()
+            nui = NonUniformImage(
+                ax, extent=(midpoints[0][0], midpoints[0][-1], midpoints[1][0], midpoints[1][-1]), origin='lower', norm=norm
+            )
+            nui.set_data(midpoints[0], midpoints[1], plothist.T)
+            ax.add_image(nui)
+            ax.set_xlim(midpoints[0][0], midpoints[0][-1])
+            ax.set_ylim(midpoints[1][0], midpoints[1][-1])
+            cb = pyplot.colorbar(nui)
+            cb.set_label(label)
+
+            pyplot.xlabel(self.dimensions[0]['label'])
+            pyplot.xlim(self.dimensions[0].get('lb'), self.dimensions[0].get('ub'))
+            pyplot.ylabel(self.dimensions[1]['label'])
+            pyplot.ylim(self.dimensions[1].get('lb'), self.dimensions[1].get('ub'))
+            if self.plottitle:
+                pyplot.title(self.plottitle)
+            if self.postprocess_function:
+                self.postprocess_function(plothist, midpoints, binbounds)
+            if self.plot_contour:
+                pyplot.contour(midpoints[0], midpoints[1], plothist.T)
+            pyplot.savefig(self.plot_output_filename)
+
+
+class InstantPlotHist(PlotSupports2D):
+    subcommand = 'instant'
+    help_text = 'plot probability distribution for a single WE iteration'
+    description = '''\
+Plot a probability distribution for a single WE iteration. The probability
+distribution must have been previously extracted with ``w_pdist`` (or, at
+least, must be compatible with the output format of ``w_pdist``; see
+``w_pdist --help`` for more information).
+'''
+
+    def add_args(self, parser):
+        self.input_arg_group.add_argument(
+            '--iter',
+            metavar='N_ITER',
+            dest='n_iter',
+            type=int,
+            help='''Plot distribution for iteration N_ITER
+                                          (default: last completed iteration).''',
+        )
+
+    def process_args(self, args):
+        if args.n_iter:
+            self.n_iter = min(args.n_iter, self.avail_iter_stop - 1)
+        else:
+            self.n_iter = self.avail_iter_stop - 1
+
+    def do_instant_plot_1d(self):
+        '''Plot the histogram for iteration self.n_iter'''
+
+        idim = self.dimensions[0]['idim']
+        n_iters = self.input_h5['n_iter'][...]
+        iiter = np.searchsorted(n_iters, self.n_iter)
+        binbounds = self.input_h5['binbounds_{}'.format(idim)][...]
+        midpoints = self.input_h5['midpoints_{}'.format(idim)][...]
+        hist = self.input_h5['histograms'][iiter]
+
+        # Average over other dimensions
+        hist = sum_except_along(hist, idim)
+        normhistnd(hist, [binbounds])
+        self._do_1d_output(hist, idim, midpoints)
+
+    def do_instant_plot_2d(self):
+        '''Plot the histogram for iteration self.n_iter'''
+
+        idim0 = self.dimensions[0]['idim']
+        idim1 = self.dimensions[1]['idim']
+
+        n_iters = self.input_h5['n_iter'][...]
+        iiter = np.searchsorted(n_iters, self.n_iter)
+        binbounds_0 = self.input_h5['binbounds_{}'.format(idim0)][...]
+        midpoints_0 = self.input_h5['midpoints_{}'.format(idim0)][...]
+        binbounds_1 = self.input_h5['binbounds_{}'.format(idim1)][...]
+        midpoints_1 = self.input_h5['midpoints_{}'.format(idim1)][...]
+
+        hist = self.input_h5['histograms'][iiter]
+
+        # Average over other dimensions
+        hist = sum_except_along(hist, [idim0, idim1])
+        normhistnd(hist, [binbounds_0, binbounds_1])
+        self._do_2d_output(hist, [idim0, idim1], [midpoints_0, midpoints_1], [binbounds_0, binbounds_1])
+
+    def go(self):
+        if len(self.dimensions) == 2:
+            self.do_instant_plot_2d()
+        else:
+            self.do_instant_plot_1d()
+
+
+class AveragePlotHist(PlotSupports2D):
+    subcommand = 'average'
+    help_text = 'plot average of a probability distribution over a WE simulation'
+    description = '''\
+Plot a probability distribution averaged over multiple iterations. The
+probability distribution must have been previously extracted with ``w_pdist``
+(or, at least, must be compatible with the output format of ``w_pdist``; see
+``w_pdist --help`` for more information).
+'''
+
+    def add_args(self, parser):
+        igroup = self.input_arg_group
+        igroup.add_argument(
+            '--first-iter',
+            dest='first_iter',
+            type=int,
+            metavar='N_ITER',
+            default=1,
+            help='''Begin averaging at iteration N_ITER (default: %(default)d).''',
+        )
+        igroup.add_argument(
+            '--last-iter',
+            dest='last_iter',
+            type=int,
+            metavar='N_ITER',
+            help='''Conclude averaging with N_ITER, inclusive (default: last completed iteration).''',
+        )
+
+    def process_args(self, args):
+        if args.first_iter:
+            self.iter_start = max(args.first_iter, self.avail_iter_start)
+        else:
+            self.iter_start = self.avail_iter_start
+
+        if args.last_iter:
+            self.iter_stop = min(args.last_iter + 1, self.avail_iter_stop)
+        else:
+            self.iter_stop = self.avail_iter_stop
+
+    def do_average_plot_1d(self):
+        '''Plot the average histogram for iterations self.iter_start to self.iter_stop'''
+
+        idim = self.dimensions[0]['idim']
+        n_iters = self.input_h5['n_iter'][...]
+        iiter_start = np.searchsorted(n_iters, self.iter_start)
+        iiter_stop = np.searchsorted(n_iters, self.iter_stop)
+        binbounds = self.input_h5['binbounds_{}'.format(idim)][...]
+        midpoints = self.input_h5['midpoints_{}'.format(idim)][...]
+        # hist = self.input_h5['histograms'][iiter_start:iiter_stop]
+
+        for iiter in range(iiter_start, iiter_stop):
+            iter_hist = sum_except_along(self.input_h5['histograms'][iiter], idim)
+            if iiter == iiter_start:
+                hist = iter_hist
+            else:
+                hist += iter_hist
+            del iter_hist
+
+        normhistnd(hist, [binbounds])
+        self._do_1d_output(hist, idim, midpoints)
+
+    def do_average_plot_2d(self):
+        '''Plot the histogram for iteration self.n_iter'''
+
+        idim0 = self.dimensions[0]['idim']
+        idim1 = self.dimensions[1]['idim']
+
+        n_iters = self.input_h5['n_iter'][...]
+        iiter_start = np.searchsorted(n_iters, self.iter_start)
+        iiter_stop = np.searchsorted(n_iters, self.iter_stop)
+
+        binbounds_0 = self.input_h5['binbounds_{}'.format(idim0)][...]
+        midpoints_0 = self.input_h5['midpoints_{}'.format(idim0)][...]
+        binbounds_1 = self.input_h5['binbounds_{}'.format(idim1)][...]
+        midpoints_1 = self.input_h5['midpoints_{}'.format(idim1)][...]
+
+        for iiter in range(iiter_start, iiter_stop):
+            iter_hist = sum_except_along(self.input_h5['histograms'][iiter], [idim0, idim1])
+            if iiter == iiter_start:
+                hist = iter_hist
+            else:
+                hist += iter_hist
+
+        normhistnd(hist, [binbounds_0, binbounds_1])
+        self._do_2d_output(hist, [idim0, idim1], [midpoints_0, midpoints_1], [binbounds_0, binbounds_1])
+
+    def go(self):
+        if len(self.dimensions) == 2:
+            self.do_average_plot_2d()
+        else:
+            self.do_average_plot_1d()
+
+
+class EvolutionPlotHist(PlotHistBase):
+    subcommand = 'evolution'
+    help_text = 'plot evolution of a probability distribution over the course of a WE simulation'
+    description = '''\
+Plot a probability distribution as it evolves over iterations. The
+probability distribution must have been previously extracted with ``w_pdist``
+(or, at least, must be compatible with the output format of ``w_pdist``; see
+``w_pdist --help`` for more information).
+'''
+
+    def add_args(self, parser):
+        igroup = self.input_arg_group
+        igroup.add_argument(
+            '--first-iter',
+            dest='first_iter',
+            type=int,
+            metavar='N_ITER',
+            default=1,
+            help='''Begin analysis at iteration N_ITER (default: %(default)d).''',
+        )
+        igroup.add_argument(
+            '--last-iter',
+            dest='last_iter',
+            type=int,
+            metavar='N_ITER',
+            help='''Conclude analysis with N_ITER, inclusive (default: last completed iteration).''',
+        )
+        igroup.add_argument(
+            '--step-iter', dest='step_iter', type=int, metavar='STEP', help='''Average in blocks of STEP iterations.'''
+        )
+
+    def process_args(self, args):
+        if args.first_iter:
+            self.iter_start = max(args.first_iter, self.avail_iter_start)
+        else:
+            self.iter_start = self.avail_iter_start
+
+        if args.last_iter:
+            self.iter_stop = min(args.last_iter + 1, self.avail_iter_stop)
+        else:
+            self.iter_stop = self.avail_iter_stop
+
+        if args.step_iter:
+            self.iter_step = max(args.step_iter, self.avail_iter_step)
+        else:
+            self.iter_step = self.avail_iter_step
+        log.info('using data for iterations {} -- {} with a step of {}'.format(self.iter_start, self.iter_stop, self.iter_step))
+
+    def go(self):
+        '''Plot the evolution of the histogram for iterations self.iter_start to self.iter_stop'''
+
+        idim = self.dimensions[0]['idim']
+        n_iters = self.input_h5['n_iter'][...]
+        iiter_start = np.searchsorted(n_iters, self.iter_start)
+        iiter_stop = np.searchsorted(n_iters, self.iter_stop)
+        binbounds = self.input_h5['binbounds_{}'.format(idim)][...]
+        midpoints = self.input_h5['midpoints_{}'.format(idim)][...]
+        hists_ds = self.input_h5['histograms']
+
+        itercount = self.iter_stop - self.iter_start
+
+        # We always round down, so that we don't have a dangling partial block at the end
+        nblocks = itercount // self.iter_step
+
+        block_iters = np.empty((nblocks, 2), dtype=n_iters.dtype)
+        blocked_hists = np.zeros((nblocks, hists_ds.shape[1 + idim]), dtype=hists_ds.dtype)
+
+        for iblock, istart in enumerate(range(iiter_start, iiter_start + nblocks * self.iter_step, self.iter_step)):
+            istop = min(istart + self.iter_step, iiter_stop)
+            histslice = hists_ds[istart:istop]
+
+            # Sum over time
+            histslice = np.add.reduce(histslice, axis=0)
+
+            # Sum over other dimensions
+            blocked_hists[iblock] = sum_except_along(histslice, idim)
+
+            # Normalize
+            normhistnd(blocked_hists[iblock], [binbounds])
+
+            block_iters[iblock, 0] = n_iters[istart]
+            block_iters[iblock, 1] = n_iters[istop - 1] + 1
+
+        # enehists = -np.log(blocked_hists)
+        enehists = self._ener_zero(blocked_hists)
+        log10hists = np.log10(blocked_hists)
+
+        if self.hdf5_output_filename:
+            with h5py.File(self.hdf5_output_filename, 'w') as output_h5:
+                h5io.stamp_creator_data(output_h5)
+                output_h5.attrs['source_data'] = os.path.abspath(self.input_h5.filename)
+                output_h5.attrs['source_dimension'] = idim
+                output_h5['midpoints'] = midpoints
+                output_h5['histograms'] = blocked_hists
+                output_h5['n_iter'] = block_iters
+
+        if self.plot_output_filename:
+            if self.plotscale == 'energy':
+                plothist = enehists
+                label = r'$-\ln\,P(x)$'
+            elif self.plotscale == 'log10':
+                plothist = log10hists
+                label = r'$\log_{10}\ P(x)$'
+            else:
+                plothist = blocked_hists
+                label = r'$P(x)$'
+
+            try:
+                vmin, vmax = self.plotrange
+            except TypeError:
+                vmin, vmax = None, None
+
+            pyplot.figure()
+            norm = matplotlib.colors.Normalize(vmin=vmin, vmax=vmax)
+            ax = pyplot.gca()
+            nui = NonUniformImage(
+                ax, extent=(midpoints[0], midpoints[-1], block_iters[0, -1], block_iters[-1, -1]), origin='lower', norm=norm
+            )
+
+            # not sure why plothist works but plothist.T doesn't, and the opposite is true
+            # for _do_2d_output
+            nui.set_data(midpoints, block_iters[:, -1], plothist)
+            ax.add_image(nui)
+            ax.set_xlim(midpoints[0], midpoints[-1])
+            ax.set_ylim(block_iters[0, -1], block_iters[-1, -1])
+            cb = pyplot.colorbar(nui)
+            cb.set_label(label)
+            pyplot.xlabel(self.dimensions[0]['label'])
+            pyplot.xlim(self.dimensions[0].get('lb'), self.dimensions[0].get('ub'))
+            pyplot.ylabel('WE Iteration')
+            if self.plottitle:
+                pyplot.title(self.plottitle)
+            if self.postprocess_function:
+                self.postprocess_function(plothist, midpoints, binbounds)
+            pyplot.savefig(self.plot_output_filename)
+
+
+class PlotHistTool(WESTMasterCommand):
+    prog = 'plothist'
+    subparsers_title = 'plotting modes'
+    subcommands = [InstantPlotHist, AveragePlotHist, EvolutionPlotHist]
+    description = '''\
+Plot probability density functions (histograms) generated by w_pdist or other
+programs conforming to the same output format. This program operates in one of
+three modes:
+
+  instant
+    Plot 1-D and 2-D histograms for an individual iteration. See
+    ``plothist instant --help`` for more information.
+
+  average
+    Plot 1-D and 2-D histograms, averaged over several iterations. See
+    ``plothist average --help`` for more information.
+
+  evolution
+    Plot the time evolution 1-D histograms as waterfall (heat map) plots.
+    See ``plothist evolution --help`` for more information.
+
+This program takes the output of ``w_pdist`` as input (see ``w_pdist --help``
+for more information), and can generate any kind of graphical output that
+matplotlib supports.
+
+
+------------------------------------------------------------------------------
+Command-line options
+------------------------------------------------------------------------------
+'''
+
+
+def entry_point():
+    PlotHistTool().main()
+
+
+if __name__ == '__main__':
+    entry_point()