utils.h

Go to the documentation of this file.

// Copyright (C) 2013-2025 Johan Hake, Jan Blechta, Garth N. Wells and Paul T.
// Kühner
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include "Constant.h"
#include "CoordinateElement.h"
#include "DofMap.h"
#include "ElementDofLayout.h"
#include "Expression.h"
#include "Form.h"
#include "Function.h"
#include "FunctionSpace.h"
#include "kernel.h"
#include "sparsitybuild.h"
#include <algorithm>
#include <array>
#include <concepts>
#include <cstddef>
#include <dolfinx/common/defines.h>
#include <dolfinx/common/types.h>
#include <dolfinx/la/SparsityPattern.h>
#include <dolfinx/mesh/EntityMap.h>
#include <dolfinx/mesh/Topology.h>
#include <dolfinx/mesh/cell_types.h>
#include <dolfinx/mesh/utils.h>
#include <format>
#include <functional>
#include <memory>
#include <numeric>
#include <optional>
#include <ranges>
#include <set>
#include <span>
#include <stdexcept>
#include <string>
#include <tuple>
#include <type_traits>
#include <ufcx.h>
#include <utility>
#include <vector>

namespace basix
{
template <std::floating_point T>
class FiniteElement;
}
 
namespace dolfinx::common
{
class IndexMap;
}
 
namespace dolfinx::fem
{
template <dolfinx::scalar T, std::floating_point U>
class Expression;
 
namespace impl
{
template <int num_cells>
std::array<std::int32_t, 2 * num_cells>
get_cell_facet_pairs(std::int32_t f, std::span<const std::int32_t> cells,
                     const graph::AdjacencyList<std::int32_t>& c_to_f)
{
  // Loop over cells sharing facet
  assert(cells.size() == num_cells);
  std::array<std::int32_t, 2 * num_cells> cell_local_facet_pairs;
  for (int c = 0; c < num_cells; ++c)
  {
    // Get local index of facet with respect to the cell
    std::int32_t cell = cells[c];
    auto cell_facets = c_to_f.links(cell);
    auto facet_it = std::find(cell_facets.begin(), cell_facets.end(), f);
    assert(facet_it != cell_facets.end());
    int local_f = std::distance(cell_facets.begin(), facet_it);
    cell_local_facet_pairs[2 * c] = cell;
    cell_local_facet_pairs[2 * c + 1] = local_f;
  }
 
  return cell_local_facet_pairs;
}

template <int num_cells>
std::array<std::int32_t, 2 * num_cells>
get_cell_entity_pairs(std::int32_t e, std::span<const std::int32_t> cells,
                      const graph::AdjacencyList<std::int32_t>& c_to_e)
{
  static_assert(num_cells == 1); // Patch assembly not supported.
 
  assert(cells.size() > 0);
 
  // Use first cell for assembly over by default
  std::int32_t cell = cells[0];
 
  // Find local index of entity within cell
  auto cell_entities = c_to_e.links(cell);
  auto it = std::ranges::find(cell_entities, e);
  assert(it != cell_entities.end());
  std::int32_t local_index = std::distance(cell_entities.begin(), it);
 
  return {cell, local_index};
}
 
} // namespace impl

std::vector<std::int32_t>
compute_integration_domains(IntegralType integral_type,
                            const mesh::Topology& topology,
                            std::span<const std::int32_t> entities);

template <dolfinx::scalar T, std::floating_point U>
std::vector<std::vector<std::array<std::shared_ptr<const FunctionSpace<U>>, 2>>>
extract_function_spaces(const std::vector<std::vector<const Form<T, U>*>>& a)
{
  std::vector<
      std::vector<std::array<std::shared_ptr<const FunctionSpace<U>>, 2>>>
      spaces(
          a.size(),
          std::vector<std::array<std::shared_ptr<const FunctionSpace<U>>, 2>>(
              a.front().size()));
  for (std::size_t i = 0; i < a.size(); ++i)
  {
    for (std::size_t j = 0; j < a[i].size(); ++j)
    {
      if (const Form<T, U>* form = a[i][j]; form)
        spaces[i][j] = {form->function_spaces()[0], form->function_spaces()[1]};
    }
  }
  return spaces;
}

template <dolfinx::scalar T, std::floating_point U>
la::SparsityPattern create_sparsity_pattern(const Form<T, U>& a)
{
  std::shared_ptr mesh = a.mesh();
  assert(mesh);
 
  // Get index maps and block sizes from the DOF maps. Note that in
  // mixed-topology meshes, despite there being multiple DOF maps, the
  // index maps and block sizes are the same.
  std::array<std::reference_wrapper<const DofMap>, 2> dofmaps{
      *a.function_spaces().at(0)->dofmaps(0),
      *a.function_spaces().at(1)->dofmaps(0)};
 
  const std::array index_maps{dofmaps[0].get().index_map,
                              dofmaps[1].get().index_map};
  const std::array bs
      = {dofmaps[0].get().index_map_bs(), dofmaps[1].get().index_map_bs()};
 
  la::SparsityPattern pattern(mesh->comm(), index_maps, bs);
  build_sparsity_pattern(pattern, a);
  return pattern;
}

template <dolfinx::scalar T, std::floating_point U>
void build_sparsity_pattern(la::SparsityPattern& pattern, const Form<T, U>& a)
{
  if (a.rank() != 2)
  {
    throw std::runtime_error(
        "Cannot create sparsity pattern. Form is not a bilinear.");
  }
 
  std::shared_ptr mesh = a.mesh();
  assert(mesh);
  std::shared_ptr mesh0 = a.function_spaces().at(0)->mesh();
  assert(mesh0);
  std::shared_ptr mesh1 = a.function_spaces().at(1)->mesh();
  assert(mesh1);
 
  const std::set<IntegralType> types = a.integral_types();
  if (types.find(IntegralType::interior_facet) != types.end()
      or types.find(IntegralType::exterior_facet) != types.end())
  {
    // FIXME: cleanup these calls? Some of the happen internally again.
    int tdim = mesh->topology()->dim();
    mesh->topology_mutable()->create_entities(tdim - 1);
    mesh->topology_mutable()->create_connectivity(tdim - 1, tdim);
  }
 
  common::Timer t0("Build sparsity");
 
  auto extract_cells = [](std::span<const std::int32_t> facets)
  {
    assert(facets.size() % 2 == 0);
    std::vector<std::int32_t> cells;
    cells.reserve(facets.size() / 2);
    for (std::size_t i = 0; i < facets.size(); i += 2)
      cells.push_back(facets[i]);
    return cells;
  };
 
  const int num_cell_types = mesh->topology()->cell_types().size();
  for (int cell_type_idx = 0; cell_type_idx < num_cell_types; ++cell_type_idx)
  {
    std::array<std::reference_wrapper<const DofMap>, 2> dofmaps{
        *a.function_spaces().at(0)->dofmaps(cell_type_idx),
        *a.function_spaces().at(1)->dofmaps(cell_type_idx)};
 
    // Create and build sparsity pattern
    for (auto type : types)
    {
      switch (type)
      {
      case IntegralType::cell:
        for (int i = 0; i < a.num_integrals(type, cell_type_idx); ++i)
        {
          sparsitybuild::cells(
              pattern,
              std::pair{a.domain_arg(type, 0, i, cell_type_idx),
                        a.domain_arg(type, 1, i, cell_type_idx)},
              {{dofmaps[0], dofmaps[1]}});
        }
        break;
      case IntegralType::interior_facet:
        for (int i = 0; i < a.num_integrals(type, cell_type_idx); ++i)
        {
          sparsitybuild::interior_facets(
              pattern,
              {extract_cells(a.domain_arg(type, 0, i, 0)),
               extract_cells(a.domain_arg(type, 1, i, 0))},
              {{dofmaps[0], dofmaps[1]}});
        }
        break;
      case IntegralType::exterior_facet:
      case IntegralType::ridge:
      case IntegralType::vertex:
        for (int i = 0; i < a.num_integrals(type, cell_type_idx); ++i)
        {
          sparsitybuild::cells(
              pattern,
              std::pair{extract_cells(a.domain_arg(type, 0, i, 0)),
                        extract_cells(a.domain_arg(type, 1, i, 0))},
              {{dofmaps[0], dofmaps[1]}});
        }
        break;
      default:
        throw std::runtime_error("Unsupported integral type");
      }
    }
  }
 
  t0.stop();
}

template <std::floating_point T>
ElementDofLayout create_element_dof_layout(const fem::FiniteElement<T>& element,
                                           const std::vector<int>& parent_map
                                           = {})
{
  // Create subdofmaps and compute offset
  std::vector<int> offsets(1, 0);
  std::vector<dolfinx::fem::ElementDofLayout> sub_doflayout;
  int bs = element.block_size();
  for (int i = 0; i < element.num_sub_elements(); ++i)
  {
    // The ith sub-element. For mixed elements this is subelements()[i]. For
    // blocked elements, the sub-element will always be the same, so we'll use
    // sub_elements()[0]
    std::shared_ptr<const fem::FiniteElement<T>> sub_e
        = element.sub_elements()[bs > 1 ? 0 : i];
 
    // In a mixed element DOFs are ordered element by element, so the offset to
    // the next sub-element is sub_e->space_dimension(). Blocked elements use
    // xxyyzz ordering, so the offset to the next sub-element is 1
 
    std::vector<int> parent_map_sub(sub_e->space_dimension(), offsets.back());
    for (std::size_t j = 0; j < parent_map_sub.size(); ++j)
      parent_map_sub[j] += bs * j;
    offsets.push_back(offsets.back() + (bs > 1 ? 1 : sub_e->space_dimension()));
    sub_doflayout.push_back(
        dolfinx::fem::create_element_dof_layout(*sub_e, parent_map_sub));
  }
 
  return ElementDofLayout(bs, element.entity_dofs(),
                          element.entity_closure_dofs(), parent_map,
                          sub_doflayout);
}

DofMap
create_dofmap(MPI_Comm comm, const ElementDofLayout& layout,
              mesh::Topology& topology,
              const std::function<void(std::span<std::int32_t>, std::uint32_t)>&
                  permute_inv,
              const std::function<std::vector<int>(
                  const graph::AdjacencyList<std::int32_t>&)>& reorder_fn);

std::vector<DofMap> create_dofmaps(
    MPI_Comm comm, const std::vector<ElementDofLayout>& layouts,
    mesh::Topology& topology,
    const std::function<void(std::span<std::int32_t>, std::uint32_t)>&
        permute_inv,
    const std::function<std::vector<int>(
        const graph::AdjacencyList<std::int32_t>&)>& reorder_fn);

std::vector<std::string> get_coefficient_names(const ufcx_form& ufcx_form);

std::vector<std::string> get_constant_names(const ufcx_form& ufcx_form);

template <dolfinx::scalar T, std::floating_point U = scalar_value_t<T>>
Form<T, U> create_form_factory(
    const std::vector<std::reference_wrapper<const ufcx_form>>& ufcx_forms,
    const std::vector<std::shared_ptr<const FunctionSpace<U>>>& spaces,
    const std::vector<std::shared_ptr<const Function<T, U>>>& coefficients,
    const std::vector<std::shared_ptr<const Constant<T>>>& constants,
    const std::map<
        IntegralType,
        std::vector<std::pair<std::int32_t, std::span<const std::int32_t>>>>&
        subdomains,
    const std::vector<std::reference_wrapper<const mesh::EntityMap>>&
        entity_maps,
    std::shared_ptr<const mesh::Mesh<U>> mesh = nullptr)
{
  for (const ufcx_form& ufcx_form : ufcx_forms)
  {
    if (ufcx_form.rank != (int)spaces.size())
      throw std::runtime_error("Wrong number of argument spaces for Form.");
    if (ufcx_form.num_coefficients != (int)coefficients.size())
    {
      throw std::runtime_error("Mismatch between number of expected and "
                               "provided Form coefficients.");
    }
 
    // Check Constants for rank and size consistency
    if (ufcx_form.num_constants != (int)constants.size())
    {
      throw std::runtime_error(std::format(
          "Mismatch between number of expected and "
          "provided Form Constants. Expected {} constants, but got {}.",
          ufcx_form.num_constants, constants.size()));
    }
    for (std::size_t c = 0; c < constants.size(); ++c)
    {
      if (ufcx_form.constant_ranks[c] != (int)constants[c]->shape.size())
      {
        throw std::runtime_error(std::format(
            "Mismatch between expected and actual rank of "
            "Form Constant. Rank of Constant {} should be {}, but got rank {}.",
            c, ufcx_form.constant_ranks[c], constants[c]->shape.size()));
      }
      if (!std::equal(constants[c]->shape.begin(), constants[c]->shape.end(),
                      ufcx_form.constant_shapes[c]))
      {
        throw std::runtime_error(
            std::format("Mismatch between expected and actual shape of Form "
                        "Constant for Constant {}.",
                        c));
      }
    }
  }
 
  // Check argument function spaces
  for (std::size_t form_idx = 0; form_idx < ufcx_forms.size(); ++form_idx)
  {
    for (std::size_t i = 0; i < spaces.size(); ++i)
    {
      assert(spaces[i]->elements(form_idx));
      if (auto element_hash
          = ufcx_forms[form_idx].get().finite_element_hashes[i];
          element_hash != 0
          and element_hash
                  != spaces[i]->elements(form_idx)->basix_element().hash())
      {
        throw std::runtime_error(
            "Cannot create form. Elements are different to "
            "those used to compile the form.");
      }
    }
  }
 
  // Extract mesh from FunctionSpace, and check they are the same
  if (!mesh and !spaces.empty())
    mesh = spaces.front()->mesh();
  if (!mesh)
    throw std::runtime_error("No mesh could be associated with the Form.");
 
  auto topology = mesh->topology();
  assert(topology);
  const int tdim = topology->dim();
 
  // NOTE: This assumes all forms in mixed-topology meshes have the same
  // integral offsets. Since the UFL forms for each type of cell should be
  // the same, I think this assumption is OK.
  const int* integral_offsets = ufcx_forms[0].get().form_integral_offsets;
  std::array<int, 5> num_integrals_type;
  for (std::size_t i = 0; i < num_integrals_type.size(); ++i)
    num_integrals_type[i] = integral_offsets[i + 1] - integral_offsets[i];
 
  // Create vertices, if required
  if (num_integrals_type[vertex] > 0)
  {
    mesh->topology_mutable()->create_connectivity(0, tdim);
    mesh->topology_mutable()->create_connectivity(tdim, 0);
  }
 
  // Create facets, if required
  // NOTE: exterior_facet and interior_facet is declared in ufcx.h
  if (num_integrals_type[exterior_facet] > 0
      or num_integrals_type[interior_facet] > 0)
  {
    mesh->topology_mutable()->create_entities(tdim - 1);
    mesh->topology_mutable()->create_connectivity(tdim - 1, tdim);
    mesh->topology_mutable()->create_connectivity(tdim, tdim - 1);
  }
 
  // Create ridges, if required
  if (num_integrals_type[ridge] > 0)
  {
    mesh->topology_mutable()->create_entities(tdim - 2);
    mesh->topology_mutable()->create_connectivity(tdim - 2, tdim);
    mesh->topology_mutable()->create_connectivity(tdim, tdim - 2);
  }
 
  // Get list of integral IDs, and load tabulate tensor into memory for
  // each
  std::map<std::tuple<IntegralType, int, int>, integral_data<T, U>> integrals;
 
  auto check_geometry_hash
      = [&geo = mesh->geometry()](const ufcx_integral& integral,
                                  std::size_t cell_idx)
  {
    if (integral.coordinate_element_hash != geo.cmap(cell_idx).hash())
    {
      throw std::runtime_error(
          "Generated integral geometry element does not match mesh geometry: "
          + std::to_string(integral.coordinate_element_hash) + ", "
          + std::to_string(geo.cmap(cell_idx).hash()));
    }
  };
 
  // Attach cell kernels
  bool needs_facet_permutations = false;
  {
    std::vector<std::int32_t> default_cells;
    std::span<const int> ids(ufcx_forms[0].get().form_integral_ids
                                 + integral_offsets[cell],
                             num_integrals_type[cell]);
    auto sd = subdomains.find(IntegralType::cell);
    for (std::size_t form_idx = 0; form_idx < ufcx_forms.size(); ++form_idx)
    {
      const ufcx_form& ufcx_form = ufcx_forms[form_idx];
      for (int i = 0; i < num_integrals_type[cell]; ++i)
      {
        const int id = ids[i];
        ufcx_integral* integral
            = ufcx_form.form_integrals[integral_offsets[cell] + i];
        assert(integral);
        check_geometry_hash(*integral, form_idx);
 
        // Build list of active coefficients
        std::vector<int> active_coeffs;
        for (int j = 0; j < ufcx_form.num_coefficients; ++j)
        {
          if (integral->enabled_coefficients[j])
            active_coeffs.push_back(j);
        }
 
        impl::kernel_t<T, U> k = impl::extract_kernel<T>(integral);
        if (!k)
        {
          throw std::runtime_error(
              "UFCx kernel function is NULL. Check requested types.");
        }
 
        // Build list of entities to assemble over
        if (id == -1)
        {
          // Default kernel, operates on all (owned) cells
          assert(topology->index_maps(tdim).at(form_idx));
          default_cells.resize(
              topology->index_maps(tdim).at(form_idx)->size_local(), 0);
          std::iota(default_cells.begin(), default_cells.end(), 0);
          integrals.insert({{IntegralType::cell, i, form_idx},
                            {k, default_cells, active_coeffs}});
        }
        else if (sd != subdomains.end())
        {
          // NOTE: This requires that pairs are sorted
          auto it = std::ranges::lower_bound(sd->second, id, std::less<>{},
                                             [](auto& a) { return a.first; });
          if (it != sd->second.end() and it->first == id)
          {
            integrals.insert({{IntegralType::cell, i, form_idx},
                              {k,
                               std::vector<std::int32_t>(it->second.begin(),
                                                         it->second.end()),
                               active_coeffs}});
          }
        }
 
        if (integral->needs_facet_permutations)
          needs_facet_permutations = true;
      }
    }
  }
 
  // Attach interior facet kernels
  {
    std::vector<std::int32_t> default_facets_int;
    std::span<const int> ids(ufcx_forms[0].get().form_integral_ids
                                 + integral_offsets[interior_facet],
                             num_integrals_type[interior_facet]);
    auto sd = subdomains.find(IntegralType::interior_facet);
    for (std::size_t form_idx = 0; form_idx < ufcx_forms.size(); ++form_idx)
    {
      const ufcx_form& ufcx_form = ufcx_forms[form_idx];
 
      // Create indicator for interprocess facets
      std::vector<std::int8_t> interprocess_marker;
      if (num_integrals_type[interior_facet] > 0)
      {
        assert(topology->index_map(tdim - 1));
        const std::vector<std::int32_t>& interprocess_facets
            = topology->interprocess_facets();
        std::int32_t num_facets = topology->index_map(tdim - 1)->size_local()
                                  + topology->index_map(tdim - 1)->num_ghosts();
        interprocess_marker.resize(num_facets, 0);
        std::ranges::for_each(interprocess_facets,
                              [&interprocess_marker](auto f)
                              { interprocess_marker[f] = 1; });
      }
 
      for (int i = 0; i < num_integrals_type[interior_facet]; ++i)
      {
        const int id = ids[i];
        ufcx_integral* integral
            = ufcx_form.form_integrals[integral_offsets[interior_facet] + i];
        assert(integral);
        check_geometry_hash(*integral, form_idx);
 
        std::vector<int> active_coeffs;
        for (int j = 0; j < ufcx_form.num_coefficients; ++j)
        {
          if (integral->enabled_coefficients[j])
            active_coeffs.push_back(j);
        }
 
        impl::kernel_t<T, U> k = impl::extract_kernel<T>(integral);
        assert(k);
 
        // Build list of entities to assembler over
        auto f_to_c = topology->connectivity(tdim - 1, tdim);
        assert(f_to_c);
        auto c_to_f = topology->connectivity(tdim, tdim - 1);
        assert(c_to_f);
        if (id == -1)
        {
          // Default kernel, operates on all (owned) interior facets
          assert(topology->index_map(tdim - 1));
          std::int32_t num_facets = topology->index_map(tdim - 1)->size_local();
          default_facets_int.reserve(4 * num_facets);
          for (std::int32_t f = 0; f < num_facets; ++f)
          {
            if (f_to_c->num_links(f) == 2)
            {
              std::array<std::int32_t, 4> pairs
                  = impl::get_cell_facet_pairs<2>(f, f_to_c->links(f), *c_to_f);
              default_facets_int.insert(default_facets_int.end(), pairs.begin(),
                                        pairs.end());
            }
            else if (interprocess_marker[f])
            {
              throw std::runtime_error(
                  "Cannot compute interior facet integral over interprocess "
                  "facet. Please use ghost mode shared facet when creating the "
                  "mesh");
            }
          }
          integrals.insert({{IntegralType::interior_facet, i, form_idx},
                            {k, default_facets_int, active_coeffs}});
        }
        else if (sd != subdomains.end())
        {
          auto it = std::ranges::lower_bound(sd->second, id, std::less{},
                                             [](auto& a) { return a.first; });
          if (it != sd->second.end() and it->first == id)
          {
            integrals.insert({{IntegralType::interior_facet, i, form_idx},
                              {k,
                               std::vector<std::int32_t>(it->second.begin(),
                                                         it->second.end()),
                               active_coeffs}});
          }
        }
 
        if (integral->needs_facet_permutations)
          needs_facet_permutations = true;
      }
    }
  }
 
  // Attach exterior entity integrals
  {
    for (IntegralType itg_type : {IntegralType::exterior_facet,
                                  IntegralType::vertex, IntegralType::ridge})
    {
      std::size_t dim;
      switch (itg_type)
      {
      case IntegralType::exterior_facet:
      {
        dim = tdim - 1;
        break;
      }
      case IntegralType::ridge:
      {
        dim = tdim - 2;
        break;
      }
      case IntegralType::vertex:
      {
        dim = 0;
        break;
      }
      default:
        throw std::runtime_error("Unsupported integral type");
      }
 
      const std::function<std::vector<std::int32_t>(const mesh::Topology&,
                                                    IntegralType)>
          get_default_integration_entities
          = [dim](const mesh::Topology& topology, IntegralType itg_type)
      {
        if (itg_type == IntegralType::exterior_facet)
        {
          // Integrate over all owned exterior facets
          return mesh::exterior_facet_indices(topology);
        }
        else
        {
          // Integrate over all owned entities
          std::int32_t num_entities = topology.index_map(dim)->size_local();
          std::vector<std::int32_t> entities(num_entities);
          std::iota(entities.begin(), entities.end(), 0);
          return entities;
        }
      };
 
      std::vector<std::int32_t> default_entities_ext;
 
      std::span<const int> ids(ufcx_forms[0].get().form_integral_ids
                                   + integral_offsets[(std::int8_t)itg_type],
                               num_integrals_type[(std::int8_t)itg_type]);
      auto sd = subdomains.find(itg_type);
      for (std::size_t form_idx = 0; form_idx < ufcx_forms.size(); ++form_idx)
      {
        const ufcx_form& ufcx_form = ufcx_forms[form_idx];
        for (int i = 0; i < num_integrals_type[(std::int8_t)itg_type]; ++i)
        {
          const int id = ids[i];
          ufcx_integral* integral
              = ufcx_form.form_integrals[integral_offsets[(std::int8_t)itg_type]
                                         + i];
          assert(integral);
          check_geometry_hash(*integral, form_idx);
 
          std::vector<int> active_coeffs;
          for (int j = 0; j < ufcx_form.num_coefficients; ++j)
          {
            if (integral->enabled_coefficients[j])
              active_coeffs.push_back(j);
          }
 
          impl::kernel_t<T, U> k = impl::extract_kernel<T>(integral);
 
          // Build list of entities to assembler over
          auto e_to_c = topology->connectivity(dim, tdim);
          assert(e_to_c);
          auto c_to_e = topology->connectivity(tdim, dim);
          assert(c_to_e);
          if (id == -1)
          {
            std::vector default_entities
                = get_default_integration_entities(*topology, itg_type);
            // Default kernel
            default_entities_ext.reserve(2 * default_entities.size());
            for (std::int32_t e : default_entities)
            {
              // There will only be one pair for an exterior facet integral
              std::array<std::int32_t, 2> pair = impl::get_cell_entity_pairs<1>(
                  e, e_to_c->links(e), *c_to_e);
              default_entities_ext.insert(default_entities_ext.end(),
                                          pair.begin(), pair.end());
            }
            integrals.insert({{itg_type, i, form_idx},
                              {k, default_entities_ext, active_coeffs}});
          }
          else if (sd != subdomains.end())
          {
            // NOTE: This requires that pairs are sorted
            auto it = std::ranges::lower_bound(sd->second, id, std::less<>{},
                                               [](auto& a) { return a.first; });
            if (it != sd->second.end() and it->first == id)
            {
              integrals.insert({{itg_type, i, form_idx},
                                {k,
                                 std::vector<std::int32_t>(it->second.begin(),
                                                           it->second.end()),
                                 active_coeffs}});
            }
          }
 
          if (integral->needs_facet_permutations)
            needs_facet_permutations = true;
        }
      }
    }
  }
 
  return Form<T, U>(spaces, std::move(integrals), mesh, coefficients, constants,
                    needs_facet_permutations, entity_maps);
}

template <dolfinx::scalar T, std::floating_point U = scalar_value_t<T>>
Form<T, U> create_form(
    const ufcx_form& ufcx_form,
    const std::vector<std::shared_ptr<const FunctionSpace<U>>>& spaces,
    const std::map<std::string, std::shared_ptr<const Function<T, U>>>&
        coefficients,
    const std::map<std::string, std::shared_ptr<const Constant<T>>>& constants,
    const std::map<
        IntegralType,
        std::vector<std::pair<std::int32_t, std::span<const std::int32_t>>>>&
        subdomains,
    const std::vector<std::reference_wrapper<const mesh::EntityMap>>&
        entity_maps,
    std::shared_ptr<const mesh::Mesh<U>> mesh = nullptr)
{
  // Place coefficients in appropriate order
  std::vector<std::shared_ptr<const Function<T, U>>> coeff_map;
  for (const std::string& name : get_coefficient_names(ufcx_form))
  {
    if (auto it = coefficients.find(name); it != coefficients.end())
      coeff_map.push_back(it->second);
    else
    {
      throw std::runtime_error("Form coefficient \"" + name
                               + "\" not provided.");
    }
  }
 
  // Place constants in appropriate order
  std::vector<std::shared_ptr<const Constant<T>>> const_map;
  for (const std::string& name : get_constant_names(ufcx_form))
  {
    if (auto it = constants.find(name); it != constants.end())
      const_map.push_back(it->second);
    else
      throw std::runtime_error("Form constant \"" + name + "\" not provided.");
  }
 
  return create_form_factory({ufcx_form}, spaces, coeff_map, const_map,
                             subdomains, entity_maps, mesh);
}

template <dolfinx::scalar T, std::floating_point U = scalar_value_t<T>>
Form<T, U> create_form(
    ufcx_form* (*fptr)(),
    const std::vector<std::shared_ptr<const FunctionSpace<U>>>& spaces,
    const std::map<std::string, std::shared_ptr<const Function<T, U>>>&
        coefficients,
    const std::map<std::string, std::shared_ptr<const Constant<T>>>& constants,
    const std::map<
        IntegralType,
        std::vector<std::pair<std::int32_t, std::span<const std::int32_t>>>>&
        subdomains,
    const std::vector<std::reference_wrapper<const mesh::EntityMap>>&
        entity_maps,
    std::shared_ptr<const mesh::Mesh<U>> mesh = nullptr)
{
  ufcx_form* form = fptr();
  Form<T, U> L = create_form<T, U>(*form, spaces, coefficients, constants,
                                   subdomains, entity_maps, mesh);
  std::free(form);
  return L;
}

template <std::floating_point T>
FunctionSpace<T> create_functionspace(
    std::shared_ptr<mesh::Mesh<T>> mesh,
    std::shared_ptr<const fem::FiniteElement<T>> e,
    std::function<std::vector<int>(const graph::AdjacencyList<std::int32_t>&)>
        reorder_fn = nullptr)
{
  // TODO: check cell type of e (need to add method to fem::FiniteElement)
  assert(e);
  assert(mesh);
  assert(mesh->topology());
  if (e->cell_type() != mesh->topology()->cell_type())
    throw std::runtime_error("Cell type of element and mesh must match.");
 
  // Create element dof layout
  fem::ElementDofLayout layout = fem::create_element_dof_layout(*e);
 
  // Create a dofmap
  std::function<void(std::span<std::int32_t>, std::uint32_t)> permute_inv
      = e->needs_dof_permutations() ? e->dof_permutation_fn(true, true)
                                    : nullptr;
  auto dofmap = std::make_shared<const DofMap>(create_dofmap(
      mesh->comm(), layout, *mesh->topology(), permute_inv, reorder_fn));
 
  return FunctionSpace(mesh, e, dofmap);
}

template <dolfinx::scalar T, std::floating_point U = scalar_value_t<T>>
Expression<T, U> create_expression(
    const ufcx_expression& e,
    const std::vector<std::shared_ptr<const Function<T, U>>>& coefficients,
    const std::vector<std::shared_ptr<const Constant<T>>>& constants,
    std::shared_ptr<const FunctionSpace<U>> argument_space = nullptr)
{
  if (!coefficients.empty())
  {
    assert(coefficients.front());
    assert(coefficients.front()->function_space());
    std::shared_ptr<const mesh::Mesh<U>> mesh
        = coefficients.front()->function_space()->mesh();
    if (mesh->geometry().cmap().hash() != e.coordinate_element_hash)
    {
      throw std::runtime_error(
          "Expression and mesh geometric maps do not match.");
    }
  }
 
  if (e.rank > 0 and !argument_space)
  {
    throw std::runtime_error("Expression has Argument but no Argument "
                             "function space was provided.");
  }
 
  std::vector<U> X(e.points, e.points + e.num_points * e.entity_dimension);
  std::array<std::size_t, 2> Xshape
      = {static_cast<std::size_t>(e.num_points),
         static_cast<std::size_t>(e.entity_dimension)};
  std::vector<std::size_t> value_shape(e.value_shape,
                                       e.value_shape + e.num_components);
  std::function<void(T*, const T*, const T*, const scalar_value_t<T>*,
                     const int*, const std::uint8_t*, void*)>
      tabulate_tensor = nullptr;
  if constexpr (std::is_same_v<T, float>)
    tabulate_tensor = e.tabulate_tensor_float32;
#ifndef DOLFINX_NO_STDC_COMPLEX_KERNELS
  else if constexpr (std::is_same_v<T, std::complex<float>>)
  {
    tabulate_tensor = reinterpret_cast<void (*)(
        T*, const T*, const T*, const scalar_value_t<T>*, const int*,
        const unsigned char*, void*)>(e.tabulate_tensor_complex64);
  }
#endif // DOLFINX_NO_STDC_COMPLEX_KERNELS
  else if constexpr (std::is_same_v<T, double>)
    tabulate_tensor = e.tabulate_tensor_float64;
#ifndef DOLFINX_NO_STDC_COMPLEX_KERNELS
  else if constexpr (std::is_same_v<T, std::complex<double>>)
  {
    tabulate_tensor = reinterpret_cast<void (*)(
        T*, const T*, const T*, const scalar_value_t<T>*, const int*,
        const unsigned char*, void*)>(e.tabulate_tensor_complex128);
  }
#endif // DOLFINX_NO_STDC_COMPLEX_KERNELS
  else
    throw std::runtime_error("Type not supported.");
 
  assert(tabulate_tensor);
  return Expression(coefficients, constants, std::span<const U>(X), Xshape,
                    tabulate_tensor, value_shape, argument_space);
}

template <dolfinx::scalar T, std::floating_point U = scalar_value_t<T>>
Expression<T, U> create_expression(
    const ufcx_expression& e,
    const std::map<std::string, std::shared_ptr<const Function<T, U>>>&
        coefficients,
    const std::map<std::string, std::shared_ptr<const Constant<T>>>& constants,
    std::shared_ptr<const FunctionSpace<U>> argument_space = nullptr)
{
  // Place coefficients in appropriate order
  std::vector<std::shared_ptr<const Function<T, U>>> coeff_map;
  std::vector<std::string> coefficient_names;
  coefficient_names.reserve(e.num_coefficients);
  for (int i = 0; i < e.num_coefficients; ++i)
    coefficient_names.push_back(e.coefficient_names[i]);
 
  for (const std::string& name : coefficient_names)
  {
    if (auto it = coefficients.find(name); it != coefficients.end())
      coeff_map.push_back(it->second);
    else
    {
      throw std::runtime_error("Expression coefficient \"" + name
                               + "\" not provided.");
    }
  }
 
  // Place constants in appropriate order
  std::vector<std::shared_ptr<const Constant<T>>> const_map;
  std::vector<std::string> constant_names;
  constant_names.reserve(e.num_constants);
  for (int i = 0; i < e.num_constants; ++i)
    constant_names.push_back(e.constant_names[i]);
 
  for (const std::string& name : constant_names)
  {
    if (auto it = constants.find(name); it != constants.end())
      const_map.push_back(it->second);
    else
    {
      throw std::runtime_error("Expression constant \"" + name
                               + "\" not provided.");
    }
  }
 
  return create_expression(e, coeff_map, const_map, argument_space);
}

template <std::floating_point T>
mesh::Mesh<T> interpolate_geometry(
    std::shared_ptr<mesh::Mesh<T>> mesh, const CoordinateElement<T>& new_cmap,
    const std::function<std::vector<int>(
        const graph::AdjacencyList<std::int32_t>&)>& reorder_fn = nullptr)
{
  assert(mesh);
  const CoordinateElement<T>& old_cmap = mesh->geometry().cmap();
  if (new_cmap.cell_shape() != old_cmap.cell_shape())
  {
    throw std::runtime_error(
        "Cell shape of new coordinate element must match input mesh.");
  }
 
  const int gdim = mesh->geometry().dim();
 
  // Build a vector-valued Lagrange FiniteElement from the coordinate element.
  basix::FiniteElement<T> b_element = basix::create_element<T>(
      basix::element::family::P,
      mesh::cell_type_to_basix_type(new_cmap.cell_shape()), new_cmap.degree(),
      new_cmap.variant(), basix::element::dpc_variant::unset, false);
  auto element = std::make_shared<const FiniteElement<T>>(
      b_element, std::vector<std::size_t>{static_cast<std::size_t>(gdim)});
 
  FunctionSpace<T> V = create_functionspace(mesh, element, reorder_fn);
 
  // Tabulate physical coordinates of the new geometry dofs.
  std::vector<T> x_new = V.tabulate_dof_coordinates(false);
 
  // Pull the geometry dofmap and index map from V.
  std::shared_ptr<const DofMap> dm = V.dofmap();
  assert(dm);
  std::shared_ptr<const common::IndexMap> new_imap = dm->index_map;
  assert(new_imap);
 
  auto map_view = dm->map();
  std::vector<std::int32_t> dofmap_flat(
      map_view.data_handle(), map_view.data_handle() + map_view.size());
 
  // Build input_global_indices as the local-to-global of the new geometry
  // dofs.
  const std::int32_t num_nodes
      = new_imap->size_local() + new_imap->num_ghosts();
  std::vector<std::int32_t> local(num_nodes);
  std::iota(local.begin(), local.end(), 0);
  std::vector<std::int64_t> igi(num_nodes);
  new_imap->local_to_global(local, igi);
 
  mesh::Geometry<T> geometry(
      new_imap, std::vector<std::vector<std::int32_t>>{std::move(dofmap_flat)},
      std::vector<CoordinateElement<T>>{new_cmap}, std::move(x_new), gdim,
      std::move(igi));
 
  return mesh::Mesh<T>(mesh->comm(), mesh->topology(), std::move(geometry));
}
 
} // namespace dolfinx::fem