Skip to content

Instantly share code, notes, and snippets.

@IgorBaratta
Created May 27, 2021 19:49
Show Gist options
  • Select an option

  • Save IgorBaratta/818fc0288203adfae6c3a93976dc5ede to your computer and use it in GitHub Desktop.

Select an option

Save IgorBaratta/818fc0288203adfae6c3a93976dc5ede to your computer and use it in GitHub Desktop.
from IPython import embed
from mpi4py import MPI
import dolfinx
import numpy
from pyvista.plotting.helpers import plot
import ufl
from dolfinx.io import XDMFFile
import dolfinx.plot
import pyvista
bc = 1
pyvista.rcParams["background"] = [bc, ]*3
pyvista.rcParams["font"]["color"] = [0, 0, 0]
x = numpy.array([[0, 0, 0],
[1, 0, 0],
[0, 1, 0],
[0, 0, 1]])
cells = numpy.array([[0, 1, 2, 3]])
ufl_mesh = ufl.Mesh(ufl.VectorElement("Lagrange", "tetrahedron", 1))
mesh = dolfinx.mesh.create_mesh(MPI.COMM_WORLD, cells, x, ufl_mesh)
W = dolfinx.FunctionSpace(mesh, ("Lagrange", 10))
point = W.element.interpolation_points()
degree = 2
element = ufl.FiniteElement("N1curl", "tetrahedron", degree)
V = dolfinx.FunctionSpace(mesh, element)
u = dolfinx.Function(V)
pos = V.dofmap.cell_dofs(0)
ndofs = V.dofmap.index_map.size_local
for dof in range(ndofs):
u.x.array[:] = 0
u.x.array[pos[dof]] = 1
c = dolfinx.fem.Expression(u, point)
vals = c.eval([0])
vals = vals/numpy.max(vals)
cloud = pyvista.PolyData(point)
cloud["Basis"] = vals.reshape(point.shape[0], 3)
glyphs = cloud.glyph("Basis", factor=.3)
plotter = pyvista.Plotter()
topology, cell_types = dolfinx.plot.create_vtk_topology(
mesh, mesh.topology.dim)
grid = pyvista.UnstructuredGrid(topology, cell_types, mesh.geometry.x)
actor = plotter.add_mesh(grid, style="wireframe", color="k", line_width=2)
actor2 = plotter.add_mesh(glyphs, show_scalar_bar=False)
# plotter.show()
B_fig = plotter.screenshot(f"basis_{degree}_{dof}.png")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment