Poisson equation

This demo illustrates how to:

• Solve a linear partial differential equation

• Create and apply Dirichlet boundary conditions

• Create a finite element function space

• Interpolate functions into a finite element space

Equation and problem definition

For a domain \(\Omega \subset \mathbb{R}^n\) with boundary \(\partial \Omega = \Gamma_{D} \cup \Gamma_{N}\), the Poisson problem reads

\[\begin{align*} - \nabla \cdot (\kappa \nabla u) &= f \quad {\rm in} \ \Omega, \\ u &= 0 \quad {\rm on} \ \Gamma_{D}, \\ \kappa \nabla u \cdot n &= g \quad {\rm on} \ \Gamma_{N}, \end{align*}\]

where \(f\) and \(g\) are input data, \(\kappa\) is a constant diffusion coefficient and \(n\) is the outward unit normal to the boundary. The variational problem reads: find \(u \in V\) such that

\[ a(u, v) = L(v) \quad \forall \ v \in V, \]

where \(V\) is a suitable function space of functions that vanish on \(\Gamma_{D}\), and

\[\begin{align*} a(u, v) &= \int_{\Omega} \kappa \nabla u \cdot \nabla v \, {\rm d} x, \\ L(v) &= \int_{\Omega} f v \, {\rm d} x + \int_{\Gamma_{N}} g v \, {\rm d} s. \end{align*}\]

This demo uses:

• \(\Omega = [0,2] \times [0,1]\)

• \(\Gamma_{D} = \{(x, y) \in \partial \Omega : x = 0 \ {\rm or} \ x = 2\}\)

• \(\Gamma_{N} = \partial \Omega \setminus \Gamma_{D}\)

• \(g = \sin(5x)\)

• \(f = 10\exp(-((x - 0.5)^2 + (y - 0.5)^2) / 0.02)\)

• \(\kappa = 2\)

Implementation

The variational forms are expressed in UFL in a Python file, which FFCx compiles into C code. The C++ program builds the mesh and function space, assembles and solves the linear system, and writes the solution to file.

Running this demo requires the files: demo_poisson/main.cpp, demo_poisson/poisson.py and demo_poisson/CMakeLists.txt.

UFL form file

The UFL forms are defined in demo_poisson/poisson.py.

UFL form implemented in Python

"""Poisson problem."""

We begin by defining the finite element:

from basix.ufl import element
from ufl import (
    Coefficient,
    Constant,
    FunctionSpace,
    Mesh,
    TestFunction,
    TrialFunction,
    ds,
    dx,
    grad,
    inner,
)

e = element("Lagrange", "triangle", 1)

The arguments to element are the finite element family, the cell type and the polynomial degree. The element e is therefore the continuous piecewise linear Lagrange element on triangles.

The mesh is described by a vector-valued linear Lagrange coordinate element. On the mesh we define the function space, the trial and test functions (\(u\) and \(v\)), the coefficients (\(f\) and \(g\)) and the constant \(\kappa\):

coord_element = element("Lagrange", "triangle", 1, shape=(2,))
mesh = Mesh(coord_element)

V = FunctionSpace(mesh, e)

u = TrialFunction(V)
v = TestFunction(V)
f = Coefficient(V)
g = Coefficient(V)
kappa = Constant(mesh)

Finally, we define the bilinear and linear forms according to the variational formulation of the equations:

a = kappa * inner(grad(u), grad(v)) * dx
L = inner(f, v) * dx + inner(g, v) * ds

C++ program

We include the DOLFINx headers and the header poisson.h generated by FFCx, which declares the compiled forms. The scalar type T is the PETSc scalar type, and U is the corresponding real type used for the mesh geometry.

#include "poisson.h"
#include <basix/finite-element.h>
#include <cmath>
#include <dolfinx.h>
#include <dolfinx/fem/Constant.h>
#include <dolfinx/fem/petsc.h>
#include <dolfinx/la/petsc.h>
#include <petscksp.h>
#include <petscmat.h>
#include <petscsys.h>
#include <petscsystypes.h>
#include <stdexcept>
#include <utility>
#include <vector>

using namespace dolfinx;
using T = PetscScalar;
using U = typename dolfinx::scalar_value_t<T>;

The main function performs the following steps:

1. Create a mesh of \(\Omega\) with dolfinx::mesh::create_rectangle(), using \(32 \times 16\) rectangles each divided into two triangles, and a continuous piecewise linear Lagrange space dolfinx::fem::FunctionSpace on it.

2. Create the forms \(a\) and \(L\) from the generated code, attaching the constant \(\kappa\) and the functions \(f\) and \(g\).

3. Locate the degrees-of-freedom on the facets of \(\Gamma_{D}\) and create a dolfinx::fem::DirichletBC that sets them to zero.

4. Interpolate \(f\) and \(g\) into \(V\) using dolfinx::fem::Function::interpolate().

5. Assemble the matrix and right-hand side vector, applying the Dirichlet condition, and solve the linear system with a direct (LU) solver.

6. Write the solution to file for visualisation in, e.g., ParaView.

int main(int argc, char* argv[])
{
  dolfinx::init_logging(argc, argv);
  common::petsc::check(PetscInitialize(&argc, &argv, nullptr, nullptr),
                       "PetscInitialize");

  {
    // Create mesh and function space
    auto mesh = std::make_shared<mesh::Mesh<U>>(mesh::create_rectangle<U>(
        MPI_COMM_WORLD, {{{0.0, 0.0}, {2.0, 1.0}}}, {32, 16},
        mesh::CellType::triangle, graph::partition_graph,
        mesh::DiagonalType::right, 2, mesh::GhostMode::shared_facet));

    auto element = basix::create_element<U>(
        basix::element::family::P, basix::cell::type::triangle, 1,
        basix::element::lagrange_variant::unset,
        basix::element::dpc_variant::unset, false);

    auto V
        = std::make_shared<fem::FunctionSpace<U>>(fem::create_functionspace<U>(
            mesh, std::make_shared<fem::FiniteElement<U>>(
                      element, mesh->geometry().dim())));

    // Create the coefficients of the forms
    auto kappa = std::make_shared<fem::Constant<T>>(2.0);
    auto f = std::make_shared<fem::Function<T>>(V);
    auto g = std::make_shared<fem::Function<T>>(V);

    // Define variational forms
    fem::Form<T> a = fem::create_form<T>(*form_poisson_a, {V, V}, {},
                                         {{"kappa", kappa}}, {}, {});
    fem::Form<T> L = fem::create_form<T>(*form_poisson_L, {V},
                                         {{"f", f}, {"g", g}}, {}, {}, {});

    // Define the boundary condition u = 0 on the facets with x = 0 or
    // x = 2
    std::vector facets = mesh::locate_entities_boundary(
        *mesh, 1,
        [](auto x)
        {
          using coord_t = typename decltype(x)::value_type;
          constexpr coord_t eps = 1.0e-8;
          std::vector<std::int8_t> marker(x.extent(1), false);
          for (std::size_t p = 0; p < x.extent(1); ++p)
          {
            auto x0 = x(0, p);
            if (std::abs(x0) < eps or std::abs(x0 - 2) < eps)
              marker[p] = true;
          }
          return marker;
        });
    std::vector bdofs = fem::locate_dofs_topological(
        *V->mesh()->topology_mutable(), *V->dofmap(), 1, facets);
    fem::DirichletBC<T> bc(0, bdofs, V);

    // Interpolate the source term f and boundary flux g
    f->interpolate(
        [](auto x) -> std::pair<std::vector<T>, std::vector<std::size_t>>
        {
          std::vector<T> f(x.extent(1));
          for (std::size_t p = 0; p < x.extent(1); ++p)
          {
            auto dx = (x(0, p) - 0.5) * (x(0, p) - 0.5);
            auto dy = (x(1, p) - 0.5) * (x(1, p) - 0.5);
            f[p] = 10 * std::exp(-(dx + dy) / 0.02);
          }

          return {std::move(f), std::vector<std::size_t>{f.size()}};
        });

    g->interpolate(
        [](auto x) -> std::pair<std::vector<T>, std::vector<std::size_t>>
        {
          std::vector<T> g(x.extent(1));
          for (std::size_t p = 0; p < x.extent(1); ++p)
            g[p] = std::sin(5 * x(0, p));
          return {std::move(g), std::vector<std::size_t>{g.size()}};
        });

    // Assemble the matrix. Rows and columns of constrained
    // degrees-of-freedom are zeroed and one is placed on the diagonal.
    auto u = std::make_shared<fem::Function<T>>(V);
    la::petsc::Matrix A(fem::petsc::create_matrix(a), false);
    la::Vector<T> b(L.function_spaces()[0]->dofmap()->index_map,
                    L.function_spaces()[0]->dofmap()->index_map_bs());

    common::petsc::check(MatZeroEntries(A.mat()), "MatZeroEntries");
    fem::assemble_matrix(la::petsc::Matrix::set_block_fn(A.mat(), ADD_VALUES),
                         a, {bc});
    fem::set_diagonal<T>(la::petsc::Matrix::set_fn(A.mat(), ADD_VALUES), *V,
                         {bc});
    common::petsc::check(MatAssemblyBegin(A.mat(), MAT_FINAL_ASSEMBLY),
                         "MatAssemblyBegin");
    common::petsc::check(MatAssemblyEnd(A.mat(), MAT_FINAL_ASSEMBLY),
                         "MatAssemblyEnd");

    // Assemble the right-hand side vector, modify it to account for the
    // Dirichlet condition, and set the constrained entries to the
    // boundary values
    std::ranges::fill(b.array(), 0);
    fem::assemble_vector(b.array(), L);
    fem::apply_lifting(b.array(), {a}, {{bc}}, {}, T(1));
    b.scatter_rev(std::plus<T>());
    bc.set(b.array(), std::nullopt);

    // Solve the linear system with a direct LU solver
    la::petsc::KrylovSolver lu(MPI_COMM_WORLD);
    common::petsc::set_option("ksp_type", "preonly");
    common::petsc::set_option("pc_type", "lu");
    lu.set_from_options();

    lu.set_operator(A.mat());
    la::petsc::Vector _u(la::petsc::create_vector_wrap(*u->x()), false);
    la::petsc::Vector _b(la::petsc::create_vector_wrap(b), false);
    if (lu.solve(_u.vec(), _b.vec()) < 0)
      throw std::runtime_error("Linear solver did not converge.");

    // The KSP object is available for anything the solver does not
    // wrap, here the number of linear solver iterations
    PetscInt num_it = 0;
    common::petsc::check(KSPGetIterationNumber(lu.ksp(), &num_it),
                         "KSPGetIterationNumber");
    std::cout << "Number of linear solver iterations: " << num_it << std::endl;

    // Update ghost values before output
    u->x()->scatter_fwd();

    // Save solution in VTK format
    io::VTKFile file(MPI_COMM_WORLD, "u.pvd", "w");
    file.write<T>({*u}, 0);

#ifdef HAS_ADIOS2
    // Save solution in VTX format
    io::VTXWriter<U> vtx(MPI_COMM_WORLD, "u.bp", {u}, "bp4");
    vtx.write(0);
#endif
  }

  common::petsc::check(PetscFinalize(), "PetscFinalize");
  return 0;
}