Prototype real space

README.md

Poisson with Real space

Build demo

    cmake -G Ninja -B build-dir -DCMAKE_BUILD_TYPE=Developer
    ninja -j8 -C build-dir

Run

./build-dir/demo_real_space N

where N is number of elements in each direction

CMakeLists.txt

# This file was generated by running
#
# python cmake/scripts/generate-cmakefiles.py from dolfinx/cpp
#
cmake_minimum_required(VERSION 3.19)

set(PROJECT_NAME demo_real_space)
project(${PROJECT_NAME} LANGUAGES C CXX)

# Set C++20 standard
set(CMAKE_CXX_STANDARD 20)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
set(CMAKE_CXX_EXTENSIONS OFF)

if(NOT TARGET dolfinx)
  find_package(DOLFINX REQUIRED)
endif()

include(CheckSymbolExists)
set(CMAKE_REQUIRED_INCLUDES ${PETSC_INCLUDE_DIRS})
check_symbol_exists(PETSC_USE_COMPLEX petscsystypes.h PETSC_SCALAR_COMPLEX)
check_symbol_exists(PETSC_USE_REAL_DOUBLE petscsystypes.h PETSC_REAL_DOUBLE)

# Add target to compile UFL files
if(PETSC_SCALAR_COMPLEX EQUAL 1)
  if(PETSC_REAL_DOUBLE EQUAL 1)
    set(SCALAR_TYPE "--scalar_type=complex128")
  else()
    set(SCALAR_TYPE "--scalar_type=complex64")
  endif()
else()
  if(PETSC_REAL_DOUBLE EQUAL 1)
    set(SCALAR_TYPE "--scalar_type=float64")
  else()
    set(SCALAR_TYPE "--scalar_type=float32")
  endif()
endif()
add_custom_command(
  OUTPUT poisson.c
  COMMAND ffcx ${CMAKE_CURRENT_SOURCE_DIR}/poisson.py ${SCALAR_TYPE}
  VERBATIM
  DEPENDS poisson.py
  COMMENT "Compile poisson.py using FFCx"
)


set(CMAKE_INCLUDE_CURRENT_DIR ON)

add_executable(${PROJECT_NAME} main.cpp ${CMAKE_CURRENT_BINARY_DIR}/poisson.c)
target_link_libraries(${PROJECT_NAME} dolfinx)

# Do not throw error for 'multi-line comments' (these are typical in rst which
# includes LaTeX)
include(CheckCXXCompilerFlag)
check_cxx_compiler_flag("-Wno-comment" HAVE_NO_MULTLINE)
set_source_files_properties(
  main.cpp
  PROPERTIES
    COMPILE_FLAGS
    "$<$<BOOL:${HAVE_NO_MULTLINE}>:-Wno-comment -Wall -Wextra -pedantic -Werror>"
)

# Test targets (used by DOLFINx testing system)
set(TEST_PARAMETERS2 -np 2 ${MPIEXEC_PARAMS} "./${PROJECT_NAME}")
set(TEST_PARAMETERS3 -np 3 ${MPIEXEC_PARAMS} "./${PROJECT_NAME}")
add_test(NAME ${PROJECT_NAME}_mpi_2 COMMAND "mpirun" ${TEST_PARAMETERS2})
add_test(NAME ${PROJECT_NAME}_mpi_3 COMMAND "mpirun" ${TEST_PARAMETERS3})
add_test(NAME ${PROJECT_NAME}_serial COMMAND ${PROJECT_NAME})

main.cpp

// Hand coded Real space
// Author: Jørgen S. Dokken
// SPDX Licence: MIT

#include <basix/finite-element.h>
#include <cmath>
#include <concepts>
#include <cstdlib>
#include <dolfinx/common/IndexMap.h>
#include <dolfinx/common/log.h>
#include <dolfinx/fem/DofMap.h>
#include <dolfinx/fem/ElementDofLayout.h>
#include <dolfinx/fem/FiniteElement.h>
#include <dolfinx/fem/Form.h>
#include <dolfinx/fem/FunctionSpace.h>
#include <dolfinx/fem/petsc.h>
#include <dolfinx/fem/utils.h>
#include <dolfinx/io/ADIOS2Writers.h>
#include <dolfinx/la/MatrixCSR.h>
#include <dolfinx/la/petsc.h>
#include <dolfinx/mesh/Mesh.h>
#include <dolfinx/mesh/cell_types.h>
#include <dolfinx/mesh/generation.h>
#include <dolfinx/mesh/utils.h>
#include <filesystem>
#include <mpi.h>
#include <numbers>
#include <poisson.h>

using namespace dolfinx;

namespace
{
    template <typename T>
    dolfinx::fem::FunctionSpace<T>
    create_real_functionspace(std::shared_ptr<const dolfinx::mesh::Mesh<T>> mesh,
                              std::vector<std::size_t> value_shape)
    {

        basix::FiniteElement e_v = basix::create_element<T>(
            basix::element::family::P,
            mesh::cell_type_to_basix_type(mesh->topology()->cell_type()), 0,
            basix::element::lagrange_variant::unset,
            basix::element::dpc_variant::unset, true);

        // NOTE: Optimize input source/dest later as we know this a priori
        std::int32_t num_dofs = (dolfinx::MPI::rank(MPI_COMM_WORLD) == 0) ? 1 : 0;
        std::int32_t num_ghosts = (dolfinx::MPI::rank(MPI_COMM_WORLD) != 0) ? 1 : 0;
        std::vector<std::int64_t> ghosts(num_ghosts, 0);
        ghosts.reserve(1);
        std::vector<int> owners(num_ghosts, 0);
        owners.reserve(1);
        std::shared_ptr<const dolfinx::common::IndexMap> imap = std::make_shared<const dolfinx::common::IndexMap>(
            MPI_COMM_WORLD, num_dofs, ghosts, owners);
        std::size_t value_size = std::accumulate(
            value_shape.begin(), value_shape.end(), 1, std::multiplies{});
        int index_map_bs = value_size;
        int bs = value_size;
        // Element dof layout
        fem::ElementDofLayout dof_layout(value_size, e_v.entity_dofs(),
                                         e_v.entity_closure_dofs(), {}, {});
        std::size_t num_cells_on_process = mesh->topology()->index_map(mesh->topology()->dim())->size_local() + mesh->topology()->index_map(mesh->topology()->dim())->num_ghosts();

        std::vector<std::int32_t> dofmap(num_cells_on_process, 0);
        dofmap.reserve(1);
        std::shared_ptr<const dolfinx::fem::DofMap> real_dofmap = std::make_shared<const dolfinx::fem::DofMap>(dof_layout, imap,
                                                                                                               index_map_bs, dofmap, bs);

        std::shared_ptr<const dolfinx::fem::FiniteElement<T>> d_el = std::make_shared<const dolfinx::fem::FiniteElement<T>>(e_v, value_size,
                                                                                                                            false);

        return dolfinx::fem::FunctionSpace<T>(mesh, d_el, real_dofmap, value_shape);
    }
}
/// @brief This program shows how to create finite element spaces without FFCx
/// generated code.
int main(int argc, char *argv[])
{
    dolfinx::init_logging(argc, argv);
    PetscInitialize(&argc, &argv, nullptr, nullptr);

    // The main body of the function is scoped to ensure that all objects
    // that depend on an MPI communicator are destroyed before MPI is
    // finalised at the end of this function.
    {
        if (argc < 2)
        {
            std::cerr << "Usage: ./demo_real_space N" << std::endl;
            return 1;
        }

        std::int64_t N = atoi(argv[1]);

        // Create mesh using PetscScalar for geometry coordinates
        auto mesh0 = std::make_shared<mesh::Mesh<PetscReal>>(mesh::create_rectangle<PetscReal>(
            MPI_COMM_WORLD, {{{0.0, 0.0}, {1, 1}}}, {N, N},
            mesh::CellType::triangle,
            mesh::create_cell_partitioner(mesh::GhostMode::none)));
        std::shared_ptr<const dolfinx::fem::FunctionSpace<PetscReal>> Q = std::make_shared<const dolfinx::fem::FunctionSpace<PetscReal>>(
            create_real_functionspace<PetscReal>(mesh0, {1}));

        auto e_u = basix::create_element<PetscReal>(
            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<const fem::FunctionSpace<PetscReal>>(
            fem::create_functionspace(mesh0, e_u));

        // Define variational forms
        auto a00 = std::make_shared<fem::Form<PetscScalar>>(
            fem::create_form<PetscScalar>(*form_poisson_a00, {V, V}, {}, {{}}, {}, {}));
        auto a01 = std::make_shared<fem::Form<PetscScalar>>(
            fem::create_form<PetscScalar>(*form_poisson_a01, {V, Q}, {}, {{}}, {}, {}));
        auto a10 = std::make_shared<fem::Form<PetscScalar>>(
            fem::create_form<PetscScalar>(*form_poisson_a10, {Q, V}, {}, {{}}, {}, {}));

        std::vector<std::vector<const dolfinx::fem::Form<PetscScalar, PetscReal> *>> as;
        std::vector<const dolfinx::fem::Form<PetscScalar, PetscReal> *> a0 = {a00.get(), a01.get()};
        as.push_back(a0);
        std::vector<const dolfinx::fem::Form<PetscScalar, PetscReal> *> a1 = {a10.get(), nullptr};
        as.push_back(a1);
        Mat A = dolfinx::fem::petsc::create_matrix_block(as);
        std::vector<std::pair<std::reference_wrapper<const common::IndexMap>, int>>
            maps;
        maps.push_back({*V->dofmap()->index_map, V->dofmap()->index_map_bs()});
        maps.push_back({*Q->dofmap()->index_map, Q->dofmap()->index_map_bs()});
        std::vector<IS> index_sets = dolfinx::la::petsc::create_index_sets(maps);
        MatZeroEntries(A);
        for (std::size_t i = 0; i < as.size(); ++i)
        {
            for (std::size_t j = 0; j < as[i].size(); ++j)
            {
                if (as[i][j])
                {
                    Mat Asub;
                    MatGetLocalSubMatrix(A, index_sets[i], index_sets[j], &Asub);
                    fem::assemble_matrix(
                        la::petsc::Matrix::set_block_fn(Asub, ADD_VALUES), *as[i][j], {});
                    MatRestoreLocalSubMatrix(A, index_sets[i], index_sets[j], &Asub);
                    MatDestroy(&Asub);
                }
            }
        }

        MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY);
        MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY);

        // Create vectors
        Vec b = dolfinx::fem::petsc::create_vector_block(maps);
        VecZeroEntries(b);
        std::vector<std::vector<PetscScalar>> local_vecs = dolfinx::la::petsc::get_local_vectors(b, maps);
        auto L = std::make_shared<fem::Form<PetscScalar, PetscReal>>(
            fem::create_form<PetscScalar, PetscReal>(*form_poisson_L, {V}, {}, {{}}, {}, {}));
        fem::assemble_vector(std::span(local_vecs[0].data(), local_vecs[0].size()),
                             *L);

        std::vector<std::span<const PetscScalar>> lv;
        lv.reserve(2);
        for (auto &l : local_vecs)
            lv.push_back(std::span<const PetscScalar>(l.data(), l.size()));
        dolfinx::la::petsc::scatter_local_vectors(b, lv, maps);
        VecGhostUpdateBegin(b, ADD_VALUES, SCATTER_REVERSE);
        VecGhostUpdateEnd(b, ADD_VALUES, SCATTER_REVERSE);

        la::petsc::KrylovSolver lu(MPI_COMM_WORLD);
        la::petsc::options::set("ksp_type", "preonly");
        la::petsc::options::set("pc_type", "lu");
        la::petsc::options::set("pc_factor_mat_solver_type", "mumps");
        lu.set_from_options();

        lu.set_operator(A);
        Vec _u = dolfinx::fem::petsc::create_vector_block(maps);

        lu.solve(_u, b);

        std::vector<std::vector<PetscScalar>> local_out = dolfinx::la::petsc::get_local_vectors(_u, maps);
        std::shared_ptr<dolfinx::fem::Function<PetscScalar>> u = std::make_shared<dolfinx::fem::Function<PetscScalar>>(V);
        std::span<PetscScalar> u_arr = u->x()->mutable_array();
        std::copy(local_out[0].begin(), local_out[0].end(), u_arr.begin());
        u->x()->scatter_fwd();

        auto J = std::make_shared<fem::Form<PetscScalar>>(fem::create_form<PetscScalar>(
            *form_poisson_J, {}, {{"uh", u}}, {{}}, {}, {}, V->mesh()));

        PetscScalar error_local = fem::assemble_scalar(*J);
        PetscScalar error_global;

        MPI_Allreduce(&error_local, &error_global, 1, dolfinx::MPI::mpi_type<PetscScalar>(), MPI_SUM,
                      MPI_COMM_WORLD);
        if (dolfinx::MPI::rank(MPI_COMM_WORLD) == 0)
            std::cout << "N" << N << " Error: " << std::sqrt(error_global)
                      << std::endl;

        io::VTXWriter<double> file(MPI_COMM_WORLD, "u.bp", {u}, "BP5");
        file.write(0.0);
        file.close();
        VecDestroy(&b);
        MatDestroy(&A);
    }
    return 0;

    PetscFinalize();
}

poisson.py

from basix.ufl import element
from ufl import (
    Coefficient,
    FacetNormal,
    FunctionSpace,
    Mesh,
    SpatialCoordinate,
    TestFunction,
    TrialFunction,
    cos,
    div,
    dot,
    ds,
    Measure,
    grad,
    inner,
    pi,
    sin,
)

e_u = element("Lagrange", "triangle", 1, discontinuous=False)
e_v = element("Lagrange", "triangle", 0, discontinuous=True)

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

V = FunctionSpace(mesh, e_u)
Q = FunctionSpace(mesh, e_v)

u = TrialFunction(V)
v = TestFunction(V)
c = TrialFunction(Q)
d = TestFunction(Q)

x = SpatialCoordinate(mesh)
u_ex = sin(x[0] * 2 * pi) * cos(x[1] * 4 * pi)
dx = Measure("dx", domain=mesh)
f = -div(grad(u_ex))
n = FacetNormal(mesh)
g = dot(grad(u_ex), n)
a00 = inner(grad(u), grad(v)) * dx
a01 = inner(c, v) * dx
a10 = inner(u, d) * dx

L = inner(f, v) * dx + inner(g, v) * ds

uh = Coefficient(V)
J = inner(u_ex - uh, u_ex - uh) * dx

forms = [a00, a01, a10, L, J]

Build with:

    cmake -G Ninja -B build-dir -DCMAKE_BUILD_TYPE=Developer 
    ninja -j8 -C build-dir

Run with

./build-dir/demo_real_space