Add contour and tricontour plots to plot types

timhoffm · timhoffm · commit ae5f936a4136 · 2021-09-18T23:46:53.000+02:00
diff --git a/plot_types/arrays/contour.py b/plot_types/arrays/contour.py
@@ -0,0 +1,25 @@
+"""
+================
+contour(X, Y, Z)
+================
+
+See `~matplotlib.axes.Axes.contour`.
+"""
+import matplotlib.pyplot as plt
+import numpy as np
+
+plt.style.use('mpl_plot_gallery')
+
+# make data
+X, Y = np.meshgrid(np.linspace(-3, 3, 256), np.linspace(-3, 3, 256))
+Z = (1 - X/2 + X**5 + Y**3) * np.exp(-X**2 - Y**2)
+Z = Z - Z.min()
+levels = np.linspace(np.min(Z), np.max(Z), 7)
+
+# plot
+fig, ax = plt.subplots()
+ax.grid(False)
+
+ax.contour(X, Y, Z, levels=levels)
+
+plt.show()
diff --git a/plot_types/arrays/contourf.py b/plot_types/arrays/contourf.py
@@ -1,9 +1,9 @@
 """
-====================================
-contour(X, Y, Z) / contourf(X, Y, Z)
-====================================
+=================
+contourf(X, Y, Z)
+=================
 
-See `~matplotlib.axes.Axes.contour` / `~matplotlib.axes.Axes.contourf`.
+See `~matplotlib.axes.Axes.contourf`.
 """
 import matplotlib.pyplot as plt
 import numpy as np
@@ -12,14 +12,13 @@
 
 # make data
 X, Y = np.meshgrid(np.linspace(-3, 3, 256), np.linspace(-3, 3, 256))
-Z = (1 - X/2. + X**5 + Y**3) * np.exp(-X**2 - Y**2)
+Z = (1 - X/2 + X**5 + Y**3) * np.exp(-X**2 - Y**2)
 Z = Z - Z.min()
 levels = np.linspace(np.min(Z), np.max(Z), 7)
 
 # plot
 fig, ax = plt.subplots()
 
 ax.contourf(X, Y, Z, levels=levels)
-ax.contour(X, Y, Z, levels=levels, colors="white", linewidths=0.5)
 
 plt.show()
diff --git a/plot_types/unstructured/tricontour.py b/plot_types/unstructured/tricontour.py
@@ -12,7 +12,7 @@
 
 # make structured data
 X, Y = np.meshgrid(np.linspace(-3, 3, 256), np.linspace(-3, 3, 256))
-Z = (1 - X/2. + X**5 + Y**3) * np.exp(-X**2 - Y**2)
+Z = (1 - X/2 + X**5 + Y**3) * np.exp(-X**2 - Y**2)
 Z = Z - Z.min()
 
 # sample it to make unstructured x, y, z
@@ -25,10 +25,11 @@
 
 # plot:
 fig, ax = plt.subplots()
+ax.grid(False)
 
-ax.plot(x, y, '.k', alpha=0.5)
+ax.plot(x, y, 'o', markersize=2, color='lightgrey')
 levels = np.linspace(np.min(Z), np.max(Z), 7)
-ax.tricontourf(x, y, z, levels=levels)
+ax.tricontour(x, y, z, levels=levels)
 
 ax.set(xlim=(-3, 3), ylim=(-3, 3))
 
diff --git a/plot_types/unstructured/tricontourf.py b/plot_types/unstructured/tricontourf.py
@@ -0,0 +1,36 @@
+"""
+==========================================
+tricontour(x, y, z) / tricontourf(x, y, z)
+==========================================
+
+See `~matplotlib.axes.Axes.tricontour` / `~matplotlib.axes.Axes.tricontourf`.
+"""
+import matplotlib.pyplot as plt
+import numpy as np
+
+plt.style.use('mpl_plot_gallery')
+
+# make structured data
+X, Y = np.meshgrid(np.linspace(-3, 3, 256), np.linspace(-3, 3, 256))
+Z = (1 - X/2. + X**5 + Y**3) * np.exp(-X**2 - Y**2)
+Z = Z - Z.min()
+
+# sample it to make unstructured x, y, z
+np.random.seed(1)
+ysamp = np.random.randint(0, 256, size=250)
+xsamp = np.random.randint(0, 256, size=250)
+y = Y[:, 0][ysamp]
+x = X[0, :][xsamp]
+z = Z[ysamp, xsamp]
+
+# plot:
+fig, ax = plt.subplots()
+ax.grid(False)
+
+ax.plot(x, y, 'o', markersize=2, color='grey')
+levels = np.linspace(np.min(Z), np.max(Z), 7)
+ax.tricontourf(x, y, z, levels=levels)
+
+ax.set(xlim=(-3, 3), ylim=(-3, 3))
+
+plt.show()
