Don't have streamplot error on a zero field (#4188)

* Don't have streamplot error on a zero field

* Use approximately zero.

Author: dham

firedrake/pyplot/mpl.py

@@ -392,7 +392,12 @@ def streamline(function, point, direction=+1, tolerance=3e-3, loc_tolerance=1e-1
     x = np.array(point)
     v1 = toreal(direction * function.at(x, tolerance=loc_tolerance), complex_component)
     r = toreal(cell_sizes.at(x, tolerance=loc_tolerance), "real")
-    dt = 0.5 * r / np.sqrt(np.sum(v1**2))
+    v1norm = np.sqrt(np.sum(v1**2))
+    if np.isclose(v1norm, 0.0):
+        # Bail early for zero fields.
+        return
+
+    dt = 0.5 * r / v1norm
 
     while True:
         try:
@@ -622,7 +627,10 @@ def streamplot(function, resolution=None, min_length=None, max_time=None,
     if max_time is None:
         area = assemble(Constant(1) * dx(mesh))
         average_speed = np.sqrt(assemble(inner(function, function) * dx) / area)
-        max_time = 50 * min_length / average_speed
+        if np.isclose(average_speed, 0.0):
+            max_time = 50 * min_length / average_speed
+        else:
+            max_time = 0.
 
     streamplotter = Streamplotter(function, resolution, min_length, max_time,
                                   tolerance, loc_tolerance,
@@ -666,9 +674,13 @@ def streamplot(function, resolution=None, min_length=None, max_time=None,
     widths = np.array(widths)
 
     points = np.asarray(points)
-    vmin = kwargs.pop("vmin", speeds.min())
-    vmax = kwargs.pop("vmax", speeds.max())
-    norm = kwargs.pop("norm", matplotlib.colors.Normalize(vmin=vmin, vmax=vmax))
+    if speeds.size > 0:
+        vmin = kwargs.pop("vmin", speeds.min())
+        vmax = kwargs.pop("vmax", speeds.max())
+        norm = kwargs.pop("norm", matplotlib.colors.Normalize(vmin=vmin, vmax=vmax))
+    else:
+        norm = None
+
     cmap = plt.get_cmap(kwargs.pop("cmap", None))
 
     collection = LineCollection(points, cmap=cmap, norm=norm, linewidth=widths)

tests/firedrake/output/test_plotting.py

@@ -191,18 +191,20 @@ def test_streamplot():
     saddle = assemble(interpolate(2 * as_vector((v[0], -v[1])), V))
     r = Constant(.5)
     sink = assemble(interpolate(center - r * v, V))
+    zero = Function(V)
 
-    fig, axes = plt.subplots(ncols=1, nrows=3, sharex=True, sharey=True)
+    fig, axes = plt.subplots(ncols=1, nrows=4, sharex=True, sharey=True)
     for ax in axes:
         ax.set_aspect("equal")
 
     color_norm = matplotlib.colors.PowerNorm(gamma=0.5)
     kwargses = [
         {'resolution': 1/48, 'tolerance': 2e-2, 'norm': color_norm, 'seed': 0},
         {'loc_tolerance': 1e-5, 'cmap': 'bone', 'vmax': 1., 'seed': 0},
-        {'min_length': 1/4, 'max_time': 5., 'seed': 0}
+        {'min_length': 1/4, 'max_time': 5., 'seed': 0},
+        {}
     ]
-    for ax, function, kwargs in zip(axes, [center, saddle, sink], kwargses):
+    for ax, function, kwargs in zip(axes, [center, saddle, sink, zero], kwargses):
         lines = streamplot(function, axes=ax, **kwargs)
         colorbar = fig.colorbar(lines, ax=ax)
         assert lines is not None