Added load plots.

I merged Strategy and ExplicitStrategy for now. I also pruned zero
probability quorums from the strategies. The load plots are cool. I want
to plot capacity plots now, though it's a little tricky.

quorums/quorum_system.py

@@ -4,7 +4,7 @@
 from . import distribution
 from .distribution import Distribution
 from .expr import Expr, Node
-from .strategy import ExplicitStrategy, Strategy
+from .strategy import Strategy
 from typing import Dict, Iterator, Generic, List, Optional, Set, TypeVar
 import collections
 import itertools
@@ -251,8 +251,17 @@ class QuorumSystem(Generic[T]):
 
         # Solve the linear program.
         problem.solve(pulp.apis.PULP_CBC_CMD(msg=False))
-        return ExplicitStrategy(nodes,
-                                read_quorums,
-                                [v.varValue for v in read_quorum_vars],
-                                write_quorums,
-                                [v.varValue for v in write_quorum_vars])
+
+        non_zero_read_quorums = [
+            (rq, v.varValue)
+            for (rq, v) in zip(read_quorums, read_quorum_vars)
+            if v.varValue != 0]
+        non_zero_write_quorums = [
+            (wq, v.varValue)
+            for (wq, v) in zip(write_quorums, write_quorum_vars)
+            if v.varValue != 0]
+        return Strategy(nodes,
+                        [rq for (rq, _) in non_zero_read_quorums],
+                        [weight for (_, weight) in non_zero_read_quorums],
+                        [wq for (wq, _) in non_zero_write_quorums],
+                        [weight for (_, weight) in non_zero_write_quorums])

quorums/strategy.py

@@ -3,6 +3,8 @@ from .distribution import Distribution
 from .expr import Node
 from typing import Dict, Generic, List, Optional, Set, TypeVar
 import collections
+import matplotlib
+import matplotlib.pyplot as plt
 import numpy as np
 
 
@@ -10,26 +12,6 @@ T = TypeVar('T')
 
 
 class Strategy(Generic[T]):
-    def load(self,
-             read_fraction: Optional[Distribution] = None,
-             write_fraction: Optional[Distribution] = None) \
-             -> float:
-        raise NotImplementedError
-
-    def capacity(self,
-                 read_fraction: Optional[Distribution] = None,
-                 write_fraction: Optional[Distribution] = None) \
-                 -> float:
-        return 1 / self.load(read_fraction, write_fraction)
-
-    def get_read_quorum(self) -> Set[T]:
-        raise NotImplementedError
-
-    def get_write_quorum(self) -> Set[T]:
-        raise NotImplementedError
-
-
-class ExplicitStrategy(Strategy[T]):
     def __init__(self,
                  nodes: Set[Node[T]],
                  reads: List[Set[T]],
@@ -63,15 +45,20 @@ class ExplicitStrategy(Strategy[T]):
         non_zero_writes = {tuple(w): p
                            for (w, p) in zip(self.writes, self.write_weights)
                            if p > 0}
-        return (f'ExplicitStrategy(reads={non_zero_reads}, ' +
-                                 f'writes={non_zero_writes})')
+        return f'Strategy(reads={non_zero_reads}, writes={non_zero_writes})'
 
     def __repr__(self) -> str:
-        return (f'ExplicitStrategy(nodes={self.nodes}, '+
-                                 f'reads={self.reads}, ' +
-                                 f'read_weights={self.read_weights},' +
-                                 f'writes={self.writes}, ' +
-                                 f'write_weights={self.write_weights})')
+        return (f'Strategy(nodes={self.nodes}, '+
+                         f'reads={self.reads}, ' +
+                         f'read_weights={self.read_weights},' +
+                         f'writes={self.writes}, ' +
+                         f'write_weights={self.write_weights})')
+
+    def get_read_quorum(self) -> Set[T]:
+        return np.random.choice(self.reads, p=self.read_weights)
+
+    def get_write_quorum(self) -> Set[T]:
+        return np.random.choice(self.writes, p=self.write_weights)
 
     def load(self,
              read_fraction: Optional[Distribution] = None,
@@ -90,6 +77,79 @@ class ExplicitStrategy(Strategy[T]):
         return sum(weight * self._node_load(x, fr)
                    for (fr, weight) in d.items())
 
+    def capacity(self,
+                 read_fraction: Optional[Distribution] = None,
+                 write_fraction: Optional[Distribution] = None) \
+                 -> 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=nodes or list(self.nodes),
+                               read_fraction=read_fraction,
+                               write_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:
+        nodes = nodes or list(self.nodes)
+        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)))
+        read_cmap = matplotlib.cm.get_cmap('Reds')
+        write_cmap = matplotlib.cm.get_cmap('Blues')
+        d = distribution.canonicalize_rw(read_fraction, write_fraction)
+        fr = sum(weight * fr for (fr, weight) in d.items())
+        fw = 1 - fr
+
+        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 / 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 / self.write_capacity[x]
+            return bar_heights
+
+        bottom = np.zeros(len(x_list))
+        for (i, (rq, weight)) in enumerate(zip(self.reads, self.read_weights)):
+            bar_heights = fr * weight * read_quorum_to_bar_heights(rq)
+            ax.bar(x_ticks,
+                   bar_heights,
+                   bottom=bottom,
+                   color=read_cmap(1 - i * 0.75 / len(self.reads)),
+                   edgecolor='white', width=0.8)
+            bottom += bar_heights
+
+        for (i, (wq, weight)) in enumerate(zip(self.writes, self.write_weights)):
+            bar_heights = fw * weight * write_quorum_to_bar_heights(wq)
+            ax.bar(x_ticks,
+                   bar_heights,
+                   bottom=bottom,
+                   color=write_cmap(1 - i * 0.75 / len(self.writes)),
+                   edgecolor='white', width=0.8)
+            bottom += bar_heights
+
+        ax.set_xticks(x_ticks)
+        ax.set_xticklabels(str(x) for x in x_list)
+
     def _node_load(self, x: T, fr: float) -> float:
         """
         _node_load returns the load on x given a fixed read fraction fr.
@@ -105,9 +165,3 @@ class ExplicitStrategy(Strategy[T]):
         return max(self._node_load(node.x, fr) for node in self.nodes)
 
     # TODO(mwhittaker): Add read/write load and capacity and read/write cap.
-
-    def get_read_quorum(self) -> Set[T]:
-        return np.random.choice(self.reads, p=self.read_weights)
-
-    def get_write_quorum(self) -> Set[T]:
-        return np.random.choice(self.writes, p=self.write_weights)