Minor improvements to quadmesh_demo.

anntzer · anntzer · commit b465289cf023 · 2018-01-25T21:23:34.000-08:00
diff --git a/examples/images_contours_and_fields/quadmesh_demo.py b/examples/images_contours_and_fields/quadmesh_demo.py
@@ -9,37 +9,38 @@
 This demo illustrates a bug in quadmesh with masked data.
 """
 
+import copy
+
+from matplotlib import cm, colors, pyplot as plt
 import numpy as np
-from matplotlib.pyplot import figure, show, savefig
-from matplotlib import cm, colors
-from numpy import ma
 
 n = 12
 x = np.linspace(-1.5, 1.5, n)
 y = np.linspace(-1.5, 1.5, n * 2)
 X, Y = np.meshgrid(x, y)
 Qx = np.cos(Y) - np.cos(X)
 Qz = np.sin(Y) + np.sin(X)
-Qx = (Qx + 1.1)
 Z = np.sqrt(X**2 + Y**2) / 5
 Z = (Z - Z.min()) / (Z.max() - Z.min())
 
-# The color array can include masked values:
-Zm = ma.masked_where(np.abs(Qz) < 0.5 * np.max(Qz), Z)
+# The color array can include masked values.
+Zm = np.ma.masked_where(np.abs(Qz) < 0.5 * np.max(Qz), Z)
+
+fig, axs = plt.subplots(1, 3)
+axs[0].pcolormesh(Qx, Qz, Z, shading='gouraud')
+axs[0].set_title('Without masked values')
 
-fig = figure()
-ax = fig.add_subplot(121)
-ax.pcolormesh(Qx, Qz, Z, shading='gouraud')
-ax.set_title('Without masked values')
+# You can control the color of the masked region.  We copy the default colormap
+# before modifying it.
+cmap = copy.copy(cm.get_cmap(plt.rcParams['image.cmap']))
+cmap.set_bad('y', 1.0)
+axs[1].pcolormesh(Qx, Qz, Zm, shading='gouraud', cmap=cmap)
+axs[1].set_title('With masked values')
 
-ax = fig.add_subplot(122)
-#  You can control the color of the masked region:
-# cmap = cm.RdBu
-# cmap.set_bad('y', 1.0)
-# ax.pcolormesh(Qx, Qz, Zm, cmap=cmap)
-#  Or use the default, which is transparent:
-col = ax.pcolormesh(Qx, Qz, Zm, shading='gouraud')
-ax.set_title('With masked values')
+# Or use the default, which is transparent.
+axs[2].pcolormesh(Qx, Qz, Zm, shading='gouraud')
+axs[2].set_title('With masked values')
 
+fig.tight_layout()
 
-show()
+plt.show()
