def mplot_expression(ax, f, mesh, **kwargs):
# TODO: Can probably avoid creating the function space here by
# restructuring mplot_function a bit so it can handle Expression
# natively
V = create_cg1_function_space(mesh, f.value_shape)
g = fem.interpolate(f, V)
return mplot_function(ax, g, **kwargs)
def _compute_uvalues(u, mesh):
# the whole purpose of this function is
# to have a scalar (or vector) for each point of the mesh
if not u: return
# print('u',u)
if hasattr(u, 'compute_vertex_values'): # old dolfin, works fine
u_values = u.compute_vertex_values(mesh)
if u.value_rank() and u.value_dimension(0) > 1:
l = u_values.shape[0]
u_values = u_values.reshape(u.value_dimension(0),
int(l / u.value_dimension(0))).T
elif hasattr(u, 'compute_point_values'): # dolfinx
u_values = u.compute_point_values()
try:
from dolfin import fem
fvec = u.vector
except RuntimeError:
fspace = u.function_space
try:
fspace = fspace.collapse()
except RuntimeError:
return []
fvec = fem.interpolate(u, fspace).vector
tdim = mesh.topology.dim
#print('fvec.getSize', fvec.getSize(), mesh.num_entities(tdim))
if fvec.getSize() == mesh.num_entities(tdim):
# DG0 cellwise function
C = fvec.get_local()
if (C.dtype.type is np.complex128):
print("Plotting real part of complex data")
C = np.real(C)
u_values = C
else:
u_values = []
if hasattr(mesh, "coordinates"):
coords = mesh.coordinates()
else:
coords = mesh.geometry.points
if u_values.shape[0] != coords.shape[0]:
printc('Warning: mismatch in vtkplotter.dolfin._compute_uvalues()',
c=1)
u_values = np.array([u(p) for p in coords])
return u_values
def mplot_function(ax, f, **kwargs):
mesh = f.function_space.mesh
gdim = mesh.geometry.dim
tdim = mesh.topology.dim
# Extract the function vector in a way that also works for
# subfunctions
try:
fvec = f.vector
except RuntimeError:
fspace = f.function_space
try:
fspace = fspace.collapse()
except RuntimeError:
return
fvec = fem.interpolate(f, fspace).vector
if fvec.getSize() == mesh.num_entities(tdim):
# DG0 cellwise function
C = fvec.get_local()
if (C.dtype.type is np.complex128):
warnings.warn("Plotting real part of complex data")
C = np.real(C)
# NB! Assuming here dof ordering matching cell numbering
if gdim == 2 and tdim == 2:
return ax.tripcolor(mesh2triang(mesh), C, **kwargs)
elif gdim == 3 and tdim == 2: # surface in 3d
# FIXME: Not tested, probably broken
xy = mesh.geometry.points
shade = kwargs.pop("shade", True)
return ax.plot_trisurf(mesh2triang(mesh),
xy[:, 2],
C,
shade=shade,
**kwargs)
elif gdim == 1 and tdim == 1:
x = mesh.geometry.points[:, 0]
nv = len(x)
# Insert duplicate points to get piecewise constant plot
xp = np.zeros(2 * nv - 2)
xp[0] = x[0]
xp[-1] = x[-1]
xp[1:2 * nv - 3:2] = x[1:-1]
xp[2:2 * nv - 2:2] = x[1:-1]
Cp = np.zeros(len(xp))
Cp[0:len(Cp) - 1:2] = C
Cp[1:len(Cp):2] = C
return ax.plot(xp, Cp, *kwargs)
# elif tdim == 1: # FIXME: Plot embedded line
else:
raise AttributeError(
'Matplotlib plotting backend only supports 2D mesh for scalar functions.'
)
elif f.value_rank == 0:
# Scalar function, interpolated to vertices
# TODO: Handle DG1?
C = f.compute_point_values()
if (C.dtype.type is np.complex128):
warnings.warn("Plotting real part of complex data")
C = np.real(C)
if gdim == 2 and tdim == 2:
mode = kwargs.pop("mode", "contourf")
if mode == "contourf":
levels = kwargs.pop("levels", 40)
return ax.tricontourf(mesh2triang(mesh), C[:, 0], levels,
**kwargs)
elif mode == "color":
shading = kwargs.pop("shading", "gouraud")
return ax.tripcolor(mesh2triang(mesh),
C[:, 0],
shading=shading,
**kwargs)
elif mode == "warp":
from matplotlib import cm
cmap = kwargs.pop("cmap", cm.jet)
linewidths = kwargs.pop("linewidths", 0)
return ax.plot_trisurf(mesh2triang(mesh),
C[:, 0],
cmap=cmap,
linewidths=linewidths,
**kwargs)
elif mode == "wireframe":
return ax.triplot(mesh2triang(mesh), **kwargs)
elif mode == "contour":
return ax.tricontour(mesh2triang(mesh), C[:, 0], **kwargs)
elif gdim == 3 and tdim == 2: # surface in 3d
# FIXME: Not tested
from matplotlib import cm
cmap = kwargs.pop("cmap", cm.jet)
return ax.plot_trisurf(mesh2triang(mesh),
C[:, 0],
cmap=cmap,
**kwargs)
elif gdim == 3 and tdim == 3:
# Volume
# TODO: Isosurfaces?
# Vertex point cloud
X = [mesh.geometrycoordinates[:, i] for i in range(gdim)]
return ax.scatter(*X, c=C, **kwargs)
elif gdim == 1 and tdim == 1:
x = mesh.geometry.points[:, 0]
ax.set_aspect('auto')
p = ax.plot(x, C[:, 0], **kwargs)
# Setting limits for Line2D objects
# Must be done after generating plot to avoid ignoring function
# range if no vmin/vmax are supplied
vmin = kwargs.pop("vmin", None)
vmax = kwargs.pop("vmax", None)
ax.set_ylim([vmin, vmax])
return p
# elif tdim == 1: # FIXME: Plot embedded line
else:
raise AttributeError(
'Matplotlib plotting backend only supports 2D mesh for scalar functions.'
)
elif f.value_rank == 1:
# Vector function, interpolated to vertices
w0 = f.compute_point_values()
if (w0.dtype.type is np.complex128):
warnings.warn("Plotting real part of complex data")
w0 = np.real(w0)
nv = mesh.num_entities(0)
if w0.shape[1] != gdim:
raise AttributeError(
'Vector length must match geometric dimension.')
X = mesh.geometry.points
X = [X[:, i] for i in range(gdim)]
U = [x for x in w0.T]
# Compute magnitude
C = U[0]**2
for i in range(1, gdim):
C += U[i]**2
C = np.sqrt(C)
mode = kwargs.pop("mode", "glyphs")
if mode == "glyphs":
args = X + U + [C]
if gdim == 3:
length = kwargs.pop("length", 0.1)
return ax.quiver(*args, length=length, **kwargs)
else:
return ax.quiver(*args, **kwargs)
elif mode == "displacement":
Xdef = [X[i] + U[i] for i in range(gdim)]
import matplotlib.tri as tri
if gdim == 2 and tdim == 2:
# FIXME: Not tested
triang = tri.Triangulation(Xdef[0], Xdef[1], mesh.cells())
shading = kwargs.pop("shading", "flat")
return ax.tripcolor(triang, C, shading=shading, **kwargs)
else:
# Return gracefully to make regression test pass without vtk
warnings.warn(
'Plotting does not support displacement for {} in {}}. Continuing without plot.'
.format(tdim, gdim))
return
