FunctionSpace.h

// Copyright (C) 2008-2023 Anders Logg and Garth N. Wells
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include "CoordinateElement.h"
#include "DofMap.h"
#include "FiniteElement.h"
#include <boost/uuid/uuid.hpp>
#include <boost/uuid/uuid_generators.hpp>
#include <concepts>
#include <cstddef>
#include <cstdint>
#include <dolfinx/common/IndexMap.h>
#include <dolfinx/mesh/Geometry.h>
#include <dolfinx/mesh/Mesh.h>
#include <dolfinx/mesh/Topology.h>
#include <map>
#include <memory>
#include <vector>
 
namespace dolfinx::fem
{
template <std::floating_point T>
class FunctionSpace
{
public:
  using geometry_type = T;
 
  FunctionSpace(std::shared_ptr<const mesh::Mesh<geometry_type>> mesh,
                std::shared_ptr<const FiniteElement<geometry_type>> element,
                std::shared_ptr<const DofMap> dofmap,
                std::vector<std::size_t> value_shape)
      : _mesh(mesh), _element(element), _dofmap(dofmap),
        _id(boost::uuids::random_generator()()), _root_space_id(_id),
        _value_shape(value_shape)
  {
    // Do nothing
  }
 
  // Copy constructor (deleted)
  FunctionSpace(const FunctionSpace& V) = delete;
 
  FunctionSpace(FunctionSpace&& V) = default;
 
  virtual ~FunctionSpace() = default;
 
  // Assignment operator (delete)
  FunctionSpace& operator=(const FunctionSpace& V) = delete;
 
  FunctionSpace& operator=(FunctionSpace&& V) = default;
 
  FunctionSpace sub(const std::vector<int>& component) const
  {
    assert(_mesh);
    assert(_element);
    assert(_dofmap);
 
    // Check that component is valid
    if (component.empty())
      throw std::runtime_error("Component must be non-empty");
 
    // Extract sub-element
    auto element = this->_element->extract_sub_element(component);
 
    // Extract sub dofmap
    auto dofmap
        = std::make_shared<DofMap>(_dofmap->extract_sub_dofmap(component));
 
    // Create new sub space
    FunctionSpace sub_space(_mesh, element, dofmap,
                            compute_value_shape(element,
                                                _mesh->topology()->dim(),
                                                _mesh->geometry().dim()));
 
    // Set root space id and component w.r.t. root
    sub_space._root_space_id = _root_space_id;
    sub_space._component = _component;
    sub_space._component.insert(sub_space._component.end(), component.begin(),
                                component.end());
    return sub_space;
  }
 
  bool contains(const FunctionSpace& V) const
  {
    if (this == std::addressof(V))
    {
      // Spaces are the same (same memory address)
      return true;
    }
    else if (_root_space_id != V._root_space_id)
    {
      // Different root spaces
      return false;
    }
    else if (_component.size() > V._component.size())
    {
      // V is a superspace of *this
      return false;
    }
    else if (!std::equal(_component.begin(), _component.end(),
                         V._component.begin()))
    {
      // Components of 'this' are not the same as the leading components
      // of V
      return false;
    }
    else
    {
      // Ok, V is really our subspace
      return true;
    }
  }
 
  std::pair<FunctionSpace, std::vector<std::int32_t>> collapse() const
  {
    if (_component.empty())
      throw std::runtime_error("Function space is not a subspace");
 
    // Create collapsed DofMap
    auto [_collapsed_dofmap, collapsed_dofs]
        = _dofmap->collapse(_mesh->comm(), *_mesh->topology());
    auto collapsed_dofmap
        = std::make_shared<DofMap>(std::move(_collapsed_dofmap));
 
    return {
        FunctionSpace(_mesh, _element, collapsed_dofmap,
                      compute_value_shape(_element, _mesh->topology()->dim(),
                                          _mesh->geometry().dim())),
        std::move(collapsed_dofs)};
  }
 
  std::vector<int> component() const { return _component; }
 
  bool symmetric() const
  {
    if (_element)
      return _element->symmetric();
    return false;
  }
 
  std::vector<geometry_type> tabulate_dof_coordinates(bool transpose) const
  {
    if (!_component.empty())
    {
      throw std::runtime_error("Cannot tabulate coordinates for a "
                               "FunctionSpace that is a subspace.");
    }
 
    assert(_element);
    if (_element->is_mixed())
    {
      throw std::runtime_error(
          "Cannot tabulate coordinates for a mixed FunctionSpace.");
    }
 
    // Geometric dimension
    assert(_mesh);
    assert(_element);
    const std::size_t gdim = _mesh->geometry().dim();
    const int tdim = _mesh->topology()->dim();
 
    // Get dofmap local size
    assert(_dofmap);
    std::shared_ptr<const common::IndexMap> index_map = _dofmap->index_map;
    assert(index_map);
    const int index_map_bs = _dofmap->index_map_bs();
    const int dofmap_bs = _dofmap->bs();
 
    const int element_block_size = _element->block_size();
    const std::size_t scalar_dofs
        = _element->space_dimension() / element_block_size;
    const std::int32_t num_dofs
        = index_map_bs * (index_map->size_local() + index_map->num_ghosts())
          / dofmap_bs;
 
    // Get the dof coordinates on the reference element
    if (!_element->interpolation_ident())
    {
      throw std::runtime_error("Cannot evaluate dof coordinates - this element "
                               "does not have pointwise evaluation.");
    }
    const auto [X, Xshape] = _element->interpolation_points();
 
    // Get coordinate map
    const CoordinateElement<geometry_type>& cmap = _mesh->geometry().cmap();
 
    // Prepare cell geometry
    auto x_dofmap = _mesh->geometry().dofmap();
    const std::size_t num_dofs_g = cmap.dim();
    std::span<const geometry_type> x_g = _mesh->geometry().x();
 
    // Array to hold coordinates to return
    const std::size_t shape_c0 = transpose ? 3 : num_dofs;
    const std::size_t shape_c1 = transpose ? num_dofs : 3;
    std::vector<geometry_type> coords(shape_c0 * shape_c1, 0);
 
    using mdspan2_t = MDSPAN_IMPL_STANDARD_NAMESPACE::mdspan<
        geometry_type,
        MDSPAN_IMPL_STANDARD_NAMESPACE::dextents<std::size_t, 2>>;
    using cmdspan4_t = MDSPAN_IMPL_STANDARD_NAMESPACE::mdspan<
        const geometry_type,
        MDSPAN_IMPL_STANDARD_NAMESPACE::dextents<std::size_t, 4>>;
 
    // Loop over cells and tabulate dofs
    std::vector<geometry_type> x_b(scalar_dofs * gdim);
    mdspan2_t x(x_b.data(), scalar_dofs, gdim);
 
    std::vector<geometry_type> coordinate_dofs_b(num_dofs_g * gdim);
    mdspan2_t coordinate_dofs(coordinate_dofs_b.data(), num_dofs_g, gdim);
 
    auto map = _mesh->topology()->index_map(tdim);
    assert(map);
    const int num_cells = map->size_local() + map->num_ghosts();
 
    std::span<const std::uint32_t> cell_info;
    if (_element->needs_dof_transformations())
    {
      _mesh->topology_mutable()->create_entity_permutations();
      cell_info = std::span(_mesh->topology()->get_cell_permutation_info());
    }
 
    const std::array<std::size_t, 4> phi_shape
        = cmap.tabulate_shape(0, Xshape[0]);
    std::vector<geometry_type> phi_b(
        std::reduce(phi_shape.begin(), phi_shape.end(), 1, std::multiplies{}));
    cmdspan4_t phi_full(phi_b.data(), phi_shape);
    cmap.tabulate(0, X, Xshape, phi_b);
    auto phi = MDSPAN_IMPL_STANDARD_NAMESPACE::submdspan(
        phi_full, 0, MDSPAN_IMPL_STANDARD_NAMESPACE::full_extent,
        MDSPAN_IMPL_STANDARD_NAMESPACE::full_extent, 0);
 
    // TODO: Check transform
    // Basis function reference-to-conforming transformation function
    auto apply_dof_transformation
        = _element->template dof_transformation_fn<geometry_type>(
            doftransform::standard);
 
    for (int c = 0; c < num_cells; ++c)
    {
      // Extract cell geometry 'dofs'
      auto x_dofs = MDSPAN_IMPL_STANDARD_NAMESPACE::submdspan(
          x_dofmap, c, MDSPAN_IMPL_STANDARD_NAMESPACE::full_extent);
      for (std::size_t i = 0; i < x_dofs.size(); ++i)
        for (std::size_t j = 0; j < gdim; ++j)
          coordinate_dofs(i, j) = x_g[3 * x_dofs[i] + j];
 
      // Tabulate dof coordinates on cell
      cmap.push_forward(x, coordinate_dofs, phi);
      apply_dof_transformation(
          x_b, std::span(cell_info.data(), cell_info.size()), c, x.extent(1));
 
      // Get cell dofmap
      auto dofs = _dofmap->cell_dofs(c);
 
      // Copy dof coordinates into vector
      if (!transpose)
      {
        for (std::size_t i = 0; i < dofs.size(); ++i)
          for (std::size_t j = 0; j < gdim; ++j)
            coords[dofs[i] * 3 + j] = x(i, j);
      }
      else
      {
        for (std::size_t i = 0; i < dofs.size(); ++i)
          for (std::size_t j = 0; j < gdim; ++j)
            coords[j * num_dofs + dofs[i]] = x(i, j);
      }
    }
 
    return coords;
  }
 
  std::shared_ptr<const mesh::Mesh<geometry_type>> mesh() const
  {
    return _mesh;
  }
 
  std::shared_ptr<const FiniteElement<geometry_type>> element() const
  {
    return _element;
  }
 
  std::shared_ptr<const DofMap> dofmap() const { return _dofmap; }
 
  std::span<const std::size_t> value_shape() const noexcept
  {
    return _value_shape;
  }
 
  int value_size() const
  {
    return std::accumulate(_value_shape.begin(), _value_shape.end(), 1,
                           std::multiplies{});
  }
 
private:
  // The mesh
  std::shared_ptr<const mesh::Mesh<geometry_type>> _mesh;
 
  // The finite element
  std::shared_ptr<const FiniteElement<geometry_type>> _element;
 
  // The dofmap
  std::shared_ptr<const DofMap> _dofmap;
 
  // The component w.r.t. to root space
  std::vector<int> _component;
 
  // Unique identifier for the space and for its root space
  boost::uuids::uuid _id;
  boost::uuids::uuid _root_space_id;
 
  std::vector<std::size_t> _value_shape;
};
 
template <dolfinx::scalar T>
std::array<std::vector<std::shared_ptr<const FunctionSpace<T>>>, 2>
common_function_spaces(
    const std::vector<
        std::vector<std::array<std::shared_ptr<const FunctionSpace<T>>, 2>>>& V)
{
  assert(!V.empty());
  std::vector<std::shared_ptr<const FunctionSpace<T>>> spaces0(V.size(),
                                                               nullptr);
  std::vector<std::shared_ptr<const FunctionSpace<T>>> spaces1(V.front().size(),
                                                               nullptr);
 
  // Loop over rows
  for (std::size_t i = 0; i < V.size(); ++i)
  {
    // Loop over columns
    for (std::size_t j = 0; j < V[i].size(); ++j)
    {
      auto& V0 = V[i][j][0];
      auto& V1 = V[i][j][1];
      if (V0 and V1)
      {
        if (!spaces0[i])
          spaces0[i] = V0;
        else
        {
          if (spaces0[i] != V0)
            throw std::runtime_error("Mismatched test space for row.");
        }
 
        if (!spaces1[j])
          spaces1[j] = V1;
        else
        {
          if (spaces1[j] != V1)
            throw std::runtime_error("Mismatched trial space for column.");
        }
      }
    }
  }
 
  // Check that there are no null entries
  if (std::find(spaces0.begin(), spaces0.end(), nullptr) != spaces0.end())
    throw std::runtime_error("Could not deduce all block test spaces.");
  if (std::find(spaces1.begin(), spaces1.end(), nullptr) != spaces1.end())
    throw std::runtime_error("Could not deduce all block trial spaces.");
 
  return {spaces0, spaces1};
}
 
template <std::floating_point T>
std::vector<std::size_t> compute_value_shape(
    std::shared_ptr<const dolfinx::fem::FiniteElement<T>> element,
    std::size_t tdim, std::size_t gdim)
{
  auto rvs = element->reference_value_shape();
  std::vector<std::size_t> value_shape(rvs.size());
  if (element->block_size() > 1)
  {
    for (std::size_t i = 0; i < rvs.size(); ++i)
    {
      value_shape[i] = rvs[i];
    }
  }
  else
  {
    for (std::size_t i = 0; i < rvs.size(); ++i)
    {
      if (rvs[i] == tdim)
        value_shape[i] = gdim;
      else
        value_shape[i] = rvs[i];
    }
  }
  return value_shape;
}
 
template <typename U, typename V, typename W, typename X>
FunctionSpace(U mesh, V element, W dofmap, X value_shape)
    -> FunctionSpace<typename std::remove_cvref<
        typename U::element_type>::type::geometry_type::value_type>;
 
} // namespace dolfinx::fem