westpa 2022.10__cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

westpa/cli/tools/ploterr.py

@@ -0,0 +1,506 @@
+import logging
+import os
+import re
+
+import h5py
+import numpy as np
+
+from westpa.tools import WESTMasterCommand, WESTSubcommand, ProgressIndicatorComponent, Plotter
+from westpa.core import h5io
+
+if os.environ.get('DISPLAY') is not None:
+    from matplotlib import pyplot
+
+log = logging.getLogger('ploterrs')
+
+
+class CommonPloterrs(WESTSubcommand):
+    def __init__(self, parent):
+        super().__init__(parent)
+
+        self.progress = ProgressIndicatorComponent()
+
+        self.xscale = None
+        self.yscale = None
+        self.xrange = None
+        self.yrange = None
+        self.xlabel = None
+        self.ylabel = None
+        self.title = None
+
+        self.plot_options_group = None
+
+    def add_args(self, parser):
+        self.progress.add_args(parser)
+
+        pogroup = self.plot_options_group = parser.add_argument_group('plot options')
+        pogroup.add_argument(
+            '--xscale',
+            choices=['linear', 'log', 'symlog'],
+            default='linear',
+            help='''Use "linear", "log", or "symlog" scaling for the x axis.
+                             (Default: %(default)s).''',
+        )
+        pogroup.add_argument(
+            '--yscale',
+            choices=['linear', 'log', 'symlog'],
+            default='linear',
+            help='''Use "linear", "log", or "symlog" scaling for the y axis.
+                             (Default: %(default)s).''',
+        )
+        pogroup.add_argument(
+            '--xrange',
+            help='''Restrict X range to XRANGE, which must be formatted as "xmin,xmax".
+                             (Default: determined by input data.)''',
+        )
+        pogroup.add_argument(
+            '--yrange',
+            help='''Restrict Y range to YRANGE, which must be formatted as "ymin,ymax".
+                             (Default: determined by input data.)''',
+        )
+        pogroup.add_argument('--xlabel', help='''Use XLABEL for the x-axis label. (Default: varies.)''')
+        pogroup.add_argument('--ylabel', help='''Use YLABEL for the y-axis label. (Default: varies.)''')
+        pogroup.add_argument('--title', help='''Use TITLE for the plot title. (Default: varies.)''')
+        pogroup.add_argument('--terminal', '-t', dest='plotting', action='store_true', help='''Plot output in terminal.''')
+
+    def process_args(self, args):
+        self.progress.process_args(args)
+
+        if args.xrange:
+            self.xrange = self.parse_range(args.xrange)
+
+        if args.yrange:
+            self.yrange = self.parse_range(args.yrange)
+
+        self.xscale = args.xscale
+        self.yscale = args.yscale
+        self.xlabel = args.xlabel or 'Iteration'
+        self.ylabel = args.ylabel
+        self.title = args.title
+        if args.plotting or os.environ.get('DISPLAY') is None:
+            self.interface = 'text'
+        else:
+            self.interface = 'matplotlib'
+
+    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}: {!s}'.format(rangespec, e))
+
+    def do_plot(self, data, output_filename, title=None, x_range=None, y_range=None, x_label=None, y_label=None):
+        if not output_filename:
+            return
+
+        title = title or self.title
+        x_range = x_range or self.xrange
+        y_range = y_range or self.yrange
+        x_label = x_label or self.xlabel
+        y_label = y_label or self.ylabel
+
+        iters = data['iter_stop'] - 1
+
+        pyplot.figure()
+        pyplot.plot(iters, data['expected'], color='black')
+        pyplot.plot(iters, data['ci_lbound'], color='gray')
+        pyplot.plot(iters, data['ci_ubound'], color='gray')
+
+        pyplot.gca().set_xscale(self.xscale)
+        pyplot.gca().set_yscale(self.yscale)
+
+        if title:
+            pyplot.title(title)
+
+        if x_range is not None:
+            pyplot.xlim(x_range)
+
+        if y_range is not None:
+            pyplot.ylim(y_range)
+
+        if x_label:
+            pyplot.xlabel(x_label)
+
+        if y_label:
+            pyplot.ylabel(y_label)
+
+        pyplot.savefig(output_filename)
+
+
+class GenericIntervalSubcommand(CommonPloterrs):
+    description = '''\
+Plots generic expectation/CI data. A path to the HDF5 file and the dataset
+within it must be provided. This path takes the form **FILENAME/PATH[SLICE]**.
+If the dataset is not a vector (one dimensional) then a slice must be provided.
+For example, to access the state 0 to state 1 rate evolution calculated by
+``w_kinavg``, one would use ``kinavg.h5/rate_evolution[:,0,1]``.
+
+
+-----------------------------------------------------------------------------
+Command-line arguments
+-----------------------------------------------------------------------------
+'''
+    subcommand = 'generic'
+    help_text = 'arbitrary HDF5 file and dataset'
+
+    def __init__(self, parent):
+        super().__init__(parent)
+        self.h5file = None
+        self.h5dset = None
+        self.dset_slice = None
+        self.output_filename = None
+
+    def add_args(self, parser):
+        iogroup = parser.add_argument_group('input/output options')
+        iogroup.add_argument(
+            '-o',
+            '--output',
+            default='errbars.pdf',
+            help='''Write plot to OUTPUT (default: %(default)s), whose format will
+                             be determined by filename extension.''',
+        )
+        iogroup.add_argument(
+            'dsspec',
+            help='''Use data located at DSSPEC, which must be formatted as
+                             FILENAME/PATH[SLICE]. FILENAME is the HDF5 file to read, PATH is the
+                             HDF5 path to the dataset, and SLICE, if provided, must be the Numpy-style
+                             slice (including brackets) which selects a vector of data of the
+                             appropriate type.''',
+        )
+
+    def process_args(self, args):
+        self.output_filename = args.output
+        (pathname, slicestr) = re.search(r'([^[]+)(\[[^\]]+\])?$', args.dsspec).groups()
+        if slicestr:
+            sl = eval('np.index_exp' + slicestr)
+        else:
+            sl = np.index_exp[...]
+        self.h5file, self.h5dset = h5io.resolve_filepath(pathname, mode='r')
+        self.dset_slice = sl
+
+    def load_and_validate_data(self):
+        reqd_fields = set(['iter_start', 'iter_stop', 'expected', 'ci_lbound', 'ci_ubound'])
+
+        self.progress.indicator.new_operation('loading data')
+        data = self.h5dset[self.dset_slice]
+
+        if data.ndim != 1:
+            raise TypeError('dataset to be plotted must be 1-dimensional')
+        try:
+            fieldnames = set(data.dtype.fields.keys())
+        except AttributeError:
+            raise TypeError('dataset has inappropriate type')
+        else:
+            if len(fieldnames & reqd_fields) < len(reqd_fields):
+                raise TypeError('dataset does not contain correct fields')
+
+        return data
+
+    def go(self):
+        with self.progress.indicator:
+            data = self.load_and_validate_data()
+            self.progress.indicator.new_operation('plotting')
+            self.do_plot(data, self.output_filename)
+
+
+class DirectKinetics(CommonPloterrs):
+    subcommand = 'd.kinetics'
+    help_text = 'output of w_direct kinetics'
+    input_filename = 'direct.h5'
+    flux_output_filename = 'flux_evolution_d_{state_label}.pdf'
+    rate_output_filename = 'rate_evolution_d_{istate_label}_{fstate_label}.pdf'
+    description = '''\
+Plot evolution of state-to-state rates and total flux into states as generated
+by ``w_{direct/reweight} kinetics`` (when used with the ``--evolution-mode``
+option). Plots are generated for all rates/fluxes calculated. Output filenames
+require (and plot titles and axis labels support) substitution based on which
+flux/rate is being plotted:
+
+  istate_label, fstate_label
+    *(String, for rates)* Names of the initial and final states, as originally
+    given to ``w_assign``.
+
+  istate_index, fstate_index
+    *(Integer, for rates)* Indices of initial and final states.
+
+  state_label
+    *(String, for fluxes)* Name of state
+
+  state_index
+    *(Integer, for fluxes)* Index of state
+'''
+
+    def __init__(self, parent):
+        super().__init__(parent)
+        self.kinavg_file = None
+
+        self.dset_slice = None
+        self.rate_output_pattern = None
+        self.flux_output_pattern = None
+
+        self.state_labels = None
+
+    def add_args(self, parser):
+        iogroup = parser.add_argument_group('input/output')
+        iogroup.add_argument(
+            '-i', '--input', default=self.input_filename, help='''Read kinetics results from INPUT (default: %(default)s).'''
+        )
+        iogroup.add_argument(
+            '--rate-output',
+            default=self.rate_output_filename,
+            help='''Filename pattern for rate evolution output. See above for valid
+                             field names. (Default: %(default)r).''',
+        )
+        iogroup.add_argument(
+            '--flux-output',
+            default=self.flux_output_filename,
+            help='''Filename pattern for flux evolution output. See above for valid
+                             field names. (Default: %(default)r).''',
+        )
+
+    def process_args(self, args):
+        self.kinavg_file = h5py.File(args.input, 'r')
+
+        self.state_labels = list(self.kinavg_file['state_labels'][...])
+
+        self.rate_output_pattern = args.rate_output
+        self.flux_output_pattern = args.flux_output
+
+    def plot_flux(self, istate):
+        label = self.state_labels[istate]
+        data = self.kinavg_file['target_flux_evolution'][:, istate]
+
+        if (data['iter_start'] == 0).all():
+            # No data
+            return
+
+        subdict = dict(state_label=label, state_index=istate)
+
+        output_filename = self.flux_output_pattern.format(**subdict) if self.flux_output_pattern else None
+
+        title = self.title if self.title is not None else 'Flux into state "{state_label}"'
+        title = title.format(**subdict)
+
+        x_label = self.xlabel.format(**subdict) if self.xlabel else None
+
+        y_label = self.ylabel if self.ylabel is not None else r'Flux $(\tau^{{-1}})$'
+        y_label = y_label.format(**subdict)
+
+        self.do_plot(data, output_filename, title, x_label=x_label, y_label=y_label)
+
+    def plot_rate(self, istate, jstate):
+        ilabel = self.state_labels[istate]
+        jlabel = self.state_labels[jstate]
+        data = self.kinavg_file['rate_evolution'][:, istate, jstate]
+
+        if (data['iter_start'] == 0).all():
+            # No data
+            return
+
+        subdict = dict(istate_label=ilabel, istate_index=istate, fstate_label=jlabel, fstate_index=jstate)
+
+        output_filename = self.rate_output_pattern.format(**subdict) if self.rate_output_pattern else None
+
+        title = self.title if self.title is not None else 'Rate from state "{istate_label}" to state "{fstate_label}"'
+        title = title.format(**subdict)
+
+        x_label = self.xlabel.format(**subdict) if self.xlabel else None
+
+        y_label = self.ylabel if self.ylabel is not None else r'Rate $(\tau^{{-1}})$'
+        y_label = y_label.format(**subdict)
+
+        self.do_plot(data, output_filename, title, x_label=x_label, y_label=y_label)
+
+    def go(self):
+        pi = self.progress.indicator
+        nstates = len(self.state_labels)
+        if self.interface == 'matplotlib':
+            with pi:
+                # if --evolution-mode wasn't specified, neither of these exist:
+                if 'target_flux_evolution' in self.kinavg_file:
+                    pi.new_operation('plotting fluxes', nstates)
+                    for istate in range(nstates):
+                        self.plot_flux(istate)
+                        pi.progress += 1
+
+                # if --evolution-mode wasn't specified, we won't get this either
+                if 'rate_evolution' in self.kinavg_file:
+                    pi.new_operation('plotting rates', nstates * nstates)
+                    for istate in range(nstates):
+                        for jstate in range(nstates):
+                            self.plot_rate(istate, jstate)
+                            pi.progress += 1
+                else:
+                    print('rate evolution not available')
+        else:
+            plotter = Plotter(self.kinavg_file, 'rate_evolution', iteration=-1, interface='text')
+            for istate in range(nstates):
+                for jstate in range(nstates):
+                    if istate != jstate:
+                        plotter.plot(istate, jstate)
+            plotter = Plotter(self.kinavg_file, 'conditional_flux_evolution', iteration=-1, interface='text')
+            for istate in range(nstates):
+                for jstate in range(nstates):
+                    if istate != jstate:
+                        plotter.plot(istate, jstate)
+
+
+class DirectStateprobs(CommonPloterrs):
+    subcommand = 'd.probs'
+    help_text = 'output of w_direct probs'
+    input_filename = 'direct.h5'
+    pop_output_filename = 'pop_evolution_d_{state_label}.pdf'
+    color_output_filename = 'color_evolution_d_{state_label}.pdf'
+    description = '''\
+Plot evolution of macrostate populations and associated uncertainties. Plots
+are generated for all states calculated. Output filenames require (and plot
+titles and axis labels support) substitution based on which state is being
+plotted:
+
+  state_label
+    *(String, for fluxes)* Name of state
+
+  state_index
+    *(Integer, for fluxes)* Index of state
+'''
+
+    def __init__(self, parent):
+        super().__init__(parent)
+        self.stateprobs_file = None
+
+        self.dset_slice = None
+        self.rate_output_pattern = None
+        self.flux_output_pattern = None
+
+        self.state_labels = None
+
+    def add_args(self, parser):
+        iogroup = parser.add_argument_group('input/output')
+        iogroup.add_argument(
+            '-i', '--input', default=self.input_filename, help='''Read w_kinavg results from INPUT (default: %(default)s).'''
+        )
+        iogroup.add_argument(
+            '--population-output',
+            default=self.pop_output_filename,
+            help='''Filename pattern for population evolution output. See above for valid
+                             field names. (Default: %(default)r).''',
+        )
+        iogroup.add_argument(
+            '--color-output',
+            default=self.color_output_filename,
+            help='''Filename pattern for ensemble evolution output. See above for valid
+                             field names. (Default: %(default)r).''',
+        )
+
+    def process_args(self, args):
+        self.stateprobs_file = h5py.File(args.input, 'r')
+
+        self.state_labels = list(self.stateprobs_file['state_labels'][...])
+        self.pop_output_pattern = args.population_output
+        self.color_output_pattern = args.color_output
+
+    def plot_pop(self, istate):
+        label = self.state_labels[istate]
+        data = self.stateprobs_file['state_pop_evolution'][:, istate]
+
+        if (data['iter_start'] == 0).all():
+            # No data
+            return
+
+        subdict = dict(state_label=label, state_index=istate)
+
+        output_filename = self.pop_output_pattern.format(**subdict) if self.pop_output_pattern else None
+
+        title = self.title if self.title is not None else 'Population in state "{state_label}"'
+        title = title.format(**subdict)
+
+        x_label = self.xlabel.format(**subdict) if self.xlabel else None
+        y_label = self.ylabel if self.ylabel is not None else r'Population'
+        y_label = y_label.format(**subdict)
+
+        self.do_plot(data, output_filename, title, x_label=x_label, y_label=y_label)
+
+    def plot_color(self, istate):
+        label = self.state_labels[istate]
+        data = self.stateprobs_file['color_prob_evolution'][:, istate]
+
+        if (data['iter_start'] == 0).all():
+            # No data
+            return
+
+        subdict = dict(state_label=label, state_index=istate)
+
+        output_filename = self.color_output_pattern.format(**subdict) if self.color_output_pattern else None
+
+        title = self.title if self.title is not None else 'Population in ensemble "{state_label}"'
+        title = title.format(**subdict)
+
+        x_label = self.xlabel.format(**subdict) if self.xlabel else None
+        y_label = self.ylabel if self.ylabel is not None else r'Population'
+        y_label = y_label.format(**subdict)
+
+        self.do_plot(data, output_filename, title, x_label=x_label, y_label=y_label)
+
+    def go(self):
+        pi = self.progress.indicator
+        nstates = len(self.state_labels)
+        if self.interface == 'matplotlib':
+            with pi:
+                if 'state_pop_evolution' in self.stateprobs_file:
+                    pi.new_operation('plotting populations', nstates)
+                    for istate in range(nstates):
+                        self.plot_pop(istate)
+                        pi.progress += 1
+
+                if 'color_prob_evolution' in self.stateprobs_file:
+                    pi.new_operation('plotting ensemble populations', nstates)
+                    for istate in range(nstates):
+                        self.plot_color(istate)
+                        pi.progress += 1
+                else:
+                    print('population evolution not available')
+        else:
+            plotter = Plotter(self.stateprobs_file, 'state_pop_evolution', iteration=-1, interface='text')
+            for istate in range(nstates):
+                plotter.plot(istate)
+            plotter = Plotter(self.stateprobs_file, 'color_prob_evolution', iteration=-1, interface='text')
+            for istate in range(nstates):
+                plotter.plot(istate)
+
+
+class ReweightStateprobs(DirectStateprobs):
+    subcommand = 'rw.probs'
+    help_text = 'output of w_reweight probs'
+    input_filename = 'reweight.h5'
+    pop_output_filename = 'pop_evolution_rw_{state_label}.pdf'
+    color_output_filename = 'color_evolution_rw_{state_label}.pdf'
+
+
+class ReweightKinetics(DirectKinetics):
+    subcommand = 'rw.kinetics'
+    help_text = 'output of w_reweight kinetics'
+    input_filename = 'reweight.h5'
+    flux_output_filename = 'flux_evolution_rw_{state_label}.pdf'
+    rate_output_filename = 'rate_evolution_rw_{istate_label}_{fstate_label}.pdf'
+
+
+class PloterrsTool(WESTMasterCommand):
+    prog = 'ploterrs'
+    subcommands = [DirectKinetics, DirectStateprobs, ReweightStateprobs, ReweightKinetics, GenericIntervalSubcommand]
+    subparsers_title = 'supported input formats'
+    description = '''\
+Plots error ranges for weighted ensemble datasets.
+
+
+-----------------------------------------------------------------------------
+Command-line options
+-----------------------------------------------------------------------------
+'''
+
+
+def entry_point():
+    PloterrsTool().main()
+
+
+if __name__ == '__main__':
+    entry_point()