Refactoring plotting to separate file.

plot_node_loads.py

@@ -3,14 +3,14 @@ import matplotlib
 import matplotlib.pyplot as plt
 import numpy as np
 
+
 def main():
     a = Node('a', write_capacity=1000, read_capacity=10000)
     b = Node('b', write_capacity=500, read_capacity=5000)
     c = Node('c', write_capacity=1000, read_capacity=10000)
     d = Node('d', write_capacity=500, read_capacity=5000)
     e = Node('e', write_capacity=1000, read_capacity=10000)
-    f = Node('f', write_capacity=500, read_capacity=5000)
+    nodes = [a, b, c, d, e]
-    nodes = [a, b, c, d, e, f]
 
     simple_majority = QuorumSystem(reads=majority([a, b, c, d, e]))
     crumbling_walls = QuorumSystem(reads=a*b + c*d*e)
@@ -20,9 +20,9 @@ def main():
     for i, qs in enumerate([simple_majority, crumbling_walls, paths]):
         fr = 0.9
         sigma = qs.strategy(read_fraction=fr)
-        sigma.plot_node_load_on(ax[0][i], nodes=nodes, read_fraction=fr)
+        plot_node_load_on(ax[0][i], sigma, nodes=nodes, read_fraction=fr)
-        sigma.plot_node_utilization_on(ax[1][i], nodes=nodes, read_fraction=fr)
+        plot_node_utilization_on(ax[1][i], sigma, nodes=nodes, read_fraction=fr)
-        sigma.plot_node_capacity_on(ax[2][i], nodes=nodes, read_fraction=fr)
+        plot_node_throughput_on(ax[2][i], sigma, nodes=nodes, read_fraction=fr)
 
     ax[0][0].set_title('Simple Majority')
     ax[0][1].set_title('Crumbling Walls')
@@ -33,5 +33,6 @@ def main():
     fig.tight_layout()
     fig.savefig('node_loads.pdf')
 
+
 if __name__ == '__main__':
     main()

quorums/__init__.py

@@ -1,2 +1,10 @@
 from .expr import Node, choose, majority
 from .quorum_system import QuorumSystem
+from .viz import (
+    plot_node_load,
+    plot_node_load_on,
+    plot_node_utilization,
+    plot_node_utilization_on,
+    plot_node_throughput,
+    plot_node_throughput_on,
+)

quorums/strategy.py

@@ -87,98 +87,6 @@ class Strategy(Generic[T]):
                  -> float:
         return 1 / self.load(read_fraction, write_fraction)
 
-
-    def plot_node_load(self,
-                       filename: str,
-                       nodes: Optional[List[Node[T]]] = None,
-                       read_fraction: Optional[Distribution] = None,
-                       write_fraction: Optional[Distribution] = None) \
-                       -> None:
-        fig, ax = plt.subplots()
-        self.plot_node_load_on(ax, nodes, read_fraction, write_fraction)
-        ax.set_xlabel('Node')
-        ax.set_ylabel('Load')
-        fig.tight_layout()
-        fig.savefig(filename)
-
-    def plot_node_load_on(self,
-                          ax: plt.Axes,
-                          nodes: Optional[List[Node[T]]] = None,
-                          read_fraction: Optional[Distribution] = None,
-                          write_fraction: Optional[Distribution] = None) \
-                          -> None:
-        self._plot_node_load_on(ax,
-                                scale=1,
-                                scale_by_node_capacity=True,
-                                nodes=nodes,
-                                read_fraction=read_fraction,
-                                write_fraction=write_fraction)
-
-    def plot_node_capacity(self,
-                           filename: str,
-                           nodes: Optional[List[Node[T]]] = None,
-                           read_fraction: Optional[Distribution] = None,
-                           write_fraction: Optional[Distribution] = None) \
-                           -> None:
-        fig, ax = plt.subplots()
-        self.plot_node_capacity_on(ax, nodes, read_fraction, write_fraction)
-        ax.set_xlabel('Node')
-        ax.set_ylabel('Throughput at Peak Throughput')
-        fig.tight_layout()
-        fig.savefig(filename)
-
-    def plot_node_capacity_on(self,
-                              ax: plt.Axes,
-                              nodes: Optional[List[Node[T]]] = None,
-                              read_fraction: Optional[Distribution] = None,
-                              write_fraction: Optional[Distribution] = None) \
-                              -> None:
-        self._plot_node_load_on(ax,
-                                scale=self.capacity(read_fraction,
-                                                    write_fraction),
-                                scale_by_node_capacity=False,
-                                nodes=nodes,
-                                read_fraction=read_fraction,
-                                write_fraction=write_fraction)
-
-    def plot_node_utilization(self,
-                              filename: str,
-                              nodes: Optional[List[Node[T]]] = None,
-                              read_fraction: Optional[Distribution] = None,
-                              write_fraction: Optional[Distribution] = None) \
-                              -> None:
-        fig, ax = plt.subplots()
-        self.plot_node_utilization_on(ax, nodes, read_fraction, write_fraction)
-        ax.set_xlabel('Node')
-        ax.set_ylabel('Utilization at Peak Throughput')
-        fig.tight_layout()
-        fig.savefig(filename)
-
-    def plot_node_utilization_on(self,
-                                 ax: plt.Axes,
-                                 nodes: Optional[List[Node[T]]] = None,
-                                 read_fraction: Optional[Distribution] = None,
-                                 write_fraction: Optional[Distribution] = None) \
-                                 -> None:
-        self._plot_node_load_on(
-                ax,
-                scale=self.capacity(read_fraction, write_fraction),
-                scale_by_node_capacity=True,
-                nodes=nodes,
-                read_fraction=read_fraction,
-                write_fraction=write_fraction)
-
-    def plot_load_distribution(self,
-                               filename: str,
-                               nodes: Optional[List[Node[T]]] = None) \
-                               -> None:
-        fig, ax = plt.subplots()
-        self.plot_load_distribution_on(ax, nodes)
-        ax.set_xlabel('Read Fraction')
-        ax.set_ylabel('Load')
-        fig.tight_layout()
-        fig.savefig(filename)
-
     def _group(self, segments: List[Tuple[Segment, T]]) -> Dict[Segment, List[T]]:
         groups: Dict[Segment, List[T]] = collections.defaultdict(list)
         for segment, x in segments:
@@ -240,69 +148,3 @@ class Strategy(Generic[T]):
         _load returns the load given a fixed read fraction fr.
         """
         return max(self._node_load(node.x, fr) for node in self.nodes)
-
-    def _plot_node_load_on(
-            self,
-            ax: plt.Axes,
-            scale: float,
-            scale_by_node_capacity: bool,
-            nodes: Optional[List[Node[T]]] = None,
-            read_fraction: Optional[Distribution] = None,
-            write_fraction: Optional[Distribution] = None) \
-            -> None:
-        nodes = nodes or list(self.nodes)
-        d = distribution.canonicalize_rw(read_fraction, write_fraction)
-        x_list = [node.x for node in nodes]
-        x_index = {x: i for (i, x) in enumerate(x_list)}
-        x_ticks = list(range(len(x_list)))
-
-        def read_quorum_to_bar_heights(quorum: Set[T]) -> np.array:
-            bar_heights = np.zeros(len(x_list))
-            for x in quorum:
-                bar_heights[x_index[x]] = 1
-                if scale_by_node_capacity:
-                    bar_heights[x_index[x]] /= self.read_capacity[x]
-            return bar_heights
-
-        def write_quorum_to_bar_heights(quorum: Set[T]) -> np.array:
-            bar_heights = np.zeros(len(x_list))
-            for x in quorum:
-                bar_heights[x_index[x]] = 1
-                if scale_by_node_capacity:
-                    bar_heights[x_index[x]] /= self.write_capacity[x]
-            return bar_heights
-
-        bottoms = np.zeros(len(x_list))
-
-        fr = sum(weight * fr for (fr, weight) in d.items())
-        read_cmap = matplotlib.cm.get_cmap('Reds')
-        for (i, (rq, weight)) in enumerate(zip(self.reads, self.read_weights)):
-            bar_heights = scale * fr * weight * read_quorum_to_bar_heights(rq)
-            ax.bar(x_ticks,
-                   bar_heights,
-                   bottom=bottoms,
-                   color=read_cmap(0.75 - i * 0.5 / len(self.reads)),
-                   edgecolor='white', width=0.8)
-            for j, (bar_height, bottom) in enumerate(zip(bar_heights, bottoms)):
-                if bar_height != 0:
-                    ax.text(x_ticks[j], bottom + bar_height / 2, i,
-                            ha='center', va='center')
-            bottoms += bar_heights
-
-        fw = 1 - fr
-        write_cmap = matplotlib.cm.get_cmap('Blues')
-        for (i, (wq, weight)) in enumerate(zip(self.writes, self.write_weights)):
-            bar_heights = scale * fw * weight * write_quorum_to_bar_heights(wq)
-            ax.bar(x_ticks,
-                   bar_heights,
-                   bottom=bottoms,
-                   color=write_cmap(0.75 - i * 0.5 / len(self.writes)),
-                   edgecolor='white', width=0.8)
-            for j, (bar_height, bottom) in enumerate(zip(bar_heights, bottoms)):
-                if bar_height != 0:
-                    ax.text(x_ticks[j], bottom + bar_height / 2, i,
-                            ha='center', va='center')
-            bottoms += bar_heights
-
-        ax.set_xticks(x_ticks)
-        ax.set_xticklabels(str(x) for x in x_list)

quorums/viz.py

@@ -0,0 +1,148 @@
+from . import distribution
+from .distribution import Distribution
+from .expr import Node
+from .strategy import Strategy
+from typing import List, Optional, Set, TypeVar
+import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+T = TypeVar('T')
+
+
+def plot_node_load(filename: str,
+                   strategy: Strategy[T],
+                   nodes: Optional[List[Node[T]]] = None,
+                   read_fraction: Optional[Distribution] = None,
+                   write_fraction: Optional[Distribution] = None):
+    fig, ax = plt.subplots()
+    plot_node_load_on(ax, strategy, nodes, read_fraction, write_fraction)
+    ax.set_xlabel('Node')
+    ax.set_ylabel('Load')
+    fig.tight_layout()
+    fig.savefig(filename)
+
+
+def plot_node_load_on(ax: plt.Axes,
+                      strategy: Strategy[T],
+                      nodes: Optional[List[Node[T]]] = None,
+                      read_fraction: Optional[Distribution] = None,
+                      write_fraction: Optional[Distribution] = None):
+    _plot_node_load_on(ax,
+                       strategy,
+                       nodes or list(strategy.nodes),
+                       scale=1,
+                       scale_by_node_capacity=True,
+                       read_fraction=read_fraction,
+                       write_fraction=write_fraction)
+
+
+def plot_node_utilization(filename: str,
+                          strategy: Strategy[T],
+                          nodes: Optional[List[Node[T]]] = None,
+                          read_fraction: Optional[Distribution] = None,
+                          write_fraction: Optional[Distribution] = None):
+    fig, ax = plt.subplots()
+    plot_node_utilization_on(ax, strategy, nodes, read_fraction, write_fraction)
+    ax.set_xlabel('Node')
+    ax.set_ylabel('Utilization')
+    fig.tight_layout()
+    fig.savefig(filename)
+
+
+def plot_node_utilization_on(ax: plt.Axes,
+                      strategy: Strategy[T],
+                      nodes: Optional[List[Node[T]]] = None,
+                      read_fraction: Optional[Distribution] = None,
+                      write_fraction: Optional[Distribution] = None):
+    _plot_node_load_on(ax,
+                       strategy,
+                       nodes or list(strategy.nodes),
+                       scale=strategy.capacity(read_fraction, write_fraction),
+                       scale_by_node_capacity=True,
+                       read_fraction=read_fraction,
+                       write_fraction=write_fraction)
+
+
+def plot_node_throughput(filename: str,
+                   strategy: Strategy[T],
+                   nodes: Optional[List[Node[T]]] = None,
+                   read_fraction: Optional[Distribution] = None,
+                   write_fraction: Optional[Distribution] = None):
+    fig, ax = plt.subplots()
+    plot_node_throughput_on(ax, strategy, nodes, read_fraction, write_fraction)
+    ax.set_xlabel('Node')
+    ax.set_ylabel('Throughput')
+    fig.tight_layout()
+    fig.savefig(filename)
+
+
+def plot_node_throughput_on(ax: plt.Axes,
+                            strategy: Strategy[T],
+                            nodes: Optional[List[Node[T]]] = None,
+                            read_fraction: Optional[Distribution] = None,
+                            write_fraction: Optional[Distribution] = None):
+    _plot_node_load_on(ax,
+                       strategy,
+                       nodes or list(strategy.nodes),
+                       scale=strategy.capacity(read_fraction, write_fraction),
+                       scale_by_node_capacity=False,
+                       read_fraction=read_fraction,
+                       write_fraction=write_fraction)
+
+
+def _plot_node_load_on(ax: plt.Axes,
+                       sigma: Strategy[T],
+                       nodes: List[Node[T]],
+                       scale: float,
+                       scale_by_node_capacity: bool,
+                       read_fraction: Optional[Distribution] = None,
+                       write_fraction: Optional[Distribution] = None):
+    d = distribution.canonicalize_rw(read_fraction, write_fraction)
+    x_list = [node.x for node in nodes]
+    x_index = {x: i for (i, x) in enumerate(x_list)}
+    x_ticks = list(range(len(x_list)))
+
+    def one_hot(quorum: Set[T]) -> np.array:
+        bar_heights = np.zeros(len(x_list))
+        for x in quorum:
+            bar_heights[x_index[x]] = 1
+        return bar_heights
+
+    def plot_quorums(quorums: List[Set[T]],
+                     weights: List[float],
+                     fraction: float,
+                     bottoms: np.array,
+                     capacities: np.array,
+                     cmap: matplotlib.colors.Colormap):
+        for (i, (quorum, weight)) in enumerate(zip(quorums, weights)):
+            bar_heights = scale * fraction * weight * one_hot(quorum)
+            if scale_by_node_capacity:
+                bar_heights /= capacities
+
+            ax.bar(x_ticks,
+                   bar_heights,
+                   bottom=bottoms,
+                   color=cmap(0.75 - i * 0.5 / len(quorums)),
+                   edgecolor='white', width=0.8)
+
+            for j, (bar_height, bottom) in enumerate(zip(bar_heights, bottoms)):
+                # TODO(mwhittaker): Fix the unhappy typechecker.
+                text = ''.join(str(x) for x in sorted(list(quorum)))
+                if bar_height != 0:
+                    ax.text(x_ticks[j], bottom + bar_height / 2, text,
+                            ha='center', va='center')
+            bottoms += bar_heights
+
+    fr = sum(weight * fr for (fr, weight) in d.items())
+    fw = 1 - fr
+    read_capacities = np.array([node.read_capacity for node in nodes])
+    write_capacities = np.array([node.write_capacity for node in nodes])
+    bottoms = np.zeros(len(x_list))
+    plot_quorums(sigma.reads, sigma.read_weights, fr, bottoms, read_capacities,
+                 matplotlib.cm.get_cmap('Reds'))
+    plot_quorums(sigma.writes, sigma.write_weights, fw, bottoms,
+                 write_capacities, matplotlib.cm.get_cmap('Blues'))
+    ax.set_xticks(x_ticks)
+    ax.set_xticklabels(str(x) for x in x_list)