Poisson equation

Download sources

• Python script

• Jupyter notebook

This demo illustrates how to:

• Solve a simple Helmholtz problem on a mixed-topology mesh.

• Create a mesh from numpy arrays using dolfinx.mesh.create_mesh()

In development

Mixed-topology meshes are a work in progress and are not yet fully supported in DOLFINx.

import sys
import typing

from mpi4py import MPI

import numpy as np
from scipy.sparse.linalg import spsolve

import basix
import dolfinx.cpp as _cpp
import ufl
from dolfinx.cpp.fem import locate_dofs_geometrical
from dolfinx.cpp.mesh import GhostMode, create_mesh
from dolfinx.fem import (
    FiniteElement,
    FunctionSpace,
    assemble_matrix,
    assemble_vector,
    coordinate_element,
    create_dofmaps,
    dirichletbc,
    mixed_topology_form,
)
from dolfinx.io.utils import cell_perm_vtk
from dolfinx.mesh import CellType, Mesh, Topology
from dolfinx.mesh import _create_cell_partitioner_from_ghost_mode as _cell_partitioner

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)

part = _cell_partitioner(GhostMode.none, 2)
mesh = create_mesh(
    MPI.COMM_WORLD,
    cells_np,
    [
        typing.cast(_cpp.fem.CoordinateElement_float64, hexahedron._cpp_object),
        typing.cast(_cpp.fem.CoordinateElement_float64, prism._cpp_object),
    ],
    geomx,
    part,
    2,
    1,
)

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(
        _cpp.fem.FiniteElement_float64(
            typing.cast(basix._basixcpp.FiniteElement_float64, e._e), None, False
        )
    )
    for e in elements
]
# 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],  # type: ignore[misc]
    [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_cpp, 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

x_scipy = spsolve(A_scipy, b_scipy)

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)