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