Fix bug in plot.py and add missing test

axelrod/plot.py

@@ -3,6 +3,7 @@
 import tqdm
 import warnings
 
+from distutils.version import LooseVersion
 from typing import List, Union
 
 matplotlib_installed = True
@@ -24,13 +25,12 @@
 namesType = List[str]
 dataType = List[List[Union[int, float]]]
 
-def default_cmap() -> str:
+
+def default_cmap(version: str = "2.0") -> str:
     """Sets a default matplotlib colormap based on the version."""
-    s = matplotlib.__version__.split('.')
-    if int(s[0]) >= 1 and int(s[1]) >= 5:
-        return "viridis"
-    else:
-        return 'YlGnBu'
+    if LooseVersion(version) >= "1.5":
+        return 'viridis'
+    return 'YlGnBu'
 
 
 class Plot(object):
@@ -40,7 +40,10 @@ def __init__(self, result_set: ResultSet) -> None:
         self.nplayers = self.result_set.nplayers
         self.players = self.result_set.players
 
-    def _violinplot(self, data: dataType, names: namesType, title: titleType =None, ax: matplotlib.axes.SubplotBase =None) -> matplotlib.figure.Figure:
+    def _violinplot(
+        self, data: dataType, names: namesType, title: titleType = None,
+        ax: matplotlib.axes.SubplotBase = None
+    ) -> matplotlib.figure.Figure:
         """For making violinplots."""
         if not self.matplotlib_installed:  # pragma: no cover
             return None
@@ -56,13 +59,15 @@ def _violinplot(self, data: dataType, names: namesType, title: titleType =None,
         spacing = 4
         positions = spacing * arange(1, self.nplayers + 1, 1)
         figure.set_size_inches(width, height)
-        plt.violinplot(data, positions=positions, widths=spacing / 2,
-                       showmedians=True, showextrema=False)
-        plt.xticks(positions, names, rotation=90)
-        plt.xlim(0, spacing * (self.nplayers + 1))
-        plt.tick_params(axis='both', which='both', labelsize=8)
+        ax.violinplot(data, positions=positions, widths=spacing / 2,
+                      showmedians=True, showextrema=False)
+        ax.set_xticks(positions)
+        ax.set_xticklabels(names, rotation=90)
+        ax.set_xlim([0, spacing * (self.nplayers + 1)])
+        ax.tick_params(axis='both', which='both', labelsize=8)
         if title:
-            plt.title(title)
+            ax.set_title(title)
+        plt.tight_layout()
         return figure
 
     # Box and Violin plots for mean score, score differences, wins, and match
@@ -81,11 +86,13 @@ def _boxplot_xticks_locations(self):
     def _boxplot_xticks_labels(self):
         return [str(n) for n in self.result_set.ranked_names]
 
-    def boxplot(self, title: titleType =None, ax: matplotlib.axes.SubplotBase =None) -> matplotlib.figure.Figure:
+    def boxplot(
+        self, title: titleType = None, ax: matplotlib.axes.SubplotBase = None
+    ) -> matplotlib.figure.Figure:
         """For the specific mean score boxplot."""
         data = self._boxplot_dataset
         names = self._boxplot_xticks_labels
-        figure = self._violinplot(data, names, title=title,  ax=ax)
+        figure = self._violinplot(data, names, title=title, ax=ax)
         return figure
 
     @property
@@ -100,7 +107,9 @@ def _winplot_dataset(self):
         ranked_names = [str(self.players[x[-1]]) for x in medians]
         return wins, ranked_names
 
-    def winplot(self, title: titleType =None, ax: matplotlib.axes.SubplotBase =None) -> matplotlib.figure.Figure:
+    def winplot(
+        self, title: titleType = None, ax: matplotlib.axes.SubplotBase = None
+    ) -> matplotlib.figure.Figure:
         """Plots the distributions for the number of wins for each strategy."""
         if not self.matplotlib_installed:  # pragma: no cover
             return None
@@ -125,7 +134,9 @@ def _sdv_plot_dataset(self):
         ranked_names = [str(self.players[i]) for i in ordering]
         return diffs, ranked_names
 
-    def sdvplot(self, title: titleType =None, ax: matplotlib.axes.SubplotBase =None) -> matplotlib.figure.Figure:
+    def sdvplot(
+        self, title: titleType = None, ax: matplotlib.axes.SubplotBase = None
+    ) -> matplotlib.figure.Figure:
         """Score difference violinplots to visualize the distributions of how
         players attain their payoffs."""
         diffs, ranked_names = self._sdv_plot_dataset
@@ -139,7 +150,9 @@ def _lengthplot_dataset(self):
                  for length in rep[playeri]] for playeri in
                 self.result_set.ranking]
 
-    def lengthplot(self, title: titleType =None, ax: matplotlib.axes.SubplotBase =None) -> matplotlib.figure.Figure:
+    def lengthplot(
+        self, title: titleType = None, ax: matplotlib.axes.SubplotBase = None
+    ) -> matplotlib.figure.Figure:
         """For the specific match length boxplot."""
         data = self._lengthplot_dataset
         names = self._boxplot_xticks_labels
@@ -165,7 +178,10 @@ def _pdplot_dataset(self):
         ranked_names = [str(players[i]) for i in ordering]
         return matrix, ranked_names
 
-    def _payoff_heatmap(self, data: dataType, names: namesType, title: titleType =None, ax: matplotlib.axes.SubplotBase =None) -> matplotlib.figure.Figure:
+    def _payoff_heatmap(
+        self, data: dataType, names: namesType, title: titleType = None,
+        ax: matplotlib.axes.SubplotBase = None
+    ) -> matplotlib.figure.Figure:
         """Generic heatmap plot"""
         if not self.matplotlib_installed:  # pragma: no cover
             return None
@@ -179,28 +195,31 @@ def _payoff_heatmap(self, data: dataType, names: namesType, title: titleType =No
         width = max(self.nplayers / 4, 12)
         height = width
         figure.set_size_inches(width, height)
-        cmap = default_cmap()
+        matplotlib_version = matplotlib.__version__
+        cmap = default_cmap(matplotlib_version)
         mat = ax.matshow(data, cmap=cmap)
-        plt.xticks(range(self.result_set.nplayers))
-        plt.yticks(range(self.result_set.nplayers))
+        ax.set_xticks(range(self.result_set.nplayers))
+        ax.set_yticks(range(self.result_set.nplayers))
         ax.set_xticklabels(names, rotation=90)
         ax.set_yticklabels(names)
-        plt.tick_params(axis='both', which='both', labelsize=16)
+        ax.tick_params(axis='both', which='both', labelsize=16)
         if title:
-            plt.xlabel(title)
-        # Make the colorbar match up with the plot
-        divider = make_axes_locatable(plt.gca())
-        cax = divider.append_axes("right", "5%", pad="3%")
-        plt.colorbar(mat, cax=cax)
+            ax.set_xlabel(title)
+        figure.colorbar(mat, ax=ax)
+        plt.tight_layout()
         return figure
 
-    def pdplot(self, title: titleType =None, ax: matplotlib.axes.SubplotBase =None):
+    def pdplot(
+        self, title: titleType = None, ax: matplotlib.axes.SubplotBase = None
+    ) -> matplotlib.figure.Figure:
         """Payoff difference heatmap to visualize the distributions of how
         players attain their payoffs."""
         matrix, names = self._pdplot_dataset
         return self._payoff_heatmap(matrix, names, title=title, ax=ax)
 
-    def payoff(self, title: titleType =None, ax: matplotlib.axes.SubplotBase =None):
+    def payoff(
+        self, title: titleType = None, ax: matplotlib.axes.SubplotBase = None
+    ) -> matplotlib.figure.Figure:
         """Payoff heatmap to visualize the distributions of how
         players attain their payoffs."""
         data = self._payoff_dataset
@@ -209,7 +228,10 @@ def payoff(self, title: titleType =None, ax: matplotlib.axes.SubplotBase =None):
 
     # Ecological Plot
 
-    def stackplot(self, eco, title: titleType =None, logscale: bool =True, ax: matplotlib.axes.SubplotBase =None) -> matplotlib.figure.Figure:
+    def stackplot(
+        self, eco, title: titleType = None, logscale: bool = True,
+        ax: matplotlib.axes.SubplotBase =None
+    ) -> matplotlib.figure.Figure:
         if not self.matplotlib_installed:  # pragma: no cover
             return None
 
@@ -222,26 +244,31 @@ def stackplot(self, eco, title: titleType =None, logscale: bool =True, ax: matpl
 
         figure = ax.get_figure()
         turns = range(len(populations))
-        pops = [[populations[iturn][ir] for iturn in turns] for ir in self.result_set.ranking]
+        pops = [
+            [populations[iturn][ir] for iturn in turns]
+            for ir in self.result_set.ranking
+        ]
         ax.stackplot(turns, *pops)
 
         ax.yaxis.tick_left()
         ax.yaxis.set_label_position("right")
         ax.yaxis.labelpad = 25.0
 
-        plt.ylim([0.0, 1.0])
-        plt.ylabel('Relative population size')
-        plt.xlabel('Turn')
+        ax.set_ylim([0.0, 1.0])
+        ax.set_ylabel('Relative population size')
+        ax.set_xlabel('Turn')
         if title is not None:
-            plt.title(title)
+            ax.set_title(title)
 
-        trans = transforms.blended_transform_factory(ax.transAxes, ax.transData)
+        trans = transforms.blended_transform_factory(
+            ax.transAxes, ax.transData)
         ticks = []
         for i, n in enumerate(self.result_set.ranked_names):
             x = -0.01
             y = (i + 0.5) * 1 / self.result_set.nplayers
-            ax.annotate(n, xy=(x, y), xycoords=trans, clip_on=False,
-                             va='center', ha='right', fontsize=5)
+            ax.annotate(
+                n, xy=(x, y), xycoords=trans, clip_on=False, va='center',
+                ha='right', fontsize=5)
             ticks.append(y)
         ax.set_yticks(ticks)
         ax.tick_params(direction='out')
@@ -250,10 +277,13 @@ def stackplot(self, eco, title: titleType =None, logscale: bool =True, ax: matpl
         if logscale:
             ax.set_xscale('log')
 
+        plt.tight_layout()
         return figure
 
-    def save_all_plots(self, prefix: str ="axelrod", title_prefix: str ="axelrod",
-                       filetype: str ="svg", progress_bar: bool =True):
+    def save_all_plots(
+        self, prefix: str ="axelrod", title_prefix: str ="axelrod",
+        filetype: str ="svg", progress_bar: bool = True
+    ) -> None:
         """
         A method to save all plots to file.