Mixed Poisson equation

This demo illustrates how to solve Poisson equation using a mixed (two-field) formulation. In particular, it illustrates how to

• Create a mixed finite element problem.

• Extract subspaces.

• Apply boundary conditions to different fields in a mixed problem.

• Create integration domain data to execute finite element kernels. over subsets of the boundary.

• Use a submesh to represent boundary data

The full implementation is in demo_mixed_poisson/main.cpp.

Mixed formulation for the Poisson equation

Equation and problem definition

A mixed formulation of Poisson equation can be formulated by introducing an additional (vector) variable, namely the (negative) flux: \(\sigma = \nabla u\). The partial differential equations then read

\[\begin{split} \begin{align} \sigma - \nabla u &= 0 \quad {\rm in} \ \Omega, \\ \nabla \cdot \sigma &= - f \quad {\rm in} \ \Omega, \end{align} \end{split}\]

with boundary conditions

\[\begin{split} u = u_0 \quad {\rm on} \ \Gamma_{D}, \\ \sigma \cdot n = g \quad {\rm on} \ \Gamma_{N}. \end{split}\]

where \(n\) denotes the outward unit normal vector on the boundary. We see that the boundary condition for the flux (\(\sigma \cdot n = g\)) is an essential boundary condition (which should be enforced in the function space), while the other boundary condition (\(u = u_0\)) is a natural boundary condition (which should be applied to the variational form). Inserting the boundary conditions, this variational problem can be phrased in the general form: find \((\sigma, u) \in \Sigma_g \times V\) such that

\[ a((\sigma, u), (\tau, v)) = L((\tau, v)) \quad \forall \ (\tau, v) \in \Sigma_0 \times V, \]

where the forms \(a\) and \(L\) are defined as

\[\begin{split} \begin{align} a((\sigma, u), (\tau, v)) &:= \int_{\Omega} \sigma \cdot \tau + \nabla \cdot \tau \ u + \nabla \cdot \sigma \ v \ {\rm d} x, \\ L((\tau, v)) &:= - \int_{\Omega} f v \ {\rm d} x + \int_{\Gamma_D} u_0 \tau \cdot n \ {\rm d} s, \end{align} \end{split}\]

and \(\Sigma_g := \{ \tau \in H({\rm div})\) such that \(\tau \cdot n|_{\Gamma_N} = g \}\) and \(V := L^2(\Omega)\).

To discretize the above formulation, discrete function spaces \(\Sigma_h \subset \Sigma\) and \(V_h \subset V\) are needed to form a mixed function space \(\Sigma_h \times V_h\). A stable choice of finite element spaces is to let \(\Sigma_h\) be a Raviart-Thomas elements of polynomial order \(k\) and \(V_h\) be discontinuous elements of polynomial order \(k-1\).

We will use the same definitions of functions and boundaries as in the demo for the Poisson equation. These are:

• \(\Omega = [0,1] \times [0,1]\) (a unit square)

• \(\Gamma_{D} = \{(0, y) \cup (1, y) \in \partial \Omega\}\)

• \(\Gamma_{N} = \{(x, 0) \cup (x, 1) \in \partial \Omega\}\)

• \(u_0 = 20 y + 1\) on \(\Gamma_{D}\)

• \(g = 10\) (flux) on \(\Gamma_{N}\)

• $f = \sin(5x - 0.5) + 1 (source term)

UFL form file

The UFL file is implemented in demo_mixed_poisson/mixed_poisson.py.

UFL form implemented in python

The first step is to define the variational problem at hand. We define the variational problem in UFL terms in a separate form file demo_mixed_poisson/mixed_poisson.py. We begin by defining the finite element:

from basix import CellType
from basix.cell import sub_entity_type
from basix.ufl import element, mixed_element
from ufl import (
    Coefficient,
    FacetNormal,
    FunctionSpace,
    Measure,
    Mesh,
    TestFunctions,
    TrialFunctions,
    div,
    inner,
)

# Cell type for the mesh
msh_cell = CellType.triangle

# Weakly enforced boundary data will be represented using a function space
# defined over a submesh of the boundary. We get the submesh cell type from
# then mesh cell type.
submsh_cell = sub_entity_type(msh_cell, dim=1, index=0)

# Define finite elements for the problem
RT = element("RT", msh_cell, 1)
P = element("DP", msh_cell, 0)
ME = mixed_element([RT, P])

# Define UFL mesh and submesh
msh = Mesh(element("Lagrange", msh_cell, 1, shape=(2,)))
submsh = Mesh(element("Lagrange", submsh_cell, 1, shape=(2,)))

n = FacetNormal(msh)
V = FunctionSpace(msh, ME)

(sigma, u) = TrialFunctions(V)
(tau, v) = TestFunctions(V)

V0 = FunctionSpace(msh, P)
f = Coefficient(V0)

# We represent boundary data using a first-degree Lagrange space
# defined over a submesh of the boundary
Q = FunctionSpace(submsh, element("Lagrange", submsh_cell, 1))
u0 = Coefficient(Q)

# Specify the weak form of the problem
dx = Measure("dx", msh)
ds = Measure("ds", msh)
a = inner(sigma, tau) * dx + inner(u, div(tau)) * dx + inner(div(sigma), v) * dx
L = -inner(f, v) * dx + inner(u0 * n, tau) * ds(1)

#include "mixed_poisson.h"
#include <basix/cell.h>
#include <basix/finite-element.h>
#include <basix/mdspan.hpp>
#include <cmath>
#include <dolfinx.h>
#include <dolfinx/fem/Constant.h>
#include <dolfinx/fem/petsc.h>
#include <dolfinx/la/petsc.h>
#include <map>
#include <memory>
#include <petscmat.h>
#include <petscsys.h>
#include <petscsystypes.h>
#include <ranges>
#include <span>
#include <utility>
#include <vector>

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

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

  {
    mesh::CellType cell_type = mesh::CellType::triangle;

    // Create mesh
    auto mesh = std::make_shared<mesh::Mesh<U>>(mesh::create_rectangle<U>(
        MPI_COMM_WORLD, {{{0, 0}, {1, 1}}}, {32, 32}, cell_type));

    // Create Basix elements
    basix::cell::type basix_cell_type
        = dolfinx::mesh::cell_type_to_basix_type(cell_type);
    auto RT
        = basix::create_element<U>(basix::element::family::RT, basix_cell_type,
                                   1, basix::element::lagrange_variant::unset,
                                   basix::element::dpc_variant::unset, false);
    auto P0
        = basix::create_element<U>(basix::element::family::P, basix_cell_type,
                                   0, basix::element::lagrange_variant::unset,
                                   basix::element::dpc_variant::unset, true);

    // Create DOLFINx mixed element
    auto ME = std::make_shared<fem::FiniteElement<U>>(
        std::vector<fem::BasixElementData<U>>{{RT}, {P0}});

    // Create FunctionSpace
    auto V = std::make_shared<fem::FunctionSpace<U>>(
        fem::create_functionspace<U>(mesh, ME));

    // Get subspaces (views into V)
    auto V0 = std::make_shared<fem::FunctionSpace<U>>(V->sub({0}));
    auto V1 = std::make_shared<fem::FunctionSpace<U>>(V->sub({1}));

    // Collapse spaces
    auto W0 = std::make_shared<fem::FunctionSpace<U>>(V0->collapse().first);
    auto W1 = std::make_shared<fem::FunctionSpace<U>>(V1->collapse().first);

    // Create source function and interpolate $\sin(5x) + 1$
    auto f = std::make_shared<fem::Function<T>>(W1);
    f->interpolate(
        [](auto x) -> std::pair<std::vector<T>, std::vector<std::size_t>>
        {
          std::vector<T> f;
          for (std::size_t p = 0; p < x.extent(1); ++p)
          {
            auto x0 = x(0, p);
            f.push_back(std::sin(5 * x0) + 1);
          }
          return {f, {f.size()}};
        });

    // Create boundary condition for $\sigma$ and interpolate such that
    // flux = 10 (for top and bottom boundaries)
    auto g = std::make_shared<fem::Function<T>>(W0);
    g->interpolate(
        [](auto x) -> std::pair<std::vector<T>, std::vector<std::size_t>>
        {
          using mspan_t
              = md::mdspan<T, md::extents<std::size_t, 2, md::dynamic_extent>>;

          std::vector<T> fdata(2 * x.extent(1), 0);
          mspan_t f(fdata.data(), 2, x.extent(1));
          for (std::size_t p = 0; p < x.extent(1); ++p)
            f(1, p) = x(1, p) < 0.5 ? -10 : 10;
          return {std::move(fdata), {2, x.extent(1)}};
        });

    // Get list of all boundary facets
    mesh->topology()->create_connectivity(1, 2);
    std::vector bfacets = mesh::exterior_facet_indices(*mesh->topology());

    // Get facets with boundary condition on u
    std::vector<std::int32_t> dfacets = mesh::locate_entities_boundary(
        *mesh, 1,
        [](auto x)
        {
          using U = typename decltype(x)::value_type;
          constexpr U eps = 1e-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 - 1) < eps)
              marker[p] = true;
          }
          return marker;
        });

    // We'd like to represent `u_0` using a function space defined only
    // on the facets in `dfacets`. To do so, we begin by calling
    // `create_submesh` to get a `submesh` of those facets. It also
    // returns a map `submesh_to_mesh` whose `i`th entry is the facet in
    // mesh corresponding to cell `i` in submesh.
    int tdim = mesh->topology()->dim();
    int fdim = tdim - 1;
    std::shared_ptr<mesh::Mesh<U>> submesh;
    std::vector<std::int32_t> submesh_to_mesh;
    {
      auto [_submesh, _submesh_to_mesh, v_map, g_map]
          = mesh::create_submesh(*mesh, fdim, dfacets);
      submesh = std::make_shared<mesh::Mesh<U>>(std::move(_submesh));
      submesh_to_mesh = std::move(_submesh_to_mesh);
    }

    // Create an element for `u_0`
    basix::cell::type submesh_cell_type
        = dolfinx::mesh::cell_type_to_basix_type(
            submesh->topology()->cell_type());

    auto Qe = std::make_shared<fem::FiniteElement<U>>(
        basix::create_element<U>(basix::element::family::P, submesh_cell_type,
                                 1, basix::element::lagrange_variant::unset,
                                 basix::element::dpc_variant::unset, false));

    // Create a function space for `u_0` on the submesh
    auto Q = std::make_shared<fem::FunctionSpace<U>>(
        fem::create_functionspace<U>(submesh, Qe));

    // Boundary condition value for u and interpolate $20 y + 1$
    auto u0 = std::make_shared<fem::Function<T>>(Q);
    u0->interpolate(
        [](auto x) -> std::pair<std::vector<T>, std::vector<std::size_t>>
        {
          std::vector<T> f;
          for (std::size_t p = 0; p < x.extent(1); ++p)
            f.push_back(20 * x(1, p) + 1);
          return {f, {f.size()}};
        });

    // Write u0 to file to visualise
    io::VTKFile u0_file(MPI_COMM_WORLD, "u0.pvd", "w");
    u0_file.write<T>({*u0}, 0);

    // Compute facets with $\sigma$ (flux) boundary condition facets,
    // which is `{all boundary facet} - {u0 boundary facets}`
    std::vector<std::int32_t> nfacets;
    std::ranges::set_difference(bfacets, dfacets, std::back_inserter(nfacets));

    // Get dofs that are constrained by \sigma
    std::array<std::vector<std::int32_t>, 2> ndofs
        = fem::locate_dofs_topological(
            *mesh->topology(), {*V0->dofmap(), *W0->dofmap()}, 1, nfacets);

    // Create boundary condition for $\sigma. $\sigma \cdot n$ will be
    // constrained to to be equal to the normal component of $g$. The
    // boundary conditions are applied to degrees-of-freedom ndofs, and
    // `V0` is the subspace that is constrained.
    fem::DirichletBC<T> bc(g, ndofs, V0);

    // Create integration domain data for `u0` boundary condition
    // (applied on the `ds(1)` in the UFL file). First we get facet data
    // integration data for facets in dfacets.
    std::vector<std::int32_t> domains = fem::compute_integration_domains(
        fem::IntegralType::exterior_facet, *mesh->topology(), dfacets);

    // Create data structure for the `ds(1)` integration domain in form
    // (see the UFL file). It is for en exterior facet integral (the key
    // in the map), and exterior facet domain marked as '1' in the UFL
    // file, and `domains` holds the necessary data to perform
    // integration of selected facets.
    std::map<
        fem::IntegralType,
        std::vector<std::pair<std::int32_t, std::span<const std::int32_t>>>>
        subdomain_data{{fem::IntegralType::exterior_facet, {{1, domains}}}};

    // Since we are doing a `ds(1)` integral on mesh and `u0` is defined
    // on the `submesh`, we must provide an "entity map" relating cells
    // in `submesh` to entities in `mesh`. This is simply the "inverse"
    // of `submesh_to_mesh`:
    auto facet_imap = mesh->topology()->index_map(fdim);
    assert(facet_imap);
    std::size_t num_facets = mesh->topology()->index_map(fdim)->size_local()
                             + mesh->topology()->index_map(fdim)->num_ghosts();
    // Since not all facets in the mesh appear in the submesh, `submesh_to_mesh`
    // is only injective. We therefore map nonexistent facets to -1 when
    // creating the "inverse" map mesh_to_submesh
    std::vector<std::int32_t> mesh_to_submesh(num_facets, -1);
    for (std::size_t i = 0; i < submesh_to_mesh.size(); ++i)
      mesh_to_submesh[submesh_to_mesh[i]] = i;

    // Create the entity map to pass to `create_form`
    std::map<std::shared_ptr<const mesh::Mesh<U>>,
             std::span<const std::int32_t>>
        entity_maps = {{submesh, mesh_to_submesh}};

    // Define variational forms and attach he required data
    fem::Form<T> a = fem::create_form<T>(*form_mixed_poisson_a, {V, V}, {}, {},
                                         subdomain_data, {});
    // Since this form involves multiple domains (i.e. both `mesh` and `submesh`
    // for the boundary condition), we must pass the entity maps just created.
    // We must also tell the form which domain to integrate with respect to (in
    // this case `mesh`)
    fem::Form<T> L = fem::create_form<T>(
        *form_mixed_poisson_L, {V}, {{"f", f}, {"u0", u0}}, {}, subdomain_data,
        entity_maps, V->mesh());

    // Create solution finite element Function
    auto u = std::make_shared<fem::Function<T>>(V);

    // Create matrix and RHS vector data structures
    auto A = la::petsc::Matrix(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());

    // Assemble the bilinear form into a matrix. The PETSc matrix is
    // 'flushed' so we can set values in it in the subsequent step.
    MatZeroEntries(A.mat());
    fem::assemble_matrix(la::petsc::Matrix::set_fn(A.mat(), ADD_VALUES), a,
                         {bc});
    MatAssemblyBegin(A.mat(), MAT_FLUSH_ASSEMBLY);
    MatAssemblyEnd(A.mat(), MAT_FLUSH_ASSEMBLY);

    // Set '1' on diagonal for Dirichlet dofs
    fem::set_diagonal<T>(la::petsc::Matrix::set_fn(A.mat(), INSERT_VALUES), *V,
                         {bc});
    MatAssemblyBegin(A.mat(), MAT_FINAL_ASSEMBLY);
    MatAssemblyEnd(A.mat(), MAT_FINAL_ASSEMBLY);

    // Assemble the linear form `L` into RHS vector
    b.set(0);
    fem::assemble_vector(b.mutable_array(), L);

    // Modify unconstrained dofs on RHS to account for Dirichlet BC dofs
    // (constrained dofs), and perform parallel update on the vector.
    fem::apply_lifting<T, U>(b.mutable_array(), {a}, {{bc}}, {}, T(1));
    b.scatter_rev(std::plus<T>());

    // Set value for constrained dofs
    bc.set(b.mutable_array(), std::nullopt);

    // Create PETSc linear solver
    la::petsc::KrylovSolver lu(MPI_COMM_WORLD);
    la::petsc::options::set("ksp_type", "preonly");
    la::petsc::options::set("pc_type", "lu");
    if (sizeof(PetscInt) == 4)
      la::petsc::options::set("pc_factor_mat_solver_type", "mumps");
    else
      la::petsc::options::set("pc_factor_mat_solver_type", "superlu_dist");
    lu.set_from_options();

    // Solve linear system Ax = b
    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);
    lu.solve(_u.vec(), _b.vec());

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

    // Save solution in VTK format
    auto u_soln = std::make_shared<fem::Function<T>>(u->sub(1).collapse());
    io::VTKFile file(MPI_COMM_WORLD, "u.pvd", "w");
    file.write<T>({*u_soln}, 0);

#ifdef HAS_ADIOS2
    // Save solution in VTX format
    io::VTXWriter<U> vtx_u(MPI_COMM_WORLD, "u.bp", {u_soln}, "bp4");
    vtx_u.write(0);
    // Save interpolated boundary condition
    io::VTXWriter<U> vtx_u0(MPI_COMM_WORLD, "u0.bp", {u0}, "bp4");
    vtx_u0.write(0);
#endif
  }

  PetscFinalize();

  return 0;
}