dwhswenson · Aug 31, 2024 · Aug 12, 2021 · Aug 23, 2021 · Apr 30, 2024 · May 1, 2024
diff --git a/openpathsampling/experimental/path_tree/__init__.py b/openpathsampling/experimental/path_tree/__init__.py
@@ -0,0 +1,3 @@
+from .path_tree import PathTree
+from . import mpl
+from . import svg
diff --git a/openpathsampling/experimental/path_tree/mpl.py b/openpathsampling/experimental/path_tree/mpl.py
@@ -0,0 +1,108 @@
+from functools import partial
+from .options import canonicalize_mover
+from .plotter import PathTreePlotter
+
+# In addition to the standard `draw_trajectory` and `draw_connector` methods
+# required of all PathTree plot styles, plot styles based on matplotlib
+# should have the attribute `ax`, which returns the `Axes` object they plot
+# into.
+
+
+
+class MPLPlottingStyle(PathTreePlotter):
+    """Abstract base class for matplotlib-based path tree plotting styles.
+
+    Parameters
+    ----------
+    ax : :class:`.matplotlib.Axes`
+        Axes object to plot into
+    """
+    def __init__(self, ax=None, options=None):
+        try:
+            import matplotlib.pyplot as plt
+        except ImportError:
+            raise RuntimeError("matplotlib is not installed")
+
+        super().__init__(options)
+        if ax is None:
+            _, ax = plt.subplots()
+        self.ax = ax
+
+    # @staticmethod
+    # def step_basics(step, options):
+    #     mover = canonicalize_mover(step.mover)
+    #     mover_options = options.movers[mover]
+    #     color = mover_options.color
+    #     plot_segments = mover_options.get_left_right(step)
+    #     return mover, color, plot_segments
+
+    def draw_trajectory(self, row, step, options):
+        raise NotImplementedError()
+
+    def draw_connector(self, x, bottom, top, step):
+        mover = canonicalize_mover(step.mover)
+        if mover in self.CONNECTOR_FIXES:
+            x = self.CONNECTOR_FIXES[mover](x)
+
+        self.ax.plot([x,x], [-bottom, -top], color='k')
+
+
+### ThinLines plotting style ###############################################
+
+class ThinLines(MPLPlottingStyle):
+    """ThinLines plotting style for matplotlib-based path trees.
+
+    This is a plotting style based on using `ax.plot` and with the style
+    choice that the shooting point is always represented at the middle of
+    its frame. This needs to us filling half of the shooting point frame so
+    that the trajectory line intersects the shooting point connector. This
+    is accomplished by the ``thin_lines_fix`` function.
+
+    """
+    PLOT_TYPE = "trajectory"
+    def draw_trajectory(self, row, step):
+        mover = canonicalize_mover(step.mover)
+        plot_segments, color = self.get_step_plot_details(step)
+
+        for left, right in plot_segments:
+            self.ax.plot([left - 0.5, right - 0.5], [-row] * 2, lw=3,
+                         color=color)
+
+
+
+### Blocks plotting style ##################################################
+
+from matplotlib.patches import Rectangle
+
+class Blocks(MPLPlottingStyle):
+    """Blocks plotting style for matplotlib-based path trees.
+
+    The Blocks plotting style uses rectangular patches for each snapshot.
+    """
+    PLOT_TYPE = "snapshot"
+    def __init__(self, ax=None, options=None, *, block_size=(0.8, 0.6)):
+        super().__init__(ax, options)
+        self.block_size = block_size
+
+    def draw_trajectory(self, row, step):
+        """
+        Parameters
+        ----------
+        row : int
+            the row number (counting from the top) for this step
+        """
+        mover = canonicalize_mover(step.mover)
+        plot_segments, color = self.get_step_plot_details(step)
+        # first we draw the background line
+        for left, right in plot_segments:
+            self.ax.plot([left - 0.5, right - 0.5], [-row]*2, color=color)
+
+        # now we draw the boxes
+        for left, right in plot_segments:
+            width, height = self.block_size
+            y = - row - height / 2
+            for snap in range(left, right):
+                self.ax.add_patch(
+                    Rectangle((snap - width / 2, y), width, height,
+                              color=color)
+                )
diff --git a/openpathsampling/experimental/path_tree/options.py b/openpathsampling/experimental/path_tree/options.py
@@ -0,0 +1,104 @@
+def canonicalize_mover(mover):
+    if mover is not None:
+        return mover.__class__.__name__
+
+
+class OptionDict:
+    def __init__(self, dictionary, default):
+        self.dictionary = dictionary
+        self.default = default
+
+    def __getitem__(self, item):
+        try:
+            return self.dictionary[item]
+        except KeyError:
+            return self.default
+
+    def __getattr__(self, attr):
+        if attr in ['dictionary', 'default']:
+            return object.__getattribute__(self, attr)
+        return self[attr]
+
+    def __setattr__(self, key, value):
+        if key not in ['dictionary', 'default']:
+            self.dictionary[key] = value
+        else:
+            object.__setattr__(self, key, value)
+
+    def __repr__(self):
+        return (f"OptionDict({repr(self.dictionary)}, "
+                f"default={repr(self.default)}")
+
+
+class PathTreeMoverOptions(object):
+    def __init__(self, show='full', color='#0f0'):
+        self.show = show
+        self.color = color
+
+        self.get_left_right_methods = {
+            'full': self.get_left_right_full,
+            'new': self.get_left_right_new,
+        }
+
+    @staticmethod
+    def get_left_right_full(step):
+        offset = step.offset if step.offset is not None else 0
+        return [(offset, offset + len(step.trajectory))]
+
+    @staticmethod
+    def get_left_right_new(step):
+        offset = step.offset if step.offset is not None else 0
+        return [(offset + start, offset + stop)
+                for start, stop in step.segments.new]
+
+    def get_left_right(self, step):
+        get_left_right = self.get_left_right_methods[self.show]
+        return get_left_right(step)
+
+    def __repr__(self):
+        return (f"{self.__class__.__name__}(show='{self.show}', "
+                f"color='{self.color})")
+
+
+class PathTreeOptions(object):
+    def __init__(self, movers, rejected, allowed_movers, forbidden_movers):
+        self.movers = movers
+        self.rejected = rejected
+        self.allowed_movers = allowed_movers
+        self.forbidden_movers = forbidden_movers
+
+    def _skip_step(self, step):
+        mover = canonicalize_mover(step.mover)
+        include_rejected = self.rejected != 'hidden'
+        skip_rejected = not include_rejected and not step.accepted
+        skip_allowed = self.allowed_movers and mover not in allowed
+        skip_forbidden = self.forbidden_movers and mover in forbidden
+        return skip_rejected or skip_allowed or skip_forbidden
+
+    def filter_tree_steps(self, tree_steps):
+        for step in tree_steps:
+            if not self._skip_step(step):
+                yield step
+
+
+def create_default_options():
+    two_way_options = PathTreeMoverOptions(color='#9A63EE')
+    movers = OptionDict(
+        {
+            'ForwardShootMover': PathTreeMoverOptions(show='new',
+                                                      color='#00f'),
+            'BackwardShootMover': PathTreeMoverOptions(show='new',
+                                                       color='#f00'),
+            'BackwardFirstTwoWayShootingMover': two_way_options,
+            'ForwardFirstTwoWayShootingMover': two_way_options,
+            None: PathTreeMoverOptions(color='#808080'),
+
+        },
+        default=PathTreeMoverOptions(color="#377E22")
+    )
+    return PathTreeOptions(
+        movers=movers,
+        rejected='hidden',
+        allowed_movers=None,
+        forbidden_movers=None,
+    )