CoordinateElement.h

// Copyright (C) 2018-2026 Garth N. Wells, Chris N. Richardson 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 "ElementDofLayout.h"
#include <algorithm>
#include <array>
#include <basix/element-families.h>
#include <basix/mdspan.hpp>
#include <cmath>
#include <concepts>
#include <cstdint>
#include <dolfinx/common/math.h>
#include <dolfinx/mesh/cell_types.h>
#include <limits>
#include <memory>
#include <span>
 
namespace basix
{
template <std::floating_point T>
class FiniteElement;
}
 
namespace dolfinx::fem
{
template <std::floating_point T>
class CoordinateElement
{
public:
  explicit CoordinateElement(
      std::shared_ptr<const basix::FiniteElement<T>> element);

  explicit CoordinateElement(mesh::CellType celltype, int degree,
                             basix::element::lagrange_variant type
                             = basix::element::lagrange_variant::unset,
                             bool discontinuous = false);

  ~CoordinateElement() = default;

  mesh::CellType cell_shape() const;

  int degree() const;

  int dim() const;

  basix::element::lagrange_variant variant() const;

  std::uint64_t hash() const;

  std::array<std::size_t, 4> tabulate_shape(std::size_t nd,
                                            std::size_t num_points) const;

  void tabulate(int nd, std::span<const T> X, std::array<std::size_t, 2> shape,
                std::span<T> basis) const;

  void permute_subentity_closure(std::span<std::int32_t> d,
                                 std::uint32_t cell_info,
                                 mesh::CellType entity_type,
                                 int entity_index) const;

  template <typename U, typename V, typename W>
  static void compute_jacobian(const U& dphi, const V& cell_geometry, W&& J)
  {
    math::dot(cell_geometry, dphi, J, true);
  }

  template <typename U, typename V>
  static void compute_jacobian_inverse(const U& J, V&& K)
  {
    int gdim = J.extent(0);
    int tdim = K.extent(0);
    if (gdim == tdim)
      math::inv(J, K);
    else
      math::pinv(J, K);
  }

  template <typename U>
  static double
  compute_jacobian_determinant(const U& J, std::span<typename U::value_type> w)
  {
    static_assert(U::rank() == 2, "Must be rank 2");
    if (J.extent(0) == J.extent(1))
      return math::det(J);
    else
    {
      assert(w.size() >= 2 * J.extent(0) * J.extent(1));
      using X = typename U::element_type;
      using mdspan2_t = md::mdspan<X, md::dextents<std::size_t, 2>>;
      mdspan2_t B(w.data(), J.extent(1), J.extent(0));
      mdspan2_t BA(w.data() + J.extent(0) * J.extent(1), B.extent(0),
                   J.extent(1));
      for (std::size_t i = 0; i < B.extent(0); ++i)
        for (std::size_t j = 0; j < B.extent(1); ++j)
          B(i, j) = J(j, i);
 
      // Zero working memory of BA
      std::fill_n(BA.data_handle(), BA.size(), 0);
      math::dot(B, J, BA);
      return std::sqrt(math::det(BA));
    }
  }

  ElementDofLayout create_dof_layout() const;

  template <typename U, typename V, typename W>
  static void push_forward(U&& x, const V& cell_geometry, const W& phi)
  {
    for (std::size_t i = 0; i < x.extent(0); ++i)
      for (std::size_t j = 0; j < x.extent(1); ++j)
        x(i, j) = 0;
 
    // Compute x = phi * cell_geometry;
    math::dot(phi, cell_geometry, x);
  }

  template <typename U, typename V, typename W>
  static void pull_back_affine(U&& X, const V& K, std::array<T, 3> x0,
                               const W& x)
  {
    assert(X.extent(0) == x.extent(0));
    assert(X.extent(1) == K.extent(0));
    assert(x.extent(1) == K.extent(1));
    for (std::size_t i = 0; i < X.extent(0); ++i)
      for (std::size_t j = 0; j < X.extent(1); ++j)
        X(i, j) = 0;
 
    // Calculate X for each point
    for (std::size_t p = 0; p < x.extent(0); ++p)
      for (std::size_t i = 0; i < K.extent(0); ++i)
        for (std::size_t j = 0; j < K.extent(1); ++j)
          X(p, i) += K(i, j) * (x(p, j) - x0[j]);
  }

  template <typename X>
  using mdspan2_t = md::mdspan<X, md::dextents<std::size_t, 2>>;

  void pull_back_nonaffine(mdspan2_t<T> X, mdspan2_t<const T> x,
                           mdspan2_t<const T> cell_geometry,
                           std::span<T> working_array, double tol,
                           int maxit) const;

  std::size_t pull_back_working_size(std::size_t gdim) const;

  void permute(std::span<std::int32_t> dofs, std::uint32_t cell_perm) const;

  void permute_inv(std::span<std::int32_t> dofs, std::uint32_t cell_perm) const;

  bool needs_dof_permutations() const;

  bool is_affine() const noexcept { return _is_affine; }

  bool is_discontinuous() const;
 
private:
  // Flag denoting affine map
  bool _is_affine;
 
  // Basix Element
  std::shared_ptr<const basix::FiniteElement<T>> _element;
};
} // namespace dolfinx::fem