Helmholtz equation on a mixed-topology mesh#
Download sources
This demo illustrates how to:
Solve a simple Helmholtz problem on a mixed-topology mesh.
Create an experimental mixed-topology mesh directly from NumPy arrays
In development
Mixed-topology meshes are a work in progress and are not yet fully supported in DOLFINx.
import sys
from mpi4py import MPI
import numpy as np
from scipy.sparse.linalg import spsolve
import basix
import basix._basixcpp
import dolfinx.cpp as _cpp
import ufl
from dolfinx.cpp.mesh import _create_mixed_mesh
from dolfinx.fem import (
FunctionSpace,
assemble_matrix,
assemble_vector,
coordinate_element,
create_dofmaps,
dirichletbc,
finiteelement,
locate_dofs_geometrical,
mixed_topology_form,
)
from dolfinx.graph import partitioner
from dolfinx.io.utils import cell_perm_vtk
from dolfinx.mesh import CellType, GhostMode, Mesh, Topology
if MPI.COMM_WORLD.size > 1:
print("Not yet running in parallel")
sys.exit(0)
Create a mixed-topology mesh#
nx = 16
ny = 16
nz = 16
n_cells = nx * ny * nz
cells: list = [[], []]
orig_idx: list = [[], []]
geom = []
if MPI.COMM_WORLD.rank == 0:
idx = 0
for i in range(n_cells):
iz = i // (nx * ny)
j = i % (nx * ny)
iy = j // nx
ix = j % nx
v0 = (iz * (ny + 1) + iy) * (nx + 1) + ix
v1 = v0 + 1
v2 = v0 + (nx + 1)
v3 = v1 + (nx + 1)
v4 = v0 + (nx + 1) * (ny + 1)
v5 = v1 + (nx + 1) * (ny + 1)
v6 = v2 + (nx + 1) * (ny + 1)
v7 = v3 + (nx + 1) * (ny + 1)
if ix < nx / 2:
cells[0] += [v0, v1, v2, v3, v4, v5, v6, v7]
orig_idx[0] += [idx]
idx += 1
else:
cells[1] += [v0, v1, v2, v4, v5, v6]
orig_idx[1] += [idx]
idx += 1
cells[1] += [v1, v2, v3, v5, v6, v7]
orig_idx[1] += [idx]
idx += 1
n_points = (nx + 1) * (ny + 1) * (nz + 1)
sqxy = (nx + 1) * (ny + 1)
for v in range(n_points):
iz = v // sqxy
p = v % sqxy
iy = p // (nx + 1)
ix = p % (nx + 1)
geom += [[ix / nx, iy / ny, iz / nz]]
cells_np = [np.array(c) for c in cells]
geomx = np.array(geom, dtype=np.float64)
hexahedron = coordinate_element(CellType.hexahedron, 1)
prism = coordinate_element(CellType.prism, 1)
hexahedron_cpp: _cpp.fem.CoordinateElement_float64 = (
hexahedron._cpp_object # type: ignore[assignment]
)
prism_cpp: _cpp.fem.CoordinateElement_float64 = prism._cpp_object # type: ignore[assignment]
part = partitioner()
mesh = _create_mixed_mesh(
MPI.COMM_WORLD,
cells_np,
[hexahedron_cpp, prism_cpp],
geomx,
part,
GhostMode.none,
2,
1,
None,
None,
)
Create a mixed-topology dofmap and function space#
Create elements and dofmaps for each cell type
elements = [
basix.create_element(basix.ElementFamily.P, basix.CellType.hexahedron, 1),
basix.create_element(basix.ElementFamily.P, basix.CellType.prism, 1),
]
dolfinx_elements = [
finiteelement(cell_type, basix.ufl.wrap_element(e), np.float64, mesh.geometry.dim)
for cell_type, e in zip([CellType.hexahedron, CellType.prism], elements, strict=True)
]
# NOTE: Both dofmaps have the same IndexMap, but different cell_dofs
dofmaps = create_dofmaps(
mesh.comm,
Topology(mesh.topology),
dolfinx_elements,
)
# Create C++ function space
V_cpp = _cpp.fem.FunctionSpace_float64(
mesh,
[e._cpp_object for e in dolfinx_elements],
[dofmap._cpp_object for dofmap in dofmaps],
)
# Select some BCs
def marker(x):
"""BC Selector."""
return np.logical_or(np.isclose(x[2], 0.0), np.isclose(x[2], 1.0))
# dirichletbc needs a function space that carries a UFL domain, to
# associate one with the (uniform) boundary value. UFL does not yet
# support mixed-topology domains (see the FIXME below), so wrap V_cpp
# with an arbitrarily chosen cell type's domain/element -- neither is
# used for anything beyond this association.
domain = ufl.Mesh(basix.ufl.element("Lagrange", "hexahedron", 1, shape=(3,)))
element = basix.ufl.wrap_element(elements[0])
V = FunctionSpace(Mesh(mesh, domain), element, V_cpp)
bcdofs = locate_dofs_geometrical(V, marker)
bc = dirichletbc(value=0.0, dofs=bcdofs, V=V)
Creating and compiling a variational formulation#
We create the variational forms for each cell type. FIXME: This hack is required at the moment because UFL does not yet know about mixed topology meshes.
a = []
L = []
for i, cell_name in enumerate(["hexahedron", "prism"]):
print(f"Creating form for {cell_name}")
element = basix.ufl.wrap_element(elements[i])
domain = ufl.Mesh(basix.ufl.element("Lagrange", cell_name, 1, shape=(3,)))
V = FunctionSpace(Mesh(mesh, domain), element, V_cpp)
u, v = ufl.TrialFunction(V), ufl.TestFunction(V)
k = 12.0
x = ufl.SpatialCoordinate(domain)
a += [(ufl.inner(ufl.grad(u), ufl.grad(v)) - k**2 * u * v) * ufl.dx]
f = ufl.sin(ufl.pi * x[0]) * ufl.sin(ufl.pi * x[1])
L += [f * v * ufl.dx]
Compile the form
FIXME: For the time being, since UFL doesn’t understand mixed topology
meshes, we have to call mixed_topology_form instead of form.
a_form = mixed_topology_form(a, dtype=np.float64)
L_form = mixed_topology_form(L, dtype=np.float64)
Assembling and solving the linear system#
We use the native matrix and
vector format in DOLFINx to assemble
the left and right hand side of the linear system.
A = assemble_matrix(a_form, bcs=[bc])
b = assemble_vector(L_form)
bc.set(b.array)
We use scipy.sparse.linalg.spsolve() to solve the
resulting linear system
A_scipy = A.to_scipy()
b_scipy = b.array
print(f"Solution vector norm {np.linalg.norm(x_scipy)}")
Mixed-topology I/O We manually build a ASCII XDMF file to store the mesh and solution NOTE: this should be replaced with VTKHDF
xdmf = """<?xml version="1.0"?>
<!DOCTYPE Xdmf SYSTEM "Xdmf.dtd" []>
<Xdmf Version="3.0" xmlns:xi="https://www.w3.org/2001/XInclude">
<Domain>
<Grid Name="mesh" GridType="Collection" CollectionType="spatial">
"""
perm = [cell_perm_vtk(CellType.hexahedron, 8), cell_perm_vtk(CellType.prism, 6)]
topologies = ["Hexahedron", "Wedge"]
for j in range(2):
vtk_topology = []
geom_dm = mesh.geometry.dofmaps[j]
for c in geom_dm:
vtk_topology += list(c[perm[j]])
topology_type = topologies[j]
xdmf += f"""
<Grid Name="{topology_type}" GridType="Uniform">
<Topology TopologyType="{topology_type}">
<DataItem Dimensions="{geom_dm.shape[0]} {geom_dm.shape[1]}"
Precision="4" NumberType="Int" Format="XML">
{" ".join(str(val) for val in vtk_topology)}
</DataItem>
</Topology>
<Geometry GeometryType="XYZ" NumberType="float" Rank="2" Precision="8">
<DataItem Dimensions="{mesh.geometry.x.shape[0]} 3" Format="XML">
{" ".join(str(val) for val in mesh.geometry.x.flatten())}
</DataItem>
</Geometry>
<Attribute Name="u" Center="Node" NumberType="float" Precision="8">
<DataItem Dimensions="{len(x_scipy)}" Format="XML">
{" ".join(str(val) for val in x_scipy)}
</DataItem>
</Attribute>
</Grid>"""
xdmf += """
</Grid>
</Domain>
</Xdmf>
"""
with open("mixed-mesh.xdmf", "w") as fd:
fd.write(xdmf)