Add plotting option to Fingerprint plot

.gitignore

@@ -81,7 +81,7 @@ target/
 
 # Jupyter Notebook
 .ipynb_checkpoints
-
+*.ipynb
 # pyenv
 .python-version
 

axelrod/fingerprint.py

@@ -280,15 +280,39 @@ def fingerprint(self, turns=50, repetitions=10, step=0.01, processes=None,
         self.data = self.generate_data(self.interactions, self.points, edges)
         return self.data
 
-    def plot(self, col_map='seismic', interpolation='none', title=None):
+    @staticmethod
+    def reshape_data(data, points, size):
+        """Shape the data so that it can be plotted easily.
+
+        Parameters
+        ----------
+        data : dictionary
+            A dictionary where the keys are Points of the form (x, y) and
+            the values are the mean score for the corresponding interactions.
+
+        points : list
+            of Point objects with coordinates (x, y).
+
+        size : int
+            The number of Points in every row/column.
+
+        Returns
+        ----------
+        plotting_data : list
+            2-D numpy array of the scores, correctly shaped to ensure that the
+            score corresponding to Point (0, 0) is in the left hand corner ie.
+            the standard origin.
+        """
+        ordered_data = [data[point] for point in points]
+        shaped_data = np.reshape(ordered_data, (size, size), order='F')
+        plotting_data = np.flipud(shaped_data)
+        return plotting_data
+
+    def plot(self, col_map='seismic', interpolation='none', title=None, colorbar=True, labels=True):
         """Plot the results of the spatial tournament.
 
         Parameters
         ----------
-        filename : str, optional
-            The location and name that the resulting plot should be saved to.
-            Defaults to the current directory with the name
-            `Strategy and Probe.pdf`
         col_map : str, optional
             A matplotlib colour map, full list can be found at
             http://matplotlib.org/examples/color/colormaps_reference.html
@@ -297,18 +321,36 @@ def plot(self, col_map='seismic', interpolation='none', title=None):
             http://matplotlib.org/examples/images_contours_and_fields/interpolation_methods.html
         title : str, optional
             A title for the plot
+        colorbar : bool, optional
+            Choose whether the colorbar should be included or not
+        labels : bool, optional
+            Choose whether the axis labels and ticks should be included
 
         Returns
         ----------
         figure : matplotlib figure
             A heat plot of the results of the spatial tournament
         """
         size = int((1 / self.step) // 1) + 1
-        ordered_data = [self.data[point] for point in self.points]
-        plotting_data = np.reshape(ordered_data, (size, size))
-        figure = plt.figure()
-        plt.imshow(plotting_data, cmap=col_map, interpolation=interpolation)
-        plt.axis('off')
+        plotting_data = self.reshape_data(self.data, self.points, size)
+        fig, ax = plt.subplots()
+        cax = ax.imshow(plotting_data, cmap=col_map, interpolation=interpolation)
+
+        if colorbar:
+            max_score = max(self.data.values())
+            min_score = min(self.data.values())
+            ticks = [min_score, (max_score + min_score) / 2, max_score]
+            fig.colorbar(cax, ticks=ticks)
+
+        plt.xlabel('$x$')
+        plt.ylabel('$y$', rotation=0)
+        ax.tick_params(axis='both', which='both', length=0)
+        plt.xticks([0, len(plotting_data) - 1], ['0', '1'])
+        plt.yticks([0, len(plotting_data) - 1], ['1', '0'])
+
+        if not labels:
+            plt.axis('off')
+
         if title is not None:
             plt.title(title)
-        return figure
+        return fig

axelrod/tests/unit/test_fingerprint.py

@@ -160,6 +160,33 @@ def test_generate_data(self):
         self.assertEqual(sorted(keys), self.expected_points)
         self.assertEqual(all(values), True)
 
+    def test_reshape_data(self):
+        test_points = [Point(x=0.0, y=0.0),
+                       Point(x=0.0, y=0.5),
+                       Point(x=0.0, y=1.0),
+                       Point(x=0.5, y=0.0),
+                       Point(x=0.5, y=0.5),
+                       Point(x=0.5, y=1.0),
+                       Point(x=1.0, y=0.0),
+                       Point(x=1.0, y=0.5),
+                       Point(x=1.0, y=1.0)]
+        test_data = {Point(x=0.0, y=0.0): 5,
+                     Point(x=0.0, y=0.5): 9,
+                     Point(x=0.0, y=1.0): 3,
+                     Point(x=0.5, y=0.0): 8,
+                     Point(x=0.5, y=0.5): 2,
+                     Point(x=0.5, y=1.0): 4,
+                     Point(x=1.0, y=0.0): 2,
+                     Point(x=1.0, y=0.5): 1,
+                     Point(x=1.0, y=1.0): 9}
+        test_shaped_data = [[3, 4, 9],
+                            [9, 2, 1],
+                            [5, 8, 2]]
+        af = AshlockFingerprint(self.strategy, self.probe)
+        plotting_data = af.reshape_data(test_data, test_points, 3)
+        for i in range(len(plotting_data)):
+            self.assertEqual(list(plotting_data[i]), test_shaped_data[i])
+
     def test_plot(self):
         af = AshlockFingerprint(self.strategy, self.probe)
         af.fingerprint(turns=10, repetitions=2, step=0.25, progress_bar=False)
@@ -170,7 +197,11 @@ def test_plot(self):
         r = af.plot(interpolation='bicubic')
         self.assertIsInstance(r, matplotlib.pyplot.Figure)
         t = af.plot(title='Title')
-        self.assertIsInstance(r, matplotlib.pyplot.Figure)
+        self.assertIsInstance(t, matplotlib.pyplot.Figure)
+        u = af.plot(colorbar=False)
+        self.assertIsInstance(u, matplotlib.pyplot.Figure)
+        v = af.plot(labels=False)
+        self.assertIsInstance(v, matplotlib.pyplot.Figure)
 
     def test_wsls_fingerprint(self):
         axl.seed(0)  # Fingerprinting is a random process

docs/tutorials/further_topics/fingerprinting.rst

@@ -53,9 +53,11 @@ We get the plot:
      :width: 100%
      :align: center
 
-We are also able to specify a matplotlib colour map and interpolation::
+We are also able to specify a matplotlib colour map, interpolation and can
+remove the colorbar and axis labels::
 
-    >>> af.plot(col_map='PuOr', interpolation='bicubic')  # doctest: +SKIP
+    >>> p = af.plot(col_map='PuOr', interpolation='bicubic', colorbar=False, labels=False)  # doctest: +SKIP
+    >>> p.show()
 
 .. image:: _static/fingerprinting/WSLS_large_alt.png
      :width: 100%