Solving PDEs with different scalar (float) types

This demo ( shows:

  • How to solve problems using different scalar types, .e.g. single or double precision, or complex numbers

  • Interfacing with SciPy sparse linear algebra functionality

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

import ufl
from dolfinx import fem, la, mesh, plot

from mpi4py import MPI

SciPy solvers do not support MPI, so all computations will be performed on a single MPI rank


Create a mesh and function space.

msh = mesh.create_rectangle(comm=comm, points=((0.0, 0.0), (2.0, 1.0)), n=(32, 16),
V = fem.FunctionSpace(msh, ("Lagrange", 1))

Define a variational problem.

u, v = ufl.TrialFunction(V), ufl.TestFunction(V)
x = ufl.SpatialCoordinate(msh)
fr = 10 * ufl.exp(-((x[0] - 0.5) ** 2 + (x[1] - 0.5) ** 2) / 0.02)
fc = ufl.sin(2 * np.pi * x[0]) + 10 * ufl.sin(4 * np.pi * x[1]) * 1j
gr = ufl.sin(5 * x[0])
gc = ufl.sin(5 * x[0]) * 1j
a = ufl.inner(ufl.grad(u), ufl.grad(v)) * ufl.dx
L = ufl.inner(fr + fc, v) * ufl.dx + ufl.inner(gr + gc, v) * ufl.ds

In preparation for constructing Dirichlet boundary conditions, locate facets on the constrained boundary and the corresponding degrees-of-freedom.

facets = mesh.locate_entities_boundary(msh, dim=1,
                                       marker=lambda x: np.logical_or(np.isclose(x[0], 0.0),
                                                                      np.isclose(x[0], 2.0)))
dofs = fem.locate_dofs_topological(V=V, entity_dim=1, entities=facets)

The below function computes the solution of the finite problem using a specified scalar type.

def solve(dtype=np.float32):
    """Solve the variational problem"""

    # Process forms. This will compile the forms for the requested type.
    a0 = fem.form(a, dtype=dtype)
    if np.issubdtype(dtype, np.complexfloating):
        L0 = fem.form(L, dtype=dtype)
        L0 = fem.form(ufl.replace(L, {fc: 0, gc: 0}), dtype=dtype)

    # Create a Dirichlet boundary condition
    bc = fem.dirichletbc(value=dtype(0), dofs=dofs, V=V)

    # Assemble forms
    A = fem.assemble_matrix(a0, [bc])
    b = fem.assemble_vector(L0)
    fem.apply_lifting(b.array, [a0], bcs=[[bc]])
    fem.set_bc(b.array, [bc])

    # Create a Scipy sparse matrix that shares data with A
    As = scipy.sparse.csr_matrix((, A.indices, A.indptr))

    # Solve the variational problem and return the solution
    uh = fem.Function(V, dtype=dtype)
    uh.x.array[:] = scipy.sparse.linalg.spsolve(As, b.array)
    return uh

This function visualises the solution.

def display(u, filter=np.real):
    """Plot the solution using pyvista"""
        import pyvista
        cells, types, x = plot.create_vtk_mesh(V)
        grid = pyvista.UnstructuredGrid(cells, types, x)
        grid.point_data["u"] = filter(u.x.array)

        plotter = pyvista.Plotter()
        plotter.add_mesh(grid, show_edges=True)
        plotter.add_title("real" if filter is np.real else "imag")
        if pyvista.OFF_SCREEN:
            plotter.screenshot(f"u_{'real' if filter is np.real else 'imag'}.png")
    except ModuleNotFoundError:
        print("'pyvista' is required to visualise the solution")

Solve the variational problem using different scalar types

uh = solve(dtype=np.float32)
uh = solve(dtype=np.float64)
uh = solve(dtype=np.complex64)
uh = solve(dtype=np.complex128)

Display the last computed solution

display(uh, np.real)
display(uh, np.imag)