assemble_expression_impl.h

// Copyright (C) 2025 Garth N. Wells
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include "Expression.h"
#include "FunctionSpace.h"
#include "traits.h"
#include "utils.h"
#include <algorithm>
#include <basix/mdspan.hpp>
#include <dolfinx/common/IndexMap.h>
#include <dolfinx/mesh/Geometry.h>
#include <dolfinx/mesh/Mesh.h>
#include <dolfinx/mesh/Topology.h>
#include <memory>
#include <vector>
 
namespace dolfinx::fem::impl
{
template <dolfinx::scalar T, std::floating_point U>
void tabulate_expression(
    std::span<T> values, fem::FEkernel<T> 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 scalar_value_t<T>> 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,
    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));
 
  // 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();
  for (std::size_t e = 0; e < entities.extent(0); ++e)
  {
    std::ranges::fill(values_local, 0);
    if constexpr (entities.rank() == 1)
    {
      std::int32_t entity = entities(e);
      auto x_dofs = md::submdspan(x_dofmap, entity, md::full_extent);
      for (std::size_t i = 0; i < x_dofs.size(); ++i)
      {
        std::copy_n(std::next(x.begin(), 3 * x_dofs[i]), 3,
                    std::next(coord_dofs.begin(), 3 * i));
      }
      fn(values_local.data(), &coeffs(e, 0), constants.data(),
         coord_dofs.data(), nullptr, nullptr, nullptr);
 
      P0(values_local, cell_info, entity, 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);
      auto x_dofs = md::submdspan(x_dofmap, entity, md::full_extent);
      for (std::size_t i = 0; i < x_dofs.size(); ++i)
      {
        std::copy_n(std::next(x.begin(), 3 * x_dofs[i]), 3,
                    std::next(coord_dofs.begin(), 3 * i));
      }
      fn(values_local.data(), &coeffs(e, 0), constants.data(),
         coord_dofs.data(), &local_entity, &perm, nullptr);
      P0(values_local, cell_info, entity, 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, fem::FEkernel<T> 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::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::shared_ptr<const mesh::Topology> topology = mesh.topology();
  assert(topology);
  std::size_t num_argument_dofs = 1;
  std::span<const std::uint32_t> cell_info;
  if (element)
  {
    num_argument_dofs = element->second;
    if (element->first.get().needs_dof_transformations())
    {
      mesh.topology_mutable()->create_entity_permutations();
      cell_info = std::span(topology->get_cell_permutation_info());
      post_dof_transform
          = element->first.get().template dof_transformation_right_fn<T>(
              doftransform::transpose);
    }
  }
  // An expression has no notion of requiring a facet permutation.
  md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>> facet_perms;
  if constexpr (entities.rank() == 2)
  {
    mesh::CellType cell_type = mesh.topology()->cell_types()[0];
    int num_facets_per_cell
        = mesh::cell_num_entities(cell_type, mesh.topology()->dim() - 1);
    mesh.topology_mutable()->create_entity_permutations();
    const std::vector<std::uint8_t>& p
        = mesh.topology()->get_facet_permutations();
    facet_perms = md::mdspan(p.data(), p.size() / num_facets_per_cell,
                             num_facets_per_cell);
  }
  tabulate_expression<T, U>(values, fn, Xshape, value_size, num_argument_dofs,
                            mesh.geometry().dofmaps().front(),
                            mesh.geometry().x(), coeffs, constants, entities,
                            cell_info, post_dof_transform, facet_perms);
}
} // namespace dolfinx::fem::impl