iris.quickplot: Ensure labels and plot use the same axes

lib/iris/quickplot.py

@@ -1,4 +1,4 @@
-# (C) British Crown Copyright 2010 - 2016, Met Office
+# (C) British Crown Copyright 2010 - 2018, Met Office
 #
 # This file is part of Iris.
 #
@@ -61,10 +61,13 @@ def _title(cube_or_coord, with_units):
     return title
 
 
-def _label(cube, mode, result=None, ndims=2, coords=None):
+def _label(cube, mode, result=None, ndims=2, coords=None, axes=None):
     """Puts labels on the current plot using the given cube."""
 
-    plt.title(_title(cube, with_units=False))
+    if axes is None:
+        axes = plt.gca()
+
+    axes.set_title(_title(cube, with_units=False))
 
     if result is not None:
         draw_edges = mode == iris.coords.POINT_MODE
@@ -89,22 +92,24 @@ def _label(cube, mode, result=None, ndims=2, coords=None):
 
     if ndims == 2:
         if not iplt._can_draw_map(plot_defn.coords):
-            plt.ylabel(_title(plot_defn.coords[0], with_units=True))
-            plt.xlabel(_title(plot_defn.coords[1], with_units=True))
+            axes.set_ylabel(_title(plot_defn.coords[0], with_units=True))
+            axes.set_xlabel(_title(plot_defn.coords[1], with_units=True))
     elif ndims == 1:
-        plt.xlabel(_title(plot_defn.coords[0], with_units=True))
-        plt.ylabel(_title(cube, with_units=True))
+        axes.set_xlabel(_title(plot_defn.coords[0], with_units=True))
+        axes.set_ylabel(_title(cube, with_units=True))
     else:
-        msg = 'Unexpected number of dimensions (%s) given to _label.' % ndims
+        msg = 'Unexpected number of dimensions ({}) given to ' \
+              '_label.'.format(ndims)
         raise ValueError(msg)
 
 
-def _label_with_bounds(cube, result=None, ndims=2, coords=None):
-    _label(cube, iris.coords.BOUND_MODE, result, ndims, coords)
+
+def _label_with_bounds(cube, result=None, ndims=2, coords=None, axes=None):
+    _label(cube, iris.coords.BOUND_MODE, result, ndims, coords, axes)
 
 
-def _label_with_points(cube, result=None, ndims=2, coords=None):
-    _label(cube, iris.coords.POINT_MODE, result, ndims, coords)
+def _label_with_points(cube, result=None, ndims=2, coords=None, axes=None):
+    _label(cube, iris.coords.POINT_MODE, result, ndims, coords, axes)
 
 
 def _get_titles(u_object, v_object):
@@ -126,15 +131,26 @@ def _get_titles(u_object, v_object):
     return xlabel, ylabel, title
 
 
-def _label_1d_plot(*args):
+def _label_1d_plot(*args, **kwargs):
     if len(args) > 1 and isinstance(args[1],
                                     (iris.cube.Cube, iris.coords.Coord)):
         xlabel, ylabel, title = _get_titles(*args[:2])
     else:
         xlabel, ylabel, title = _get_titles(None, args[0])
-    plt.title(title)
-    plt.xlabel(xlabel)
-    plt.ylabel(ylabel)
+
+    axes = kwargs.pop('axes', None)
+
+    if len(kwargs) != 0:
+        msg = 'Unexpected kwargs {} given to _label_1d_plot'.format(
+            kwargs.keys())
+        raise ValueError(msg)
+
+    if axes is None:
+        axes = plt.gca()
+
+    axes.set_title(title)
+    axes.set_xlabel(xlabel)
+    axes.set_ylabel(ylabel)
 
 
 def contour(cube, *args, **kwargs):
@@ -158,8 +174,9 @@ def contour(cube, *args, **kwargs):
 
     """
     coords = kwargs.get('coords')
+    axes = kwargs.get('axes')
     result = iplt.contour(cube, *args, **kwargs)
-    _label_with_points(cube, coords=coords)
+    _label_with_points(cube, coords=coords, axes=axes)
     return result
 
 
@@ -184,12 +201,13 @@ def contourf(cube, *args, **kwargs):
 
     """
     coords = kwargs.get('coords')
+    axes = kwargs.get('axes')
     result = iplt.contourf(cube, *args, **kwargs)
-    _label_with_points(cube, result, coords=coords)
+    _label_with_points(cube, result, coords=coords, axes=axes)
     return result
 
 
-def outline(cube, coords=None, color='k', linewidth=None):
+def outline(cube, coords=None, color='k', linewidth=None, axes=None):
     """
     Draws cell outlines on a labelled plot based on the given Cube.
 
@@ -211,9 +229,9 @@ def outline(cube, coords=None, color='k', linewidth=None):
 
     """
     result = iplt.outline(cube, color=color, linewidth=linewidth,
-                          coords=coords)
+                          coords=coords, axes=axes)
 
-    _label_with_bounds(cube, coords=coords)
+    _label_with_bounds(cube, coords=coords, axes=axes)
     return result
 
 
@@ -225,8 +243,9 @@ def pcolor(cube, *args, **kwargs):
 
     """
     coords = kwargs.get('coords')
+    axes = kwargs.get('axes')
     result = iplt.pcolor(cube, *args, **kwargs)
-    _label_with_bounds(cube, result, coords=coords)
+    _label_with_bounds(cube, result, coords=coords, axes=axes)
     return result
 
 
@@ -238,8 +257,9 @@ def pcolormesh(cube, *args, **kwargs):
 
     """
     coords = kwargs.get('coords')
+    axes = kwargs.get('axes')
     result = iplt.pcolormesh(cube, *args, **kwargs)
-    _label_with_bounds(cube, result, coords=coords)
+    _label_with_bounds(cube, result, coords=coords, axes=axes)
     return result
 
 
@@ -251,8 +271,9 @@ def points(cube, *args, **kwargs):
 
     """
     coords = kwargs.get('coords')
+    axes = kwargs.get('axes')
     result = iplt.points(cube, *args, **kwargs)
-    _label_with_points(cube, coords=coords)
+    _label_with_points(cube, coords=coords, axes=axes)
     return result
 
 
@@ -265,8 +286,9 @@ def plot(*args, **kwargs):
     keyword arguments.
 
     """
+    axes = kwargs.get('axes')
     result = iplt.plot(*args, **kwargs)
-    _label_1d_plot(*args)
+    _label_1d_plot(*args, axes=axes)
     return result
 
 
@@ -279,8 +301,9 @@ def scatter(x, y, *args, **kwargs):
     keyword arguments.
 
     """
+    axes = kwargs.get('axes')
     result = iplt.scatter(x, y, *args, **kwargs)
-    _label_1d_plot(x, y)
+    _label_1d_plot(x, y, axes=axes)
     return result
 
 

lib/iris/tests/test_quickplot.py

@@ -146,6 +146,30 @@ def test_contourf(self):
         qplt.contourf(self._small(), coords=['grid_longitude', 'model_level_number'])
         self.check_graphic()
 
+    def test_contourf_axes_specified(self):
+        # Check that the contourf function does not modify the matplotlib
+        # pyplot state machine.
+        fig1 = plt.figure()
+        axes = plt.axes()
+
+        # Create another figure and axes which will be the new results
+        # of plt.gcf and plt.gca.
+        fig2 = plt.figure()
+        axes2 = plt.axes()
+
+        # Add a title to the axes we don't care about.
+        plt.title('This should not be changed')
+        # Draw the contour on a specific axes.
+        qplt.contourf(self._small(), axes=axes)
+
+        # Ensure that the correct axes got the appropriate title.
+        self.assertEqual(axes2.get_title(), 'This should not be changed')
+        self.assertEqual(axes.get_title(), 'Air potential temperature')
+
+        # Check that the axes labels were set correctly.
+        self.assertEqual(axes.get_xlabel(), 'Grid longitude / degrees')
+        self.assertEqual(axes.get_ylabel(), 'Altitude / m')
+
     def test_contourf_nameless(self):
         cube = self._small()
         cube.standard_name = None