discreteoperators.h

// Copyright (C) 2015-2026 Garth N. Wells, Jørgen S. Dokken
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include "DofMap.h"
#include "ElementDofLayout.h"
#include "FiniteElement.h"
#include "FunctionSpace.h"
#include "sparsitybuild.h"
#include <algorithm>
#include <array>
#include <basix/sobolev-spaces.h>
#include <concepts>
#include <cstdint>
#include <dolfinx/common/IndexMap.h>
#include <dolfinx/common/math.h>
#include <dolfinx/la/utils.h>
#include <dolfinx/mesh/Mesh.h>
#include <memory>
#include <span>
#include <stdexcept>
#include <utility>
#include <vector>
 
namespace dolfinx::fem
{
 
namespace impl
{
template <std::floating_point T>
void check_derivative_source(const FiniteElement<T>& element)
{
  const basix::sobolev::space space = element.basix_element().sobolev_space();
  const bool h1
      = basix::sobolev::space_intersection(space, basix::sobolev::space::H1)
        == basix::sobolev::space::H1;
  if (element.map_type() == basix::maps::type::identity and !h1)
  {
    throw std::invalid_argument(
        "Source element must be H1-conforming for the discrete gradient.");
  }
  if (element.map_type() == basix::maps::type::covariantPiola and !h1
      and space != basix::sobolev::space::HCurl)
  {
    throw std::invalid_argument(
        "Source element must be H(curl)-conforming for the discrete curl.");
  }
}

template <dolfinx::scalar T>
void insert_entity_blocks(
    std::span<const std::pair<std::span<const int>, std::span<const int>>>
        blocks,
    std::span<const T> A, std::size_t ncols,
    std::span<const std::int32_t> cell_rows,
    std::span<const std::int32_t> cell_cols, std::span<std::int32_t> rows,
    std::span<std::int32_t> cols, std::span<T> Ae, auto&& mat_set)
{
  for (const auto& [rdofs, cdofs] : blocks)
  {
    std::ranges::transform(cdofs, cols.begin(),
                           [cell_cols](int d) { return cell_cols[d]; });
    for (std::size_t i = 0; i < rdofs.size(); ++i)
    {
      rows[i] = cell_rows[rdofs[i]];
      for (std::size_t j = 0; j < cdofs.size(); ++j)
        Ae[i * cdofs.size() + j] = A[rdofs[i] * ncols + cdofs[j]];
    }
    mat_set(rows.first(rdofs.size()), cols.first(cdofs.size()),
            Ae.first(rdofs.size() * cdofs.size()));
  }
}
} // namespace impl

template <std::floating_point T, dolfinx::scalar U = T>
void discrete_curl(const FunctionSpace<T>& V0, const FunctionSpace<T>& V1,
                   la::MatSet<U> auto&& mat_set)
{
  // Get mesh
  auto mesh = V0.mesh();
  assert(mesh);
  assert(V1.mesh());
  if (mesh != V1.mesh())
    throw std::invalid_argument("Meshes must be the same.");
 
  if (mesh->geometry().dim() != 3)
    throw std::invalid_argument("Geometric must be equal to 3..");
  if (mesh->geometry().dim() != mesh->topology()->dim())
  {
    throw std::invalid_argument(
        "Geometric and topological dimensions must be equal.");
  }
  constexpr int gdim = 3;
 
  // Get elements
  std::shared_ptr<const FiniteElement<T>> e0 = V0.element();
  assert(e0);
  if (e0->map_type() != basix::maps::type::covariantPiola)
  {
    throw std::invalid_argument(
        "Finite element for parent space must be covariant Piola.");
  }
  impl::check_derivative_source(*e0);
 
  std::shared_ptr<const FiniteElement<T>> e1 = V1.element();
  assert(e1);
  if (e1->map_type() != basix::maps::type::contravariantPiola)
  {
    throw std::invalid_argument(
        "Finite element for target space must be contracovariant Piola.");
  }
 
  // Get cell orientation information
  std::span<const std::uint32_t> cell_info;
  if (e1->needs_dof_transformations() or e0->needs_dof_transformations())
  {
    mesh->topology_mutable()->create_cell_permutations();
    cell_info = std::span(mesh->topology()->get_cell_permutation_info());
  }
 
  // Get dofmaps
  auto dofmap0 = V0.dofmap();
  assert(dofmap0);
  auto dofmap1 = V1.dofmap();
  assert(dofmap1);
 
  // Get dof transformation operators
  auto apply_dof_transformation0
      = e0->template dof_transformation_fn<T>(doftransform::standard, false);
  auto apply_inverse_dof_transform1 = e1->template dof_transformation_fn<U>(
      doftransform::inverse_transpose, false);
 
  // Get sizes of elements
  const std::size_t space_dim0 = e0->space_dimension();
  const std::size_t space_dim1 = e1->space_dimension();
  if (e0->reference_value_size() != 3)
    throw std::invalid_argument("Value size for parent space should be 3.");
  if (e1->reference_value_size() != 3)
    throw std::invalid_argument("Value size for target space should be 3.");
 
  // Get the V1 reference interpolation points
  const auto [X, Xshape] = e1->interpolation_points();
 
  // Evaluate V0 basis function derivatives at reference interpolation
  // points for V1, (deriv, pt_idx, phi (dof), comp)
  const auto [Phi0_b, Phi0_shape] = e0->tabulate(X, Xshape, 1);
  md::mdspan<const T, md::extents<std::size_t, 4, md::dynamic_extent,
                                  md::dynamic_extent, 3>>
      Phi0(Phi0_b.data(), Phi0_shape);
 
  // Create working arrays, (point, phi (dof), deriv, comp)
  md::extents<std::size_t, md::dynamic_extent, md::dynamic_extent, 3, 3>
      dPhi0_ext(Phi0.extent(1), Phi0.extent(2), Phi0.extent(0) - 1,
                Phi0.extent(3));
  std::vector<T> dPhi0_b(dPhi0_ext.extent(0) * dPhi0_ext.extent(1)
                         * dPhi0_ext.extent(2) * dPhi0_ext.extent(3));
  md::mdspan<T, decltype(dPhi0_ext)> dPhi0(dPhi0_b.data(), dPhi0_ext);
 
  // Get the interpolation operator (matrix) Pi that maps a function
  // evaluated at the interpolation points to the V1 element degrees of
  // freedom, i.e. dofs = Pi f_x
  const auto [Pi1_b, pi_shape] = e1->interpolation_operator();
  md::mdspan<const T, md::dextents<std::size_t, 2>> Pi_1(Pi1_b.data(),
                                                         pi_shape);
 
  // curl data structure, (pt_idx, phi (dof), comp)
  std::vector<T> curl_b(dPhi0.extent(0) * dPhi0.extent(1) * dPhi0.extent(3));
  md::mdspan<
      T, md::extents<std::size_t, md::dynamic_extent, md::dynamic_extent, 3>>
      curl(curl_b.data(), dPhi0.extent(0), dPhi0.extent(1), dPhi0.extent(3));
 
  std::vector<U> Ab(space_dim0 * space_dim1);
 
  // Insert the element matrix one mesh entity at a time, so that the
  // operator gets its exact sparsity
  const std::vector<std::pair<std::span<const int>, std::span<const int>>>
      blocks = sparsitybuild::entity_closure_blocks(
          dofmap1->element_dof_layout(), dofmap0->element_dof_layout());
  std::vector<U> Ab_e(Ab.size());
  std::vector<std::int32_t> rows(space_dim1), cols(space_dim0);
 
  // Iterate over mesh and interpolate on each cell
  assert(mesh->topology()->index_map(gdim));
  for (std::int32_t c = 0; c < mesh->topology()->index_map(gdim)->size_local();
       ++c)
  {
    // TODO: re-order loops and/or re-pack Phi0 to allow a simple flat
    // copy?
 
    // Copy (d)Phi0 (on reference) and apply DOF transformation
    // Phi0:  (deriv, pt_idx, phi (dof), comp)
    // dPhi0: (pt_idx, phi (dov), deriv, comp)
    for (std::size_t p = 0; p < Phi0.extent(1); ++p)           // point
      for (std::size_t phi = 0; phi < Phi0.extent(2); ++phi)   // phi_i
        for (std::size_t d = 0; d < Phi0.extent(3); ++d)       // Comp. of phi
          for (std::size_t dx = 0; dx < dPhi0.extent(2); ++dx) // dx
            dPhi0(p, phi, dx, d) = Phi0(dx + 1, p, phi, d);
 
    for (std::size_t p = 0; p < dPhi0.extent(0); ++p) // point
    {
      // Size: num_phi * num_derivs * num_components
      std::size_t size = dPhi0.extent(1) * dPhi0.extent(2) * dPhi0.extent(3);
      std::size_t offset = p * size; // Offset for point p
 
      // Shape: (num_phi , (value_size * num_derivs))
      if (apply_dof_transformation0)
      {
        apply_dof_transformation0(std::span(dPhi0.data_handle() + offset, size),
                                  cell_info, c,
                                  dPhi0.extent(2) * dPhi0.extent(3));
      }
    }
 
    // Compute curl
    // dPhi0: (pt_idx, phi_idx, deriv, comp)
    // curl: (pt_idx, phi_idx, comp)
    for (std::size_t p = 0; p < curl.extent(0); ++p) // point
    {
      for (std::size_t i = 0; i < curl.extent(1); ++i) // phi_i
      {
        curl(p, i, 0) = dPhi0(p, i, 1, 2) - dPhi0(p, i, 2, 1);
        curl(p, i, 1) = dPhi0(p, i, 2, 0) - dPhi0(p, i, 0, 2);
        curl(p, i, 2) = dPhi0(p, i, 0, 1) - dPhi0(p, i, 1, 0);
      }
    }
 
    // Apply interpolation matrix to basis derivative values of V0 at
    // the interpolation points of V1. Pi_1 does not depend on the
    // basis function index i, so all space_dim0 applications are
    // combined into a single loop nest (rather than calling
    // interpolation_apply once per basis function) so that each row
    // of Pi_1 is read from memory once instead of space_dim0 times.
    std::ranges::fill(Ab, U(0));
    for (std::size_t j = 0; j < space_dim1; ++j)         // V1 dof
      for (std::size_t k = 0; k < curl.extent(2); ++k)   // component
        for (std::size_t p = 0; p < curl.extent(0); ++p) // point
        {
          T pi_val = Pi_1(j, k * curl.extent(0) + p);
          for (std::size_t i = 0; i < space_dim0; ++i) // V0 basis function
            Ab[space_dim0 * j + i] += static_cast<U>(pi_val * curl(p, i, k));
        }
 
    if (apply_inverse_dof_transform1)
      apply_inverse_dof_transform1(Ab, cell_info, c, space_dim0);
 
    impl::insert_entity_blocks<U>(blocks, Ab, space_dim0, dofmap1->cell_dofs(c),
                                  dofmap0->cell_dofs(c), rows, cols, Ab_e,
                                  mat_set);
  }
}

template <dolfinx::scalar T, std::floating_point U = dolfinx::scalar_value_t<T>>
void discrete_gradient(mesh::Topology& topology,
                       std::pair<std::reference_wrapper<const FiniteElement<U>>,
                                 std::reference_wrapper<const DofMap>>
                           V0,
                       std::pair<std::reference_wrapper<const FiniteElement<U>>,
                                 std::reference_wrapper<const DofMap>>
                           V1,
                       auto&& mat_set)
{
  auto& e0 = V0.first.get();
  const DofMap& dofmap0 = V0.second.get();
  auto& e1 = V1.first.get();
  const DofMap& dofmap1 = V1.second.get();
 
  using cmdspan2_t = md::mdspan<const U, md::dextents<std::size_t, 2>>;
  using cmdspan4_t = md::mdspan<const U, md::dextents<std::size_t, 4>>;
 
  // Check elements
  if (e0.map_type() != basix::maps::type::identity)
    throw std::invalid_argument("Wrong finite element space for V0.");
  if (e0.block_size() != 1)
    throw std::invalid_argument("Block size is greater than 1 for V0.");
  if (e0.reference_value_size() != 1)
    throw std::invalid_argument("Wrong value size for V0.");
  impl::check_derivative_source(e0);
 
  if (e1.map_type() != basix::maps::type::covariantPiola)
    throw std::invalid_argument("Wrong finite element space for V1.");
  if (e1.block_size() != 1)
    throw std::invalid_argument("Block size is greater than 1 for V1.");
 
  // Get V1 (H(curl)) space interpolation points
  const auto [X, Xshape] = e1.interpolation_points();
 
  // Tabulate first derivatives of Lagrange space at V1 interpolation
  // points
  const int ndofs0 = e0.space_dimension();
  const int tdim = topology.dim();
  std::vector<U> phi0_b((tdim + 1) * Xshape[0] * ndofs0 * 1);
  cmdspan4_t phi0(phi0_b.data(), tdim + 1, Xshape[0], ndofs0, 1);
  e0.tabulate(phi0_b, X, Xshape, 1);
 
  // Reshape Lagrange basis derivatives as a matrix of shape (tdim *
  // num_points, num_dofs_per_cell)
  cmdspan2_t dphi_reshaped(
      phi0_b.data() + phi0.extent(3) * phi0.extent(2) * phi0.extent(1),
      tdim * phi0.extent(1), phi0.extent(2));
 
  // Get inverse DOF transform function
  auto apply_inverse_dof_transform = e1.template dof_transformation_fn<T>(
      doftransform::inverse_transpose, false);
 
  // Generate cell permutations
  topology.create_cell_permutations();
  const std::vector<std::uint32_t>& cell_info
      = topology.get_cell_permutation_info();
 
  // Create element kernel function
 
  // Build the element interpolation matrix
  std::vector<T> Ab(e1.space_dimension() * ndofs0);
  {
    md::mdspan<T, md::dextents<std::size_t, 2>> A(Ab.data(),
                                                  e1.space_dimension(), ndofs0);
    const auto [Pi, shape] = e1.interpolation_operator();
    cmdspan2_t _Pi(Pi.data(), shape);
    math::dot(_Pi, dphi_reshaped, A);
  }
 
  // Insert the element matrix one mesh entity at a time, so that the
  // operator gets its exact sparsity
  const std::vector<std::pair<std::span<const int>, std::span<const int>>>
      blocks = sparsitybuild::entity_closure_blocks(
          dofmap1.element_dof_layout(), dofmap0.element_dof_layout());
  auto cell_map = topology.index_map(tdim);
  assert(cell_map);
  std::int32_t num_cells = cell_map->size_local();
  std::vector<T> Ae(Ab.size()), Ab_e(Ab.size());
  std::vector<std::int32_t> rows(e1.space_dimension()), cols(ndofs0);
  for (std::int32_t c = 0; c < num_cells; ++c)
  {
    std::ranges::copy(Ab, Ae.begin());
    if (apply_inverse_dof_transform)
      apply_inverse_dof_transform(Ae, cell_info, c, ndofs0);
    impl::insert_entity_blocks<T>(blocks, Ae, ndofs0, dofmap1.cell_dofs(c),
                                  dofmap0.cell_dofs(c), rows, cols, Ab_e,
                                  mat_set);
  }
}

template <dolfinx::scalar T, std::floating_point U>
void interpolation_matrix(const FunctionSpace<U>& V0,
                          const FunctionSpace<U>& V1, auto&& mat_add)
{
  // Get mesh
  auto mesh = V0.mesh();
  assert(mesh);
 
  // Mesh dims
  const int tdim = mesh->topology()->dim();
  const int gdim = mesh->geometry().dim();
 
  // Get elements
  std::shared_ptr<const FiniteElement<U>> e0 = V0.element();
  assert(e0);
  std::shared_ptr<const FiniteElement<U>> e1 = V1.element();
  assert(e1);
 
  if (!std::ranges::equal(e0->value_shape(), e1->value_shape()))
  {
    throw std::invalid_argument(
        "Interpolation operator: elements have different value dimensions");
  }
 
  std::span<const std::uint32_t> cell_info;
  if (e1->needs_dof_transformations() or e0->needs_dof_transformations())
  {
    mesh->topology_mutable()->create_cell_permutations();
    cell_info = std::span(mesh->topology()->get_cell_permutation_info());
  }
 
  // Get dofmaps
  auto dofmap0 = V0.dofmap();
  assert(dofmap0);
  auto dofmap1 = V1.dofmap();
  assert(dofmap1);
 
  // Get block sizes and dof transformation operators
  const int bs0 = e0->block_size();
  const int bs1 = e1->block_size();
  auto apply_dof_transformation0
      = e0->template dof_transformation_fn<U>(doftransform::standard, false);
  auto apply_inverse_dof_transform1 = e1->template dof_transformation_fn<T>(
      doftransform::inverse_transpose, false);
 
  // Get sizes of elements
  const std::size_t space_dim0 = e0->space_dimension();
  const std::size_t space_dim1 = e1->space_dimension();
  const std::size_t dim0 = space_dim0 / bs0;
  const std::size_t value_size_ref0 = e0->reference_value_size();
  // basis0 holds one block of V0's basis pushed forward to the physical
  // cell; basis_values holds all bs0 blocks; mapped_values holds them
  // pulled back to the reference cell of e1.
  const std::size_t value_size0 = e0->physical_base_value_size();
  const std::size_t value_size_phys = e0->value_size();
  const std::size_t value_size_ref1
      = e1->reference_value_size() * static_cast<std::size_t>(bs1);
 
  // Get geometry data
  const CoordinateElement<U>& cmap = mesh->geometry().cmaps().front();
  auto x_dofmap = mesh->geometry().dofmaps().front();
  const std::size_t num_dofs_g = cmap.dim();
  std::span<const U> x_g = mesh->geometry().x();
 
  using mdspan2_t = md::mdspan<U, md::dextents<std::size_t, 2>>;
  using cmdspan2_t = md::mdspan<const U, md::dextents<std::size_t, 2>>;
  using cmdspan4_t = md::mdspan<const U, md::dextents<std::size_t, 4>>;
  using mdspan3_t = md::mdspan<U, md::dextents<std::size_t, 3>>;
 
  // Evaluate coordinate map basis at reference interpolation points
  const auto [X, Xshape] = e1->interpolation_points();
  std::array<std::size_t, 4> phi_shape = cmap.tabulate_shape(1, Xshape[0]);
  std::vector<U> phi_b(
      std::reduce(phi_shape.begin(), phi_shape.end(), 1, std::multiplies{}));
  cmdspan4_t phi(phi_b.data(), phi_shape);
  cmap.tabulate(1, X, Xshape, phi_b);
 
  // Evaluate V0 basis functions at reference interpolation points for V1
  std::vector<U> basis_derivatives_reference0_b(Xshape[0] * dim0
                                                * value_size_ref0);
  cmdspan4_t basis_derivatives_reference0(basis_derivatives_reference0_b.data(),
                                          1, Xshape[0], dim0, value_size_ref0);
  e0->tabulate(basis_derivatives_reference0_b, X, Xshape, 0);
 
  // Clamp values
  std::ranges::transform(
      basis_derivatives_reference0_b, basis_derivatives_reference0_b.begin(),
      [atol = 1e-14](auto x) { return std::abs(x) < atol ? 0.0 : x; });
 
  // Create working arrays
  std::vector<U> basis_reference0_b(Xshape[0] * dim0 * value_size_ref0);
  mdspan3_t basis_reference0(basis_reference0_b.data(), Xshape[0], dim0,
                             value_size_ref0);
  std::vector<U> J_b(Xshape[0] * gdim * tdim);
  mdspan3_t J(J_b.data(), Xshape[0], gdim, tdim);
  std::vector<U> K_b(Xshape[0] * tdim * gdim);
  mdspan3_t K(K_b.data(), Xshape[0], tdim, gdim);
  std::vector<U> detJ(Xshape[0]);
  std::vector<U> det_scratch(2 * tdim * gdim);
 
  // Get the interpolation operator (matrix) `Pi` that maps a function
  // evaluated at the interpolation points to the element degrees of
  // freedom, i.e. dofs = Pi f_x
  const auto [_Pi_1, pi_shape] = e1->interpolation_operator();
  cmdspan2_t Pi_1(_Pi_1.data(), pi_shape);
 
  bool interpolation_ident = e1->interpolation_ident();
 
  using u_t = md::mdspan<U, md::dextents<std::size_t, 2>>;
  using U_t = md::mdspan<const U, md::dextents<std::size_t, 2>>;
  using J_t = md::mdspan<const U, md::dextents<std::size_t, 2>>;
  using K_t = md::mdspan<const U, md::dextents<std::size_t, 2>>;
  auto push_forward_fn0
      = e0->basix_element().template map_fn<u_t, U_t, J_t, K_t>();
 
  // Basis values of Lagrange space unrolled for block size
  // (num_quadrature_points, Lagrange dof, value_size)
  std::vector<U> basis_values_b(Xshape[0] * bs0 * dim0 * value_size_phys);
  mdspan3_t basis_values(basis_values_b.data(), Xshape[0], bs0 * dim0,
                         value_size_phys);
  std::vector<U> mapped_values_b(Xshape[0] * bs0 * dim0 * value_size_ref1);
  mdspan3_t mapped_values(mapped_values_b.data(), Xshape[0], bs0 * dim0,
                          value_size_ref1);
 
  auto pull_back_fn1
      = e1->basix_element().template map_fn<u_t, U_t, K_t, J_t>();
 
  std::vector<U> coord_dofs_b(num_dofs_g * gdim);
  mdspan2_t coord_dofs(coord_dofs_b.data(), num_dofs_g, gdim);
  std::vector<U> basis0_b(Xshape[0] * dim0 * value_size0);
  mdspan3_t basis0(basis0_b.data(), Xshape[0], dim0, value_size0);
 
  // Buffers
  std::vector<T> Ab(space_dim0 * space_dim1);
 
  // Iterate over mesh and interpolate on each cell
  auto cell_map = mesh->topology()->index_map(tdim);
  assert(cell_map);
  std::int32_t num_cells = cell_map->size_local();
  int row_bs = dofmap1->bs();
  std::int32_t num_owned_rows
      = dofmap1->index_map->size_local() * dofmap1->index_map_bs() / row_bs;
  std::int32_t num_ghosted_rows
      = dofmap1->index_map->num_ghosts() * dofmap1->index_map_bs() / row_bs;
  std::vector<std::int8_t> row_added(num_owned_rows + num_ghosted_rows, 0);
 
  for (std::int32_t c = 0; c < num_cells; ++c)
  {
    // Get cell geometry (coordinate dofs)
    auto x_dofs = md::submdspan(x_dofmap, c, md::full_extent);
    for (std::size_t i = 0; i < x_dofs.size(); ++i)
    {
      for (int j = 0; j < gdim; ++j)
        coord_dofs(i, j) = x_g[3 * x_dofs[i] + j];
    }
 
    // Compute Jacobians and reference points for current cell. For an
    // affine map the Jacobian (and hence its inverse and determinant)
    // is constant over the cell, so it is computed once and broadcast
    // to the remaining interpolation points rather than being
    // recomputed Xshape[0] times.
    std::ranges::fill(J_b, 0);
    if (cmap.is_affine() and Xshape[0] > 0)
    {
      auto dphi
          = md::submdspan(phi, std::pair(1, tdim + 1), 0, md::full_extent, 0);
      auto _J = md::submdspan(J, 0, md::full_extent, md::full_extent);
      cmap.compute_jacobian(dphi, coord_dofs, _J);
      auto _K = md::submdspan(K, 0, md::full_extent, md::full_extent);
      cmap.compute_jacobian_inverse(_J, _K);
      detJ[0] = cmap.compute_jacobian_determinant(_J, det_scratch);
      for (std::size_t p = 1; p < Xshape[0]; ++p)
      {
        std::copy_n(J_b.begin(), gdim * tdim, J_b.begin() + p * gdim * tdim);
        std::copy_n(K_b.begin(), tdim * gdim, K_b.begin() + p * tdim * gdim);
        detJ[p] = detJ[0];
      }
    }
    else
    {
      for (std::size_t p = 0; p < Xshape[0]; ++p)
      {
        auto dphi
            = md::submdspan(phi, std::pair(1, tdim + 1), p, md::full_extent, 0);
        auto _J = md::submdspan(J, p, md::full_extent, md::full_extent);
        cmap.compute_jacobian(dphi, coord_dofs, _J);
        auto _K = md::submdspan(K, p, md::full_extent, md::full_extent);
        cmap.compute_jacobian_inverse(_J, _K);
        detJ[p] = cmap.compute_jacobian_determinant(_J, det_scratch);
      }
    }
 
    // Copy evaluated basis on reference (a fresh copy is needed each
    // cell as DOF transformations below are applied in place), and
    // push forward to physical element. The source and destination
    // have identical memory layout (the leading, unit-length
    // derivative axis of basis_derivatives_reference0 is a no-op for
    // indexing purposes), so a flat copy is used instead of an
    // element-by-element mdspan loop.
    std::ranges::copy(basis_derivatives_reference0_b,
                      basis_reference0_b.begin());
    if (apply_dof_transformation0)
    {
      for (std::size_t p = 0; p < Xshape[0]; ++p)
      {
        apply_dof_transformation0(std::span(basis_reference0.data_handle()
                                                + p * dim0 * value_size_ref0,
                                            dim0 * value_size_ref0),
                                  cell_info, c, value_size_ref0);
      }
    }
 
    for (std::size_t p = 0; p < basis0.extent(0); ++p)
    {
      auto _u = md::submdspan(basis0, p, md::full_extent, md::full_extent);
      auto _U = md::submdspan(basis_reference0, p, md::full_extent,
                              md::full_extent);
      auto _K = md::submdspan(K, p, md::full_extent, md::full_extent);
      auto _J = md::submdspan(J, p, md::full_extent, md::full_extent);
      push_forward_fn0(_u, _U, _J, detJ[p], _K);
    }
 
    // Unroll basis function for input space for block size
    for (std::size_t p = 0; p < Xshape[0]; ++p)
      for (std::size_t i = 0; i < dim0; ++i)
        for (std::size_t j = 0; j < value_size0; ++j)
          for (int k = 0; k < bs0; ++k)
            basis_values(p, i * bs0 + k, j * bs0 + k) = basis0(p, i, j);
 
    // Pull back the physical values to the reference of output space
    for (std::size_t p = 0; p < basis_values.extent(0); ++p)
    {
      auto _u
          = md::submdspan(basis_values, p, md::full_extent, md::full_extent);
      auto _U
          = md::submdspan(mapped_values, p, md::full_extent, md::full_extent);
      auto _K = md::submdspan(K, p, md::full_extent, md::full_extent);
      auto _J = md::submdspan(J, p, md::full_extent, md::full_extent);
      pull_back_fn1(_U, _u, _K, 1.0 / detJ[p], _J);
    }
 
    // Apply interpolation matrix to basis values of V0 at the
    // interpolation points of V1
    if (interpolation_ident)
    {
      md::mdspan<T, md::dextents<std::size_t, 3>> A(
          Ab.data(), Xshape[0], value_size_ref1, space_dim0);
      for (std::size_t i = 0; i < mapped_values.extent(0); ++i)
        for (std::size_t j = 0; j < mapped_values.extent(1); ++j)
          for (std::size_t k = 0; k < mapped_values.extent(2); ++k)
            A(i, k, j) = mapped_values(i, j, k);
    }
    else
    {
      // Apply interpolation matrix to basis values of V0 at the
      // interpolation points of V1. Pi_1 does not depend on the basis
      // function index i, so all space_dim0 applications are combined
      // into a single loop nest (rather than calling
      // interpolation_apply once per basis function) so that each row
      // of Pi_1 is read from memory once instead of space_dim0 times.
      // This mirrors the two cases handled by interpolation_apply
      // (interpolate.h): bs1 == 1 (columns of Pi_1 fold the point and
      // component indices together) and bs1 != 1 (each block/component
      // is handled by a separate outer loop, columns of Pi_1 index
      // points only).
      std::ranges::fill(Ab, T(0));
      std::size_t num_pts = mapped_values.extent(0);
      if (bs1 == 1)
      {
        for (std::size_t idof = 0; idof < Pi_1.extent(0); ++idof)
          for (std::size_t k = 0; k < mapped_values.extent(2); ++k)
            for (std::size_t p = 0; p < num_pts; ++p)
            {
              U pi_val = Pi_1(idof, k * num_pts + p);
              for (std::size_t i = 0; i < space_dim0; ++i) // V0 basis fn
              {
                Ab[space_dim0 * idof + i]
                    += static_cast<T>(pi_val * mapped_values(p, i, k));
              }
            }
      }
      else
      {
        for (std::size_t idof = 0; idof < Pi_1.extent(0); ++idof)
          for (int k = 0; k < bs1; ++k)
            for (std::size_t p = 0; p < num_pts; ++p)
            {
              U pi_val = Pi_1(idof, p);
              for (std::size_t i = 0; i < space_dim0; ++i) // V0 basis fn
              {
                Ab[space_dim0 * (bs1 * idof + k) + i]
                    += static_cast<T>(pi_val * mapped_values(p, i, k));
              }
            }
      }
    }
 
    if (apply_inverse_dof_transform1)
      apply_inverse_dof_transform1(Ab, cell_info, c, space_dim0);
 
    // Zero out all rows that have already been added on this process
    // and only add owned rows
    {
      md::mdspan<T, md::dextents<std::size_t, 2>> A(Ab.data(), space_dim1,
                                                    space_dim0);
      auto row_dofs = dofmap1->cell_dofs(c);
      for (std::size_t i = 0; i < row_dofs.size(); ++i)
      {
        std::int32_t r = row_dofs[i];
        if (r >= num_owned_rows || row_added[r])
        {
          for (std::size_t j = 0; j < space_dim0; ++j)
            for (int k = 0; k < row_bs; ++k)
              A(i * row_bs + k, j) = 0.0;
        }
        row_added[r] = 1;
      }
    }
    mat_add(dofmap1->cell_dofs(c), dofmap0->cell_dofs(c), Ab);
  }
}
 
} // namespace dolfinx::fem