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)
      : _mesh(mesh), _elements{element}, _dofmaps{std::move(dofmap)},
        _id(boost::uuids::random_generator()()), _root_space_id(_id)
  {
    // Do nothing
  }

  FunctionSpace(
      std::shared_ptr<const mesh::Mesh<geometry_type>> mesh,
      std::vector<std::shared_ptr<const FiniteElement<geometry_type>>> elements,
      std::vector<std::shared_ptr<const DofMap>> dofmaps)
      : _mesh(mesh), _elements(elements), _dofmaps(std::move(dofmaps)),
        _id(boost::uuids::random_generator()()), _root_space_id(_id)
  {
    std::vector<mesh::CellType> cell_types = mesh->topology()->cell_types();
    std::size_t num_cell_types = cell_types.size();
    if (elements.size() != num_cell_types)
    {
      throw std::runtime_error(
          "Number of elements must match number of cell types");
    }
    if (_dofmaps.size() != num_cell_types)
    {
      throw std::runtime_error(
          "Number of dofmaps must match number of cell types");
    }
    for (std::size_t i = 0; i < num_cell_types; ++i)
    {
      if (elements.at(i)->cell_type() != cell_types.at(i))
        throw std::runtime_error(
            "Element cell types must match mesh cell types");
    }
  }
 
  // 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(_elements.size() > 0);
    assert(_dofmaps.size() > 0);
 
    // Check that component is valid
    if (component.empty())
      throw std::runtime_error("Component must be non-empty");
 
    // Extract sub-element
    std::vector<std::shared_ptr<const FiniteElement<geometry_type>>>
        sub_elements;
    for (auto e : _elements)
    {
      assert(e);
      sub_elements.push_back(e->extract_sub_element(component));
    }
    // Extract sub dofmap
    std::vector<std::shared_ptr<const DofMap>> sub_dofmaps;
    for (auto d : _dofmaps)
    {
      assert(d);
      sub_dofmaps.push_back(
          std::make_shared<const DofMap>(d->extract_sub_dofmap(component)));
    }
 
    // Create new sub space
    FunctionSpace sub_space(_mesh, sub_elements, sub_dofmaps);
 
    // 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::vector<std::int32_t>>>
  collapse() const
  {
    spdlog::debug("FunctionSpace::collapse");
    if (_component.empty())
      throw std::runtime_error("Function space is not a subspace");
 
    // Create collapsed DofMap
    std::vector<std::shared_ptr<const DofMap>> collapsed_dofmaps;
    std::vector<std::vector<std::int32_t>> collapsed_dofs;
    for (auto d : _dofmaps)
    {
      assert(d);
      spdlog::debug("Call DofMap::collapse");
      auto [_collapsed_dofmap, _collapsed_dofs]
          = d->collapse(_mesh->comm(), *_mesh->topology());
      spdlog::debug("Got collapsed dofmap");
      collapsed_dofmaps.push_back(
          std::make_shared<DofMap>(std::move(_collapsed_dofmap)));
      collapsed_dofs.push_back(_collapsed_dofs);
    }
 
    return {FunctionSpace(_mesh, _elements, collapsed_dofmaps),
            std::move(collapsed_dofs)};
  }

  std::vector<int> component() const { return _component; }

  bool symmetric() const
  {
    if (_elements.front())
      return _elements.front()->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(_elements.front());
    if (_elements.front()->is_mixed())
    {
      throw std::runtime_error(
          "Cannot tabulate coordinates for a mixed FunctionSpace.");
    }
 
    // Geometric dimension
    assert(_mesh);
    assert(_elements.front());
    const std::size_t gdim = _mesh->geometry().dim();
    const int tdim = _mesh->topology()->dim();
 
    // Get dofmap local size
    assert(_dofmaps.front());
    std::shared_ptr<const common::IndexMap> index_map
        = _dofmaps.front()->index_map;
    assert(index_map);
    const int index_map_bs = _dofmaps.front()->index_map_bs();
    const int dofmap_bs = _dofmaps.front()->bs();
 
    const int element_block_size = _elements.front()->block_size();
    const std::size_t scalar_dofs
        = _elements.front()->space_dimension() / element_block_size;
    const std::int32_t num_dofs
        = index_map_bs * (index_map->size_local() + index_map->num_ghosts())
          / dofmap_bs;
 
    std::size_t num_cell_types = _elements.size();
 
    // 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 = md::mdspan<geometry_type, md::dextents<std::size_t, 2>>;
    using cmdspan4_t
        = md::mdspan<const geometry_type, md::dextents<std::size_t, 4>>;
    std::span<const geometry_type> x_g = _mesh->geometry().x();
    std::vector<geometry_type> x_b(scalar_dofs * gdim);
    mdspan2_t x(x_b.data(), scalar_dofs, gdim);
 
    for (std::size_t i = 0; i < num_cell_types; ++i)
    {
      // Get the dof coordinates on the reference element
      if (!_elements[i]->interpolation_ident())
      {
        throw std::runtime_error(
            "Cannot evaluate dof coordinates - this element "
            "does not have pointwise evaluation.");
      }
      const auto [X, Xshape] = _elements[i]->interpolation_points();
 
      // Get coordinate map
      const CoordinateElement<geometry_type>& cmap
          = _mesh->geometry().cmaps()[i];
 
      // Prepare cell geometry
      auto x_dofmap = _mesh->geometry().dofmap(i);
      const std::size_t num_dofs_g = cmap.dim();
      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_maps(tdim)[i];
      assert(map);
      const int num_cells = map->size_local() + map->num_ghosts();
 
      std::span<const std::uint32_t> cell_info;
      if (_elements[i]->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
          = md::submdspan(phi_full, 0, md::full_extent, md::full_extent, 0);
 
      // TODO: Check transform
      // Basis function reference-to-conforming transformation function
      auto apply_dof_transformation
          = _elements[i]->template dof_transformation_fn<geometry_type>(
              doftransform::standard);
 
      for (int c = 0; c < num_cells; ++c)
      {
        // Extract cell geometry 'dofs'
        auto x_dofs = md::submdspan(x_dofmap, c, md::full_extent);
        for (std::size_t j = 0; j < x_dofs.size(); ++j)
          for (std::size_t k = 0; k < gdim; ++k)
            coordinate_dofs(j, k) = x_g[3 * x_dofs[j] + k];
 
        // 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 = _dofmaps[i]->cell_dofs(c);
 
        // Copy dof coordinates into vector
        if (!transpose)
        {
          for (std::size_t j = 0; j < dofs.size(); ++j)
            for (std::size_t k = 0; k < gdim; ++k)
              coords[dofs[j] * 3 + k] = x(j, k);
        }
        else
        {
          for (std::size_t j = 0; j < dofs.size(); ++j)
            for (std::size_t k = 0; k < gdim; ++k)
              coords[k * num_dofs + dofs[j]] = x(j, k);
        }
      }
    }
 
    return coords;
  }

  std::shared_ptr<const mesh::Mesh<geometry_type>> mesh() const
  {
    return _mesh;
  }

  std::shared_ptr<const FiniteElement<geometry_type>> element() const
  {
    if (_elements.size() > 1)
    {
      throw std::runtime_error(
          "FunctionSpace has multiple elements, call `elements` instead.");
    }
 
    return elements(0);
  }

  std::shared_ptr<const FiniteElement<geometry_type>>
  elements(int cell_type_idx) const
  {
    return _elements.at(cell_type_idx);
  }

  std::shared_ptr<const DofMap> dofmap() const
  {
    if (_dofmaps.size() > 1)
    {
      throw std::runtime_error(
          "FunctionSpace has multiple dofmaps, call `dofmaps` instead.");
    }
 
    return dofmaps(0);
  }

  std::shared_ptr<const DofMap> dofmaps(int cell_type_idx) const
  {
    return _dofmaps.at(cell_type_idx);
  }
 
private:
  // The mesh
  std::shared_ptr<const mesh::Mesh<geometry_type>> _mesh;
 
  // The finite element
  std::vector<std::shared_ptr<const FiniteElement<geometry_type>>> _elements;
 
  // The dofmap
  std::vector<std::shared_ptr<const DofMap>> _dofmaps;
 
  // 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;
};

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 <typename U, typename V, typename W>
FunctionSpace(U mesh, V element, W dofmap)
    -> FunctionSpace<typename std::remove_cvref<
        typename U::element_type>::type::geometry_type::value_type>;
 
} // namespace dolfinx::fem