New plotting methods for VertexOnlyMesh

firedrake/pyplot/__init__.py

@@ -1,10 +1,10 @@
 from firedrake.pyplot.mpl import (
-    plot, triplot, tricontourf, tricontour, trisurf, tripcolor, quiver,
+    plot, triplot, pointplot, tricontourf, tricontour, trisurf, tripcolor, quiver,
     streamplot, FunctionPlotter
 )
 from firedrake.pyplot.pgf import pgfplot
 
 __all__ = [
-    "plot", "triplot", "tricontourf", "tricontour", "trisurf", "tripcolor",
+    "plot", "triplot", "pointplot", "tricontourf", "tricontour", "trisurf", "tripcolor",
     "quiver", "streamplot", "FunctionPlotter", "pgfplot"
 ]

firedrake/pyplot/mpl.py

@@ -21,13 +21,13 @@
 from firedrake import (interpolate, sqrt, inner, Function, SpatialCoordinate,
                        FunctionSpace, VectorFunctionSpace, PointNotInDomainError,
                        Constant, assemble, dx)
-from firedrake.mesh import MeshGeometry
+from firedrake.mesh import MeshGeometry, VertexOnlyMeshTopology
 from firedrake.petsc import PETSc
 from ufl.domain import extract_unique_domain
 
 
 __all__ = [
-    "plot", "triplot", "tricontourf", "tricontour", "trisurf", "tripcolor",
+    "plot", "triplot", "pointplot", "tricontourf", "tricontour", "trisurf", "tripcolor",
     "quiver", "streamplot", "FunctionPlotter"
 ]
 
@@ -83,6 +83,54 @@ def _get_collection_types(gdim, tdim):
     raise ValueError("Geometric dimension must be either 2 or 3!")
 
 
+@PETSc.Log.EventDecorator()
+def pointplot(vom, axes=None, **kwargs):
+    r"""Plot a VertexOnlyMesh as a scatter plot.
+
+    :arg vom: either `~.MeshGeometry` object that is a VertexOnlyMesh or a `~.Function` defined on it
+    :arg axes: matplotlib :class:`Axes <matplotlib.axes.Axes>` object on which to plot the mesh
+    :arg kwargs: keyword arguments passed to :meth:`scatter <matplotlib.axes.Axes.scatter>`
+    :return: a :class:`matplotlib.collections.PathCollection` artist.
+    """
+    if not (isinstance(vom, Function)
+            and isinstance(vom.function_space().mesh().topology, VertexOnlyMeshTopology)) \
+            and not isinstance(vom.topology, VertexOnlyMeshTopology):
+        raise TypeError("Expected a VertexOnlyMesh or a Function defined on it.")
+
+    if isinstance(vom, Function):
+        if len(vom.ufl_shape) == 0:
+            # scalar field: colour points by value
+            kwargs["c"] = vom.dat.data_ro
+        elif len(vom.ufl_shape) == 1:
+            # vector field: use quiver instead
+            raise TypeError("Use `quiver` to plot vector-valued fields on a VertexOnlyMesh.")
+        else:
+            raise TypeError(
+                f"Cannot plot a rank-{len(vom.ufl_shape)} tensor field; "
+                "only scalar fields are supported in this method. "
+                "Alternatively, use `quiver` to plot vector-valued fields.")
+        vom = vom.function_space().mesh()
+
+    gdim = vom.geometric_dimension
+    coords = toreal(vom.coordinates.dat.data_ro_with_halos, "real")
+
+    if axes is None:
+        fig = plt.figure()
+        if gdim == 3:
+            axes = fig.add_subplot(111, projection="3d")
+        else:
+            axes = fig.add_subplot(111)
+
+    kwargs.setdefault("zorder", 5)  # this makes sure that points are drawn on top of the parent mesh lines
+    kwargs.setdefault("s", 10)  # controls scatter dot size
+    kwargs.setdefault("c", "red")  # default colour if no function provided
+
+    collection = axes.scatter(*(coords.T), **kwargs)
+
+    _autoscale_view(axes, coords)
+    return collection
+
+
 @PETSc.Log.EventDecorator()
 def triplot(mesh, axes=None, interior_kw={}, boundary_kw={}):
     r"""Plot a mesh colouring marked facet segments
@@ -180,6 +228,7 @@ def facet_data(typ):
     if tdim == 3:
         boundary_kw["edgecolors"] = boundary_kw.get("edgecolors", "k")
         boundary_kw["linewidths"] = boundary_kw.get("linewidths", 1.0)
+
     for marker, color in zip(markers, colors):
         vertices = []
         for typ in ["interior", "exterior"]:
@@ -353,10 +402,15 @@ def quiver(function, *, complex_component="real", **kwargs):
         figure = plt.figure()
         axes = figure.add_subplot(111)
 
-    coords = toreal(extract_unique_domain(function).coordinates.dat.data_ro, "real")
-    V = extract_unique_domain(function).coordinates.function_space()
-    function_interp = assemble(interpolate(function, V))
-    vals = toreal(function_interp.dat.data_ro, complex_component)
+    mesh = function.function_space().mesh()
+    coords = toreal(mesh.coordinates.dat.data_ro, "real")
+    if isinstance(mesh.topology, VertexOnlyMeshTopology):
+        vals = toreal(function.dat.data_ro, complex_component)
+    else:
+        V = mesh.coordinates.function_space()
+        function_interp = assemble(interpolate(function, V))
+        vals = toreal(function_interp.dat.data_ro, complex_component)
+
     C = np.linalg.norm(vals, axis=1)
     return axes.quiver(*(coords.T), *(vals.T), C, **kwargs)
 

tests/firedrake/output/test_plotting.py

@@ -195,6 +195,21 @@ def test_quiver_plot():
     fig.colorbar(arrows)
 
 
+@pytest.mark.skipplot
+def test_quiver_plot_vom():
+    mesh = UnitSquareMesh(10, 10)
+    vom = VertexOnlyMesh(mesh, [[0.5, 0.5], [0.2, 0.8], [0.9, 0.1]])
+    V = VectorFunctionSpace(vom, "DG", 0)
+    f = Function(V)
+    x = SpatialCoordinate(mesh)
+    f.interpolate(as_vector((-x[1], x[0])))
+
+    fig, axes = plt.subplots()
+    arrows = quiver(f, axes=axes)
+    assert arrows is not None
+    fig.colorbar(arrows)
+
+
 @pytest.mark.skipplot
 def test_streamplot():
     mesh = UnitSquareMesh(10, 10)
@@ -371,3 +386,41 @@ def animate(time):
 
     # Use a method of the animation to prevent warning about it being unused
     movie.to_jshtml()
+
+
+@pytest.mark.skipplot
+def test_pointplot():
+    mesh = UnitSquareMesh(10, 10)
+    vom = VertexOnlyMesh(mesh, [[0.5, 0.5], [0.2, 0.8], [0.9, 0.1]])
+
+    fig, axes = plt.subplots()
+    sc = pointplot(vom, axes=axes)
+
+    assert sc is not None
+    assert len(sc.get_offsets()) == vom.num_vertices()
+
+
+@pytest.mark.skipplot
+def test_pointplot_3d():
+    mesh = UnitCubeMesh(5, 5, 5)
+    coords_3d = np.random.rand(20, 3)
+    vom = VertexOnlyMesh(mesh, coords_3d)
+
+    fig = plt.figure()
+    axes = fig.add_subplot(111, projection='3d')
+    sc = pointplot(vom, axes=axes)
+    assert sc is not None
+    assert len(sc.get_offsets()) == vom.num_vertices()
+
+
+@pytest.mark.skipplot
+def test_pointplot_scalar_field():
+    mesh = UnitSquareMesh(10, 10)
+    vom = VertexOnlyMesh(mesh, [[0.5, 0.5], [0.2, 0.8], [0.9, 0.1]])
+    V = FunctionSpace(vom, "DG", 0)
+    f = Function(V)
+    f.dat.data[:] = [1.0, 2.0, 3.0]
+
+    fig, axes = plt.subplots()
+    sc = pointplot(vom=f, axes=axes)
+    assert np.allclose(sc.get_array(), f.dat.data_ro)