def surface(self,
f,
axis=None,
mesh=None,
derivative=(0, ),
shading=True,
grid=False,
resolution=(100, 100),
edge_resolution=10,
flag=None):
"""
Plot the surface of a function defined on the finite element mesh
Inputs:
ax: axis (don't forget to initialize it using projection='3d')
f: Function, function to be plotted
mesh: Mesh, on which to plot the function
*derivatives [(0,)]: int, tuple specifying what derivatives to
plot (see Function.eval for details).
*shading [True]: bool, shade surface or use wire plot?
*grid [False]: bool, display grid?
*resolution [(100,100)]: int, tuple (nx,ny) number of points
in the x- and y directions.
*edge_resolution: int, number of points along each each edge
*flag [None]: str/int marker for submesh TODO: Not implemented
Output:
ax: Axis, containing plot.
"""
#
# Check if input is a Function object
#
assert isinstance(f, Map), 'Can only plot Map objects.'
if mesh is None:
if f.mesh is not None:
mesh = f.mesh
else:
mesh_error = 'Mesh must be specified, either explicitly, '+\
'or as part of the Function.'
raise Exception(mesh_error)
x0, x1, y0, y1 = mesh.bounding_box()
system = Assembler()
if shading:
#
# Colormap
#
# Define Grid
nx, ny = resolution
x, y = np.linspace(x0, x1, nx), np.linspace(y0, y1, ny)
xx, yy = np.meshgrid(x, y)
xy = np.array([xx.ravel(), yy.ravel()]).transpose()
# Evaluate function
zz = f.eval(xy, derivative=derivative)
z_min, z_max = zz.min(), zz.max()
if grid:
alpha = 0.5
else:
alpha = 1
axis.plot_surface(xx,yy,zz.reshape(xx.shape),cmap='viridis', \
linewidth=1, antialiased=True, alpha=alpha)
self.exit(axis=axis)
if grid:
#
# Wirefunction
#
ne = edge_resolution
lines = []
node_count = 0
initialize_min_max = True
for node in mesh.root_node().get_leaves():
#
# Function type
#
if callable(f):
#
# Explicit function
#
assert derivative==(0,),\
'Discretize before plotting derivatives.'
f_loc = f
elif isinstance(f, Map):
#
# Function object
#
f_loc = f
elif len(f) == system.n_dofs():
#
# Nodal function
#
f_loc = f[system.get_global_dofs(node)]
elif len(f) == mesh.n_nodes():
#
# Mesh function
#
f_loc = f[node_count]
cell = node.cell()
for edge in cell.get_edges():
# Points on edges
v = edge.vertex_coordinates()
x0, y0 = v[0]
x1, y1 = v[1]
t = np.linspace(0, 1, ne)
xx = (1 - t) * x0 + t * x1
yy = (1 - t) * y0 + t * y1
# Evaluate function at edge points
zz = system.f_eval_loc(f_loc, node, x=np.array([xx,yy]).T, \
derivatives=derivative)
if initialize_min_max:
z_min = zz.min()
z_max = zz.max()
initialize_min_max = False
else:
z_max = max(zz.max(), z_max)
z_min = min(zz.min(), z_min)
for i in range(ne - 1):
lines.append([(xx[i], yy[i], zz[i]),
(xx[i + 1], yy[i + 1], zz[i + 1])])
node_count += 1
axis.add_collection(
Line3DCollection(lines, colors='k', linewidth=0.5))
x0, x1, y0, y1 = mesh.box()
hx = x1 - x0
hy = y1 - y0
hz = z_max - z_min
spc = 0.1
#print(z_min,z_max)
axis.set_xlim(x0 - spc * hx, x1 + spc * hx)
axis.set_ylim(y0 - spc * hy, y1 + spc * hy)
axis.set_zlim(z_min - spc * hz, z_max + spc * hz)
return axis
