westpa 2022.12__cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

westpa/tools/plot.py

@@ -0,0 +1,283 @@
+# import warnings
+# warnings.filterwarnings('ignore', category=DeprecationWarning)
+# warnings.filterwarnings('ignore', category=RuntimeWarning)
+# warnings.filterwarnings('ignore', category=FutureWarning)
+import numpy as np
+
+# We don't care if we're reimporting blessings, under this context.
+# warnings.filterwarnings('ignore')
+
+
+class Plotter:
+    '''
+    This is a semi-generic plotting interface that has a built in curses based terminal plotter.
+    It's fairly specific to what we're using it for here, but we could (and maybe should) build it out into
+    a little library that we can use via the command line to plot things.  Might be useful for looking at data later.
+    That would also cut the size of this tool down by a good bit.
+    '''
+
+    # def __init__(self, kinavg, kinrw, iteration, bin_labels, state_labels, state_pops, bin_pops, interface='matplotlib'):
+    def __init__(self, h5file, h5key, iteration=-1, interface='matplotlib'):
+        # Need to sort through and fix all this, but hey.
+        self.iteration = iteration
+        # These two are important for... reasons.
+        try:
+            # Commenting out below lines since `bin_labels` and `state_labels` are uninitialized variables.
+            # This will always fail so just raise explicitly
+
+            # self.bin_labels = list(bin_labels[...])
+            # self.state_labels = list(state_labels[...]) + ['unknown']
+            raise
+        except Exception:
+            try:
+                self.state_labels = list(h5file['state_labels'][...]) + ['unknown']
+            except Exception:
+                self.state_labels = None
+        # unless we totally fail out.
+        self.interface = interface
+        # What we should ACTUALLY do is just... yeah, just have it sub in what we need.
+        # We'll need to throw in the state labels or whatever, but.
+        self.h5file = h5file
+        self.h5key = h5key
+        # We should determine the number of dimensions of our dataset...
+        # This has time data, so an i to j is a 3 dim, and an i is 2.
+        try:
+            self.dim = len(h5file[h5key].shape)
+        except Exception:
+            self.dim = 1
+        try:
+            # does the ci exist?
+            _ = h5file[h5key]['expected']
+        except Exception:
+            self.dim = 1
+
+    def plot(self, i=0, j=1, tau=1, iteration=None, dim=0, interface=None):
+        if iteration is None:
+            iteration = self.iteration
+            self.__generic_ci__(self.h5file, iteration, i, j, tau=tau, h5key=self.h5key, dim=dim, interface=interface)
+
+    def __generic_ci__(self, h5file, iteration, i, j, tau, h5key='rate_evolution', dim=0, interface=None):
+        # This function just calls the appropriate plot function for our available
+        # interface.
+        if (interface is None and self.interface == 'text') or interface == 'text':
+            if self.dim > 1:
+                self.__terminal_ci__(h5file, iteration, i, j, tau, h5key)
+            else:
+                self.__terminal_expected__(h5file, iteration, i, j, tau, h5key, dim)
+        else:
+            try:
+                import matplotlib
+
+                matplotlib.use('TkAgg')
+                from matplotlib import pyplot as plt
+
+                if self.dim == 3:
+                    plt.plot(h5file[h5key]['expected'][:iteration, i, j] / tau, color='black')
+                    plt.plot(h5file[h5key]['ci_ubound'][:iteration, i, j] / tau, color='grey')
+                    plt.plot(h5file[h5key]['ci_lbound'][:iteration, i, j] / tau, color='grey')
+                else:
+                    plt.plot(h5file[h5key]['expected'][:iteration, i] / tau, color='black')
+                    plt.plot(h5file[h5key]['ci_ubound'][:iteration, i] / tau, color='grey')
+                    plt.plot(h5file[h5key]['ci_lbound'][:iteration, i] / tau, color='grey')
+                plt.show()
+            except Exception:
+                print('Unable to import plotting interface.  An X server ($DISPLAY) is required.')
+                if self.dim > 1:
+                    self.__terminal_ci__(h5file, iteration, i, j, tau)
+                else:
+                    self.__terminal_expected__(h5file, iteration, i, j, tau, h5key, dim)
+                return 1
+
+    def __generic_histo__(self, vector, labels):
+        # This function just calls the appropriate plot function for our available
+        # interface.  Same thing as generic_ci, but for a histogram.
+        if self.interface == 'text':
+            self.__terminal_histo__(vector, labels)
+        else:
+            try:
+                import matplotlib
+
+                matplotlib.use('TkAgg')
+                from matplotlib import pyplot as plt
+
+                plt.bar(
+                    list(range(0, np.array(vector).shape[0])), vector, linewidth=0, align='center', color='gold', tick_label=labels
+                )
+                plt.show()
+            except Exception:
+                print('Unable to import plotting interface.  An X server ($DISPLAY) is required.')
+                self.__terminal_histo__(vector, labels)
+                return 1
+
+    def __terminal_histo__(self, vector, labels, fullscreen_mode=True):
+        from blessings import Terminal
+
+        self.t = Terminal()
+        h = int(self.t.height / 4) * 3
+        w = self.t.width
+        cols = np.array(vector).shape[0]
+        # Let's print this business!
+
+        colwidth = w / cols
+        with self.t.fullscreen():
+            for y in range(0, h):
+                for x in range(0, cols):
+                    if x == 0:
+                        with self.t.location(0, y):
+                            print(self.t.red('{0:.4f}|'.format(float(h - y) / float(h))))
+                    with self.t.location((x * colwidth) + 8 + len(labels[x]) / 2, y):
+                        if vector[x] >= (float(h - y) / float(h)):
+                            # print(float(h-y)/float(h))
+                            print(self.t.on_blue(' '))
+            for x in range(0, cols):
+                if x == 0:
+                    with self.t.location(x, h):
+                        print('States| ')
+                with self.t.location((x * colwidth) + 8, h):
+                    print(self.t.blue(labels[x]))
+
+            if fullscreen_mode:
+                input("Press enter to continue.")
+
+    def __terminal_ci__(self, h5file, iteration, si, sj, tau, h5key):
+        from blessings import Terminal
+
+        self.t = Terminal()
+        h = int(self.t.height / 4 * 3.75)
+        # We'll figure out how to subsample the timepoints...
+        w = self.t.width
+        if self.dim == 3:
+            in_tup = (iteration - 1, si, sj)
+        else:
+            in_tup = (iteration - 1, si)
+        yupper = (h5file[h5key]['ci_ubound'][in_tup] / tau) * 2
+        ylower = (h5file[h5key]['ci_lbound'][in_tup] / tau) / 2
+        # Here are points pertaining to height.
+        scale = np.array([0.0] + [ylower + i * (yupper - ylower) / float(h) for i in range(0, h)])[::-1]
+        if iteration > w:
+            block_size = iteration / w
+        else:
+            block_size = 1
+
+        with self.t.fullscreen():
+            try:
+                for x in range(0, w - 12):
+                    iter = x * block_size
+                    if self.dim == 3:
+                        in_tup = (iter - 1, si, sj)
+                    else:
+                        in_tup = (iter - 1, si)
+                    yupper = h5file[h5key]['ci_ubound'][in_tup] / tau
+                    ylower = h5file[h5key]['ci_lbound'][in_tup] / tau
+                    ci = np.digitize([yupper, ylower], scale)
+                    if x == 0:
+                        for y in range(0, h + 1):
+                            with self.t.location(0, y):
+                                print(self.t.bold(self.t.red('{0:.7f}|'.format(scale[y]))))
+                    for y in range(ci[0], ci[1]):
+                        # with self.t.location(x+12, y):
+                        print(self.t.move(y, x + 12) + self.t.on_blue(' '))
+                        # print(self.t.on_blue(' '))
+                    # with self.t.location(x+12, np.digitize(h5file['rate_evolution']['expected'][iter-1, si, sj]/tau, scale)):
+                    #        print(self.t.on_blue('-'))
+                    print(self.t.move(np.digitize(h5file[h5key]['expected'][in_tup] / tau, scale), x + 12) + self.t.on_blue('-'))
+
+                for x in range(0, w - 12, w / 10):
+                    if x == 0:
+                        with self.t.location(x, h + 1):
+                            print('Iteration| ')
+                    with self.t.location(x + 12, h + 1):
+                        iter = x * block_size
+                        print(self.t.blue(str(iter)))
+            except Exception:
+                pass
+
+            with self.t.location(0, h + 2):
+                # We need to improve this.
+                # if h5key == 'rate_evolution':
+                #    print("k_ij from {} to {} from iter 1 to {}".format(self.state_labels[si], self.state_labels[sj], self.iteration))
+                # elif h5key == 'conditional_flux_evolution':
+                #    print("i->j flux from {} to {} from iter 1 to {}".format(self.state_labels[si], self.state_labels[sj], self.iteration))
+                if self.dim == 3:
+                    print(
+                        "{} from {} to {} from iter 1 to {}".format(
+                            h5key, self.state_labels[si], self.state_labels[sj], self.iteration
+                        )
+                    )
+                else:
+                    print("{} of state {} from iter 1 to {}".format(h5key, self.state_labels[si], self.iteration))
+            with self.t.location(0, h + 3):
+                input("Press enter to continue.")
+
+    def __terminal_expected__(self, h5file, iteration, si, sj, tau, h5key, dim):
+        from blessings import Terminal
+
+        self.t = Terminal()
+        h = int(self.t.height / 4 * 3.75)
+        # We'll figure out how to subsample the timepoints...
+        w = self.t.width
+        if self.dim == 3:
+            in_tup = (iteration - 1, si, sj)
+        else:
+            in_tup = (iteration - 1, si)
+        in_tup = (iteration - 1, dim)
+        try:
+            yupper = (np.max(h5file) / tau) * 2
+        except Exception:
+            in_tup = iteration - 1
+            yupper = (np.max(h5file) / tau) * 2
+        ylower = (np.min(h5file) / tau) * 2
+        # Here are points pertaining to height.
+        if yupper > 0:
+            yupper = (np.max(h5file) / tau) * 1.2
+            ylower = (np.min(h5file) / tau) / 2
+            scale = np.array([0.0] + [ylower + i * (yupper - ylower) / float(h) for i in range(0, h)])[::-1]
+        else:
+            yupper = (np.max(h5file) / tau) / 2
+            ylower = (np.min(h5file) / tau) * 1.2
+            scale = np.array([ylower + i * (yupper - ylower) / float(h) for i in range(0, h)] + [0.0])[::-1]
+        if iteration > w:
+            block_size = iteration / w
+        else:
+            block_size = 1
+
+        with self.t.fullscreen():
+            try:
+                for x in range(0, w - 12):
+                    iter = x * block_size
+                    if self.dim == 3:
+                        in_tup = (iter - 1, si, sj)
+                    else:
+                        in_tup = (iter - 1, si)
+                    in_tup = (iter - 1, dim)
+                    try:
+                        yupper = h5file[in_tup] / tau
+                    except Exception:
+                        in_tup = iter - 1
+                        yupper = h5file[in_tup] / tau
+                    ylower = h5file[in_tup] / tau
+                    ci = np.digitize([yupper, ylower], scale)
+                    if x == 0:
+                        for y in range(0, h + 1):
+                            with self.t.location(0, y):
+                                print(self.t.bold(self.t.red('{0:.7f}|'.format(scale[y]))))
+                    for y in range(ci[0], ci[1]):
+                        print(self.t.move(y, x + 12) + self.t.on_blue(' '))
+                        # print(self.t.on_blue(' '))
+                    print(self.t.move(np.digitize(h5file[in_tup] / tau, scale), x + 12) + self.t.on_blue('-'))
+
+                for x in range(0, w - 12, w / 10):
+                    if x == 0:
+                        with self.t.location(x, h + 1):
+                            print('Iteration| ')
+                    with self.t.location(x + 12, h + 1):
+                        iter = x * block_size
+                        print(self.t.blue(str(iter)))
+            except Exception:
+                pass
+
+            with self.t.location(0, h + 2):
+                # We need to improve this.
+                print("{} from iter 1 to {}".format(h5key, self.iteration))
+            with self.t.location(0, h + 3):
+                input("Press enter to continue.")

westpa/tools/progress.py

@@ -0,0 +1,17 @@
+import westpa
+from westpa.core.progress import ProgressIndicator
+from westpa.tools.core import WESTToolComponent
+
+
+class ProgressIndicatorComponent(WESTToolComponent):
+    def __init__(self):
+        super().__init__()
+        self.indicator = None
+
+    def add_args(self, parser):
+        pass
+
+    def process_args(self, args):
+        self.indicator = ProgressIndicator()
+        if westpa.rc.quiet_mode or westpa.rc.verbose_mode or westpa.rc.debug_mode:
+            self.indicator.fancy = False

westpa/tools/selected_segs.py

@@ -0,0 +1,154 @@
+import re
+
+import numpy as np
+
+import westpa
+from .core import WESTToolComponent
+from westpa.core.data_manager import seg_id_dtype
+
+re_split_segspec = re.compile(r'\s*:\s*|\s+')
+
+
+class SegmentSelection:
+    def __init__(self, iterable=None):
+        '''Initialize this segment selection from an iterable of (n_iter,seg_id) pairs.'''
+
+        self._segments = set()
+        self._segs_by_iter = {}
+        self._start_iter = None
+        self._stop_iter = None
+
+        if iterable is not None:
+            add = self.add
+            for pair in iterable:
+                add(tuple(pair))
+
+    def __len__(self):
+        return len(self._segments)
+
+    def __contains__(self, pair):
+        return tuple(pair) in self._segments
+
+    def add(self, pair):
+        (n_iter, seg_id) = int(pair[0]), int(pair[1])
+        self._segments.add((n_iter, seg_id))
+        self._segs_by_iter.setdefault(n_iter, set()).add(seg_id)
+        if self._start_iter is None:
+            self._start_iter = n_iter
+        else:
+            self._start_iter = min(self._start_iter, n_iter)
+        if self._stop_iter is None:
+            self._stop_iter = n_iter + 1
+        else:
+            self._stop_iter = max(self._stop_iter, n_iter + 1)
+
+    def from_iter(self, n_iter):
+        return self._segs_by_iter.get(n_iter, set())
+
+    @property
+    def start_iter(self):
+        return self._start_iter
+
+    @property
+    def stop_iter(self):
+        return self._stop_iter
+
+    @classmethod
+    def from_text(cls, filename):
+        segsel = cls()
+
+        segfile = open(filename, 'rt')
+        for line in segfile:
+            wline = line.strip()
+            if wline[0] == '#' or wline == '':
+                continue
+
+            fields = re_split_segspec.split(wline)
+            if len(fields) != 2:
+                raise ValueError('malformed segment selection {!r}'.format(line))
+
+            try:
+                n_iter, seg_id = list(map(int, fields))
+            except ValueError:
+                raise ValueError('malformed segment selection {!r}'.format(line))
+
+            segsel.add((n_iter, seg_id))
+        return segsel
+
+
+class AllSegmentSelection(SegmentSelection):
+    def __init__(self, start_iter=None, stop_iter=None, data_manager=None):
+        self.data_manager = data_manager or westpa.rc.get_data_manager()
+        self._start_iter = start_iter or 1
+        self._stop_iter = stop_iter or self.data_manager.current_iteration
+        self._segcount_by_iter = {}
+
+    def _count_from_iter(self, n_iter):
+        try:
+            segcount = self._segcount_by_iter[n_iter]
+        except KeyError:
+            segcount = self.data_manager.get_iter_group(n_iter)['seg_index'].shape[0]
+            self._segcount_by_iter[n_iter] = segcount
+        return segcount
+
+    def add(self, pair):
+        raise TypeError('cannot add segments to an AllSegmentSelection')
+
+    def from_iter(self, n_iter):
+        return np.arange(0, self._count_from_iter(n_iter), dtype=seg_id_dtype)
+
+    def __len__(self):
+        segcount = 0
+        count_from_iter = self._count_from_iter
+        for n_iter in range(self._start_iter, self._stop_iter):
+            segcount += count_from_iter(n_iter)
+        return segcount
+
+    def __contains__(self, pair):
+        (n_iter, seg_id) = pair
+        return n_iter >= self._start_iter and n_iter < self._stop_iter and seg_id < self._count_from_iter(n_iter)
+
+
+class SegSelector(WESTToolComponent):
+    def __init__(self):
+        super().__init__()
+        self.segment_selection = None
+        self.segsel_filename = None
+
+    def add_args(self, parser):
+        group = parser.add_argument_group('WEST input data options')
+        sgroup = group.add_mutually_exclusive_group()
+        sgroup.add_argument(
+            '--segments-from',
+            metavar='SEGLIST_FILE',
+            help='''Include only segments listed in SEGLIST_FILE, as generated by w_select
+                            (default: include all segments.''',
+        )
+
+    def process_args(self, args):
+        if args.segments_from:
+            self.segsel_filename = args.segments_from
+            self.segment_selection = self.parse_segsel_file(self.segsel_filename)
+        else:
+            self.segment_selection = AllSegmentSelection()
+
+    def parse_segsel_file(self, filename):
+        segsel = self.segment_selection = SegmentSelection()
+
+        segfile = open(filename, 'rt')
+        for line in segfile:
+            wline = line.strip()
+            if wline[0] == '#' or wline == '':
+                continue
+
+            fields = re_split_segspec.split(wline)
+            if len(fields) != 2:
+                raise ValueError('malformed segment selection {!r}'.format(line))
+
+            try:
+                n_iter, seg_id = list(map(int, fields))
+            except ValueError:
+                raise ValueError('malformed segment selection {!r}'.format(line))
+
+            segsel.add((n_iter, seg_id))
+        return segsel