DOC Cleaning up mplot3d/tri* examples: comments/docstrings, tweaks, and refactoring. Notably, viewing angles on tricontour examples are tweaked and the graphs of trisurf3d_demo2 are made into subplots. [MEP12]

TrishGillett · TrishGillett · commit 6a5def2d4d23 · 2016-07-05T02:20:17.000-04:00
diff --git a/examples/mplot3d/tricontour3d_demo.py b/examples/mplot3d/tricontour3d_demo.py
@@ -1,27 +1,30 @@
 """
 Contour plots of unstructured triangular grids.
+
+The data used is the same as in the second plot of trisurf3d_demo2.
+tricontourf3d_demo shows the filled version of this example.
 """
+
 import matplotlib.pyplot as plt
 from mpl_toolkits.mplot3d import Axes3D
 import matplotlib.tri as tri
 import numpy as np
-import math
 
-# First create the x and y coordinates of the points.
 n_angles = 48
 n_radii = 8
 min_radius = 0.25
-radii = np.linspace(min_radius, 0.95, n_radii)
 
-angles = np.linspace(0, 2*math.pi, n_angles, endpoint=False)
+# Create the mesh in polar coordinates and compute x, y, z.
+radii = np.linspace(min_radius, 0.95, n_radii)
+angles = np.linspace(0, 2*np.pi, n_angles, endpoint=False)
 angles = np.repeat(angles[..., np.newaxis], n_radii, axis=1)
-angles[:, 1::2] += math.pi/n_angles
+angles[:, 1::2] += np.pi/n_angles
 
 x = (radii*np.cos(angles)).flatten()
 y = (radii*np.sin(angles)).flatten()
 z = (np.cos(radii)*np.cos(angles*3.0)).flatten()
 
-# Create a custom triangulation
+# Create a custom triangulation.
 triang = tri.Triangulation(x, y)
 
 # Mask off unwanted triangles.
@@ -30,7 +33,11 @@
 mask = np.where(xmid*xmid + ymid*ymid < min_radius*min_radius, 1, 0)
 triang.set_mask(mask)
 
-plt.figure()
-plt.gca(projection='3d')
-plt.tricontour(triang, z)
+fig = plt.figure()
+ax = fig.gca(projection='3d')
+ax.tricontour(triang, z, cmap=plt.cm.CMRmap)
+
+# Customize the view angle so it's easier to understand the plot.
+ax.view_init(elev=45.)
+
 plt.show()
diff --git a/examples/mplot3d/tricontourf3d_demo.py b/examples/mplot3d/tricontourf3d_demo.py
@@ -1,27 +1,31 @@
 """
-Contour plots of unstructured triangular grids.
+Filled contour plots of unstructured triangular grids.
+
+The data used is the same as in the second plot of trisurf3d_demo2.
+tricontour3d_demo shows the unfilled version of this example.
 """
+
 import matplotlib.pyplot as plt
 from mpl_toolkits.mplot3d import Axes3D
 import matplotlib.tri as tri
 import numpy as np
-import math
 
-# First create the x and y coordinates of the points.
+# First create the x, y, z coordinates of the points.
 n_angles = 48
 n_radii = 8
 min_radius = 0.25
-radii = np.linspace(min_radius, 0.95, n_radii)
 
-angles = np.linspace(0, 2*math.pi, n_angles, endpoint=False)
+# Create the mesh in polar coordinates and compute x, y, z.
+radii = np.linspace(min_radius, 0.95, n_radii)
+angles = np.linspace(0, 2*np.pi, n_angles, endpoint=False)
 angles = np.repeat(angles[..., np.newaxis], n_radii, axis=1)
-angles[:, 1::2] += math.pi/n_angles
+angles[:, 1::2] += np.pi/n_angles
 
 x = (radii*np.cos(angles)).flatten()
 y = (radii*np.sin(angles)).flatten()
 z = (np.cos(radii)*np.cos(angles*3.0)).flatten()
 
-# Create a custom triangulation
+# Create a custom triangulation.
 triang = tri.Triangulation(x, y)
 
 # Mask off unwanted triangles.
@@ -30,7 +34,11 @@
 mask = np.where(xmid*xmid + ymid*ymid < min_radius*min_radius, 1, 0)
 triang.set_mask(mask)
 
-plt.figure()
-plt.gca(projection='3d')
-plt.tricontourf(triang, z)
+fig = plt.figure()
+ax = fig.gca(projection='3d')
+ax.tricontourf(triang, z, cmap=plt.cm.CMRmap)
+
+# Customize the view angle so it's easier to understand the plot.
+ax.view_init(elev=45.)
+
 plt.show()
diff --git a/examples/mplot3d/trisurf3d_demo.py b/examples/mplot3d/trisurf3d_demo.py
@@ -1,27 +1,29 @@
+'''
+Plot a 3D surface with a triangular mesh.
+'''
+
 from mpl_toolkits.mplot3d import Axes3D
-from matplotlib import cm
 import matplotlib.pyplot as plt
 import numpy as np
 
-n_angles = 36
+
 n_radii = 8
+n_angles = 36
 
-# An array of radii
-# Does not include radius r=0, this is to eliminate duplicate points
+# Make radii and angles spaces (radius r=0 omitted to eliminate duplication).
 radii = np.linspace(0.125, 1.0, n_radii)
-
-# An array of angles
 angles = np.linspace(0, 2*np.pi, n_angles, endpoint=False)
 
-# Repeat all angles for each radius
+# Repeat all angles for each radius.
 angles = np.repeat(angles[..., np.newaxis], n_radii, axis=1)
 
-# Convert polar (radii, angles) coords to cartesian (x, y) coords
-# (0, 0) is added here. There are no duplicate points in the (x, y) plane
+# Convert polar (radii, angles) coords to cartesian (x, y) coords.
+# (0, 0) is manually added at this stage,  so there will be no duplicate
+# points in the (x, y) plane.
 x = np.append(0, (radii*np.cos(angles)).flatten())
 y = np.append(0, (radii*np.sin(angles)).flatten())
 
-# Pringle surface
+# Compute z to make the pringle surface.
 z = np.sin(-x*y)
 
 fig = plt.figure()
diff --git a/examples/mplot3d/trisurf3d_demo2.py b/examples/mplot3d/trisurf3d_demo2.py
@@ -1,11 +1,28 @@
+'''
+Two additional examples of plotting surfaces with triangular mesh.
+
+The first demonstrates use of plot_trisurf's triangles argument, and the
+second sets a Triangulation object's mask and passes the object directly
+to plot_trisurf.
+'''
+
 import numpy as np
 import matplotlib.pyplot as plt
 from mpl_toolkits.mplot3d import Axes3D
 import matplotlib.tri as mtri
 
-# u, v are parameterisation variables
-u = (np.linspace(0, 2.0 * np.pi, endpoint=True, num=50) * np.ones((10, 1))).flatten()
-v = np.repeat(np.linspace(-0.5, 0.5, endpoint=True, num=10), repeats=50).flatten()
+
+fig = plt.figure(figsize=plt.figaspect(0.5))
+
+#============
+# First plot
+#============
+
+# Make a mesh in the space of parameterisation variables u and v
+u = np.linspace(0, 2.0 * np.pi, endpoint=True, num=50)
+v = np.linspace(-0.5, 0.5, endpoint=True, num=10)
+u, v = np.meshgrid(u, v)
+u, v = u.flatten(), v.flatten()
 
 # This is the Mobius mapping, taking a u, v pair and returning an x, y, z
 # triple
@@ -16,16 +33,18 @@
 # Triangulate parameter space to determine the triangles
 tri = mtri.Triangulation(u, v)
 
-fig = plt.figure()
-ax = fig.add_subplot(1, 1, 1, projection='3d')
-
-# The triangles in parameter space determine which x, y, z points are
-# connected by an edge
+# Plot the surface.  The triangles in parameter space determine which x, y, z
+# points are connected by an edge.
+ax = fig.add_subplot(1, 2, 1, projection='3d')
 ax.plot_trisurf(x, y, z, triangles=tri.triangles, cmap=plt.cm.Spectral)
-
 ax.set_zlim(-1, 1)
 
-# First create the x and y coordinates of the points.
+
+#============
+# Second plot
+#============
+
+# Make parameter spaces radii and angles.
 n_angles = 36
 n_radii = 8
 min_radius = 0.25
@@ -35,6 +54,7 @@
 angles = np.repeat(angles[..., np.newaxis], n_radii, axis=1)
 angles[:, 1::2] += np.pi/n_angles
 
+# Map radius, angle pairs to x, y, z points.
 x = (radii*np.cos(angles)).flatten()
 y = (radii*np.sin(angles)).flatten()
 z = (np.cos(radii)*np.cos(angles*3.0)).flatten()
@@ -45,11 +65,12 @@
 # Mask off unwanted triangles.
 xmid = x[triang.triangles].mean(axis=1)
 ymid = y[triang.triangles].mean(axis=1)
-mask = np.where(xmid*xmid + ymid*ymid < min_radius*min_radius, 1, 0)
+mask = np.where(xmid**2 + ymid**2 < min_radius**2, 1, 0)
 triang.set_mask(mask)
 
-# tripcolor plot.
-fig = plt.figure()
-ax = fig.add_subplot(1, 1, 1, projection='3d')
+# Plot the surface.
+ax = fig.add_subplot(1, 2, 2, projection='3d')
 ax.plot_trisurf(triang, z, cmap=plt.cm.CMRmap)
+
+
 plt.show()
