Prevent color bar from being misplaced (#4894)

* Prevent color bar from being misplaced

Currently, if `plt.gca()` gives a different Axes object than `axes` in the function `_label`, the color bar will be malpositioned, ending up either on another subfigure, or in another figure window altogether. Specifying the `ax` argument when calling `plt.colorbar` prevents this.

* add whatsnew and test coverage

* check axes and figure

* tidy and add extra axes2 test

Co-authored-by: Bill Little <[email protected]>

Author: krikru

docs/src/whatsnew/latest.rst

@@ -25,7 +25,7 @@ This document explains the changes made to Iris for this release
 📢 Announcements
 ================
 
-#. N/A
+#. Welcome to `@krikru`_ who made their first contribution to Iris 🎉
 
 
 ✨ Features
@@ -136,6 +136,11 @@ This document explains the changes made to Iris for this release
    array type.  This prevents masks being lost in some cases and therefore
    resolves :issue:`2987`.  (:pull:`3790`)
 
+#. `@krikru`_ and `@rcomer`_ updated :mod:`iris.quickplot` such that the
+   colorbar is added to the correct ``axes`` when specified as a keyword
+   argument to a plotting routine. Otherwise, by default the colorbar will be
+   added to the current axes of the current figure. (:pull:`4894`)
+
 #. `@rcomer`_ and `@bjlittle`_ (reviewer) modified :func:`iris.util.mask_cube` so it
    either works in place or returns a new cube (:issue:`3717`, :pull:`4889`)
 
@@ -279,6 +284,7 @@ This document explains the changes made to Iris for this release
     core dev names are automatically included by the common_links.inc:
 
 .. _@evertrol: https://github.com/evertrol
+.. _@krikru: https://github.com/krikru
 
 
 .. comment

lib/iris/quickplot.py

@@ -71,7 +71,7 @@ def _label(cube, mode, result=None, ndims=2, coords=None, axes=None):
     if result is not None:
         draw_edges = mode == iris.coords.POINT_MODE
         bar = plt.colorbar(
-            result, orientation="horizontal", drawedges=draw_edges
+            result, ax=axes, orientation="horizontal", drawedges=draw_edges
         )
         has_known_units = not (
             cube.units.is_unknown() or cube.units.is_no_unit()

lib/iris/tests/test_quickplot.py

@@ -247,5 +247,39 @@ def test_not_reference_time_units(self):
         self.check_graphic()
 
 
+@tests.skip_data
+@tests.skip_plot
+class TestSubplotColorbar(tests.IrisTest):
+    def setUp(self):
+        theta = _load_theta()
+        coords = ["model_level_number", "grid_longitude"]
+        self.data = next(theta.slices(coords))
+        spec = (1, 1, 1)
+        self.figure1 = plt.figure()
+        self.axes1 = self.figure1.add_subplot(*spec)
+        self.figure2 = plt.figure()
+        self.axes2 = self.figure2.add_subplot(*spec)
+
+    def _check(self, mappable, figure, axes):
+        self.assertIs(mappable.axes, axes)
+        self.assertIs(mappable.colorbar.mappable, mappable)
+        self.assertIs(mappable.colorbar.ax.get_figure(), figure)
+
+    def test_with_axes1(self):
+        # plot using the first figure subplot axes (explicit)
+        mappable = qplt.contourf(self.data, axes=self.axes1)
+        self._check(mappable, self.figure1, self.axes1)
+
+    def test_with_axes2(self):
+        # plot using the second figure subplot axes (explicit)
+        mappable = qplt.contourf(self.data, axes=self.axes2)
+        self._check(mappable, self.figure2, self.axes2)
+
+    def test_without_axes__default(self):
+        # plot using the second/last figure subplot axes (default)
+        mappable = qplt.contourf(self.data)
+        self._check(mappable, self.figure2, self.axes2)
+
+
 if __name__ == "__main__":
     tests.main()