nb.py 文件源码

def mplot_function(f, vmin, vmax, logscale):
    mesh = f.function_space().mesh()
    if (mesh.geometry().dim() != 2):
        raise AttributeError('Mesh must be 2D')
    # DG0 cellwise function
    if f.vector().size() == mesh.num_cells():
        C = f.vector().get_local()
        if logscale:
            return plt.tripcolor(mesh2triang(mesh), C, vmin=vmin, vmax=vmax, norm=cls.LogNorm() )
        else:
            return plt.tripcolor(mesh2triang(mesh), C, vmin=vmin, vmax=vmax)
    # Scalar function, interpolated to vertices
    elif f.value_rank() == 0:
        C = f.compute_vertex_values(mesh)
        if logscale:
            return plt.tripcolor(mesh2triang(mesh), C, vmin=vmin, vmax=vmax, norm=cls.LogNorm() )
        else:
            return plt.tripcolor(mesh2triang(mesh), C, shading='gouraud', vmin=vmin, vmax=vmax)
    # Vector function, interpolated to vertices
    elif f.value_rank() == 1:
        w0 = f.compute_vertex_values(mesh)
        if (len(w0) != 2*mesh.num_vertices()):
            raise AttributeError('Vector field must be 2D')
        X = mesh.coordinates()[:, 0]
        Y = mesh.coordinates()[:, 1]
        U = w0[:mesh.num_vertices()]
        V = w0[mesh.num_vertices():]
        C = np.sqrt(U*U+V*V)
        return plt.quiver(X,Y,U,V, C, units='x', headaxislength=7, headwidth=7, headlength=7, scale=4, pivot='middle')

# Plot a generic dolfin object (if supported)