assemble_expression_impl.h

// Copyright (C) 2025-2026 Garth N. Wells and 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 "FiniteElement.h"
#include "traits.h"
#include <algorithm>
#include <basix/mdspan.hpp>
#include <cassert>
#include <cstdint>
#include <dolfinx/common/IndexMap.h>
#include <dolfinx/mesh/Geometry.h>
#include <dolfinx/mesh/Mesh.h>
#include <dolfinx/mesh/Topology.h>
#include <functional>
#include <memory>
#include <optional>
#include <span>
#include <type_traits>
#include <utility>
#include <vector>
 
namespace dolfinx::fem::impl
{
template <dolfinx::scalar T, std::floating_point U>
void tabulate_expression(
    std::span<T> values, const fem::FEkernel<T, U> auto& fn,
    std::array<std::size_t, 2> Xshape, std::size_t value_size,
    std::size_t num_argument_dofs,
    md::mdspan<const std::int32_t, md::dextents<std::size_t, 2>> x_dofmap,
    std::span<const U> x,
    md::mdspan<const T, md::dextents<std::size_t, 2>> coeffs,
    std::span<const T> constants, fem::MDSpan2 auto entities,
    std::span<const std::uint32_t> cell_info,
    std::span<const std::int32_t> argument_cells,
    const fem::DofTransformKernel<T> auto& P0,
    md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>> perms)
{
  static_assert(entities.rank() == 1 or entities.rank() == 2);
 
  // Create data structures used in evaluation
  std::vector<U> coord_dofs(3 * x_dofmap.extent(1));
  GeometryPack geometry{
      x_dofmap,
      md::mdspan<const U, md::extents<std::size_t, md::dynamic_extent, 3>>(
          x.data(), x.size() / 3, 3)};
 
  // Iterate over cells and 'assemble' into values
  int size0 = Xshape[0] * value_size;
  std::vector<T> values_local(size0 * num_argument_dofs, 0);
  std::size_t offset = values_local.size();
 
  const T* coeffs_data = coeffs.data_handle();
  const std::size_t cstride = coeffs.extent(1);
 
  for (std::size_t e = 0; e < entities.extent(0); ++e)
  {
    std::ranges::fill(values_local, T(0));
    if constexpr (entities.rank() == 1)
    {
      std::int32_t entity = entities(e);
      gather_cell_coordinates(geometry, entity, coord_dofs.data());
      fn(values_local.data(), coeffs_data + e * cstride, constants.data(),
         coord_dofs.data(), nullptr, nullptr, nullptr);
 
      P0(values_local, cell_info,
         argument_cells.empty() ? entity : argument_cells[e], size0);
    }
    else
    {
      std::int32_t entity = entities(e, 0);
      std::int32_t local_entity = entities(e, 1);
      std::uint8_t perm = perms.empty() ? 0 : perms(entity, local_entity);
      gather_cell_coordinates(geometry, entity, coord_dofs.data());
      fn(values_local.data(), coeffs_data + e * cstride, constants.data(),
         coord_dofs.data(), &local_entity, &perm, nullptr);
      P0(values_local, cell_info,
         argument_cells.empty() ? entity : argument_cells[e], size0);
    }
 
    for (std::size_t j = 0; j < values_local.size(); ++j)
      values[e * offset + j] = values_local[j];
  }
}

template <dolfinx::scalar T, std::floating_point U>
void tabulate_expression(
    std::span<T> values, const fem::FEkernel<T, U> auto& fn,
    std::array<std::size_t, 2> Xshape, std::size_t value_size,
    md::mdspan<const T, md::dextents<std::size_t, 2>> coeffs,
    std::span<const T> constants, const mesh::Mesh<U>& mesh,
    fem::MDSpan2 auto entities,
    std::optional<
        std::pair<std::reference_wrapper<const FiniteElement<U>>, std::size_t>>
        element,
    std::span<const std::uint32_t> cell_info,
    std::span<const std::int32_t> argument_cells)
{
  assert(argument_cells.empty() or argument_cells.size() == entities.extent(0));
 
  std::function<void(std::span<T>, std::span<const std::uint32_t>, std::int32_t,
                     int)>
      post_dof_transform
      = [](std::span<T>, std::span<const std::uint32_t>, std::int32_t, int)
  {
    // Do nothing
  };
 
  std::size_t num_argument_dofs = 1;
  if (element)
  {
    num_argument_dofs = element->second;
    if (element->first.get().needs_dof_transformations())
    {
      post_dof_transform
          = element->first.get().template dof_transformation_right_fn<T>(
              doftransform::transpose);
    }
  }
 
  // The kernel is passed the permutation of the entity it is evaluated
  // on, so that the evaluation points are ordered consistently with a
  // neighbouring cell's. The points are given on the reference cell of
  // the entity, so `Xshape[1]` is the entity dimension. The mesh has a
  // single cell type, see fem::impl::check_expression_entities.
  md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>> perms;
  if constexpr (std::remove_cvref_t<decltype(entities)>::rank() == 2)
  {
    perms = mesh::entity_permutations(*mesh.topology_mutable(), Xshape[1], 0);
  }
 
  tabulate_expression(values, fn, Xshape, value_size, num_argument_dofs,
                      mesh.geometry().dofmaps().front(), mesh.geometry().x(),
                      coeffs, constants, entities, cell_info, argument_cells,
                      post_dof_transform, perms);
}
} // namespace dolfinx::fem::impl