dofhandler.h

#ifndef INCG_LF_DOFHD_H
#define INCG_LF_DOFHD_H
/***************************************************************************
 * LehrFEM++ - A simple C++ finite element libray for teaching
 * Developed from 2018 at the Seminar of Applied Mathematics of ETH Zurich,
 * lead developers Dr. R. Casagrande and Prof. R. Hiptmair
 ***************************************************************************/

 
#include <lf/mesh/mesh.h>
 
#include <map>
 
#include "assembly_types.h"
 
namespace lf::assemble {
class DofHandler {
 protected:
  DofHandler() = default;

  DofHandler(const DofHandler &) = default;
  DofHandler(DofHandler &&) = default;
  DofHandler &operator=(const DofHandler &) = default;
  DofHandler &operator=(DofHandler &&) = default;
 public:
  virtual ~DofHandler() = default;

  [[nodiscard]] virtual size_type NumDofs() const = 0;

  [[nodiscard]] virtual size_type NumLocalDofs(
      const lf::mesh::Entity &entity) const = 0;

  [[nodiscard]] virtual size_type NumInteriorDofs(
      const lf::mesh::Entity &entity) const = 0;

  [[nodiscard]] virtual std::span<const gdof_idx_t> GlobalDofIndices(
      const lf::mesh::Entity &entity) const = 0;

  [[nodiscard]] virtual std::span<const gdof_idx_t> InteriorGlobalDofIndices(
      const lf::mesh::Entity &entity) const = 0;

  [[nodiscard]] virtual const lf::mesh::Entity &Entity(
      gdof_idx_t dofnum) const = 0;

  [[nodiscard]] virtual std::shared_ptr<const lf::mesh::Mesh> Mesh() const = 0;
};

std::ostream &operator<<(std::ostream &o, const DofHandler &dof_handler);

void PrintInfo(std::ostream &stream, const DofHandler &dof_handler,
               unsigned int ctrl = 0);
 
/* ====================================================================== */

class UniformFEDofHandler : public DofHandler {
 public:
  using dof_map_t = std::map<lf::base::RefEl, base::size_type>;

  UniformFEDofHandler(std::shared_ptr<const lf::mesh::Mesh> mesh,
                      dof_map_t dofmap, bool check_edge_orientation = true);

  UniformFEDofHandler(UniformFEDofHandler &&) = default;

  UniformFEDofHandler(const UniformFEDofHandler &) = delete;

  UniformFEDofHandler &operator=(UniformFEDofHandler &&) = default;

  UniformFEDofHandler &operator=(const UniformFEDofHandler &) = delete;

  ~UniformFEDofHandler() override = default;
 
  [[nodiscard]] size_type NumDofs() const override { return num_dof_; }

  [[nodiscard]] size_type NumLocalDofs(
      const lf::mesh::Entity &entity) const override;

  [[nodiscard]] size_type NumInteriorDofs(
      const lf::mesh::Entity &entity) const override;

  [[nodiscard]] std::span<const gdof_idx_t> GlobalDofIndices(
      const lf::mesh::Entity &entity) const override;

  [[nodiscard]] std::span<const gdof_idx_t> InteriorGlobalDofIndices(
      const lf::mesh::Entity &entity) const override;

  [[nodiscard]] const lf::mesh::Entity &Entity(
      gdof_idx_t dofnum) const override {
    LF_VERIFY_MSG(dofnum < dof_entities_.size(),
                  "Illegal dof index " << dofnum << ", max = " << num_dof_);
    return *dof_entities_[dofnum];
  }

  [[nodiscard]] std::shared_ptr<const lf::mesh::Mesh> Mesh() const override {
    return mesh_;
  }
 
 private:
  void initIndexArrays();
  void InitTotalNumDofs();
 
  // Access method to numbers and values of indices of shape functions
  [[nodiscard]] std::span<const gdof_idx_t> GlobalDofIndices(
      lf::base::RefEl ref_el_type, glb_idx_t entity_index) const;
 
  [[nodiscard]] std::span<const gdof_idx_t> InteriorGlobalDofIndices(
      lf::base::RefEl ref_el_type, glb_idx_t entity_index) const;
 
  [[nodiscard]] size_type GetNumLocalDofs(lf::base::RefEl ref_el_type,
                                          glb_idx_t /*unused*/) const {
    return NumCoveredDofs(ref_el_type);
  }

  size_type kNodeOrd = 2; 
  size_type kEdgeOrd = 1; 
  size_type kCellOrd = 0; 

  [[nodiscard]] size_type NumCoveredDofs(lf::base::RefEl ref_el_type) const {
    size_type no_covered_dofs;
    switch (ref_el_type) {
      case lf::base::RefEl::kPoint(): {
        no_covered_dofs = num_dofs_[kNodeOrd];
        break;
      }
      case lf::base::RefEl::kSegment(): {
        no_covered_dofs = num_dofs_[kEdgeOrd];
        break;
      }
      case lf::base::RefEl::kTria(): {
        no_covered_dofs = num_dofs_tria_;
        break;
      }
      case lf::base::RefEl::kQuad(): {
        no_covered_dofs = num_dofs_quad_;
        break;
      }
      default: {
        LF_VERIFY_MSG(false, "Illegal entity type");
        break;
      }
    }  // end switch
    return no_covered_dofs;
  }

  [[nodiscard]] size_type NumInteriorDofs(lf::base::RefEl ref_el_type) const {
    size_type no_loc_dofs;
    switch (ref_el_type) {
      case lf::base::RefEl::kPoint(): {
        no_loc_dofs = num_loc_dof_point_;
        break;
      }
      case lf::base::RefEl::kSegment(): {
        no_loc_dofs = num_loc_dof_segment_;
        break;
      }
      case lf::base::RefEl::kTria(): {
        no_loc_dofs = num_loc_dof_tria_;
        break;
      }
      case lf::base::RefEl::kQuad(): {
        no_loc_dofs = num_loc_dof_quad_;
        break;
      }
      default: {
        LF_VERIFY_MSG(false, "Illegal entity type");
        break;
      }
    }  // end switch
    return no_loc_dofs;
  }

  std::shared_ptr<const lf::mesh::Mesh> mesh_;
  size_type num_dof_{0};
  std::vector<const lf::mesh::Entity *> dof_entities_;
  std::array<std::vector<gdof_idx_t>, 3> dofs_;
  std::array<size_type, 3> num_dofs_;
  size_type num_dofs_tria_{0}, num_dofs_quad_{0};
  size_type num_loc_dof_point_{0};
  size_type num_loc_dof_segment_{0};
  size_type num_loc_dof_tria_{0};
  size_type num_loc_dof_quad_{0};
  bool check_edge_orientation_;
};
 
/* ====================================================================== */

class DynamicFEDofHandler : public DofHandler {
 public:
  DynamicFEDofHandler(DynamicFEDofHandler &&) = default;

  DynamicFEDofHandler(const DynamicFEDofHandler &) = delete;

  DynamicFEDofHandler &operator=(DynamicFEDofHandler &&) = default;

  DynamicFEDofHandler &operator=(const DynamicFEDofHandler &) = delete;

  ~DynamicFEDofHandler() override = default;

  template <typename LOCALDOFINFO>
  // NOLINTNEXTLINE(readability-function-cognitive-complexity)
  DynamicFEDofHandler(std::shared_ptr<const lf::mesh::Mesh> mesh_p,
                      LOCALDOFINFO &&locdof)
      : mesh_p_(std::move(mesh_p)) {
    LF_ASSERT_MSG((mesh_p_->DimMesh() == 2), "Can handle 2D meshes only");
 
    // Index counter for global shape functions = global dof
    gdof_idx_t dof_idx = 0;
 
    // Step I: Set indices for shape functions on nodes
    // Run through node indices (entities of co-dimension 2)
    const size_type no_nodes = mesh_p_->NumEntities(2);
    num_int_dofs_[2].resize(no_nodes, 0);
    offsets_[2].resize(no_nodes + 1, 0);
    // Traverse nodes (co-dimension-2 entities) based on indices
    for (glb_idx_t node_idx = 0; node_idx < no_nodes; node_idx++) {
      // Obtain pointer to node entity
      const mesh::Entity *node_p{mesh_p_->EntityByIndex(2, node_idx)};
      LF_ASSERT_MSG(mesh_p_->Index(*node_p) == node_idx, "Node index mismatch");
      // Offset for indices of node in index vector
      const glb_idx_t node_dof_offset = dof_idx;
      offsets_[2][node_idx] = node_dof_offset;
      // Request number of local shape functions associated with the node
      const size_type no_int_dof_node = locdof(*node_p);
      num_int_dofs_[2][node_idx] = no_int_dof_node;
 
      // Store dof indices in array
      for (unsigned j = 0; j < no_int_dof_node; j++) {
        dofs_[2].push_back(dof_idx);
        dof_entities_.push_back(node_p);  // Store entity for current dof
        dof_idx++;                        // Move on to next index
      }
    }
    // Set sentinel
    offsets_[2][no_nodes] = dof_idx;
 
    // Step II: Set indices for shape functions on edges (co-dimension = 1)
    const size_type no_edges = mesh_p_->NumEntities(1);
    // Set length of edge-related index vectors
    num_int_dofs_[1].resize(no_edges, 0);
    offsets_[1].resize(no_edges + 1, 0);
    // points to the position of the current dof index
    size_type edge_dof_offset = 0;
    for (glb_idx_t edge_idx = 0; edge_idx < no_edges; edge_idx++) {
      // Obtain pointer to edge entity
      const mesh::Entity *edge{mesh_p_->EntityByIndex(1, edge_idx)};
      LF_ASSERT_MSG(mesh_p_->Index(*edge) == edge_idx, "Edge index mismatch");
      // Offset for indices of edge dof in index vector
      offsets_[1][edge_idx] = edge_dof_offset;
      const size_type no_int_dof_edge = locdof(*edge);
      num_int_dofs_[1][edge_idx] = no_int_dof_edge;
 
      // Obtain indices for basis functions sitting at endpoints
      // Endpoints are mesh entities with co-dimension = 2
      for (const lf::mesh::Entity *endpoint : edge->SubEntities(1)) {
        const glb_idx_t ep_idx(mesh_p_->Index(*endpoint));
        const glb_idx_t ep_dof_offset = offsets_[2][ep_idx];
        const size_type no_int_dofs_ep = num_int_dofs_[2][ep_idx];
        // Copy indices of shape functions from nodes to edge
        for (unsigned j = 0; j < no_int_dofs_ep; j++) {
          dofs_[1].push_back(dofs_[2][ep_dof_offset + j]);
          edge_dof_offset++;
        }
      }
      // Set indices for interior edge degrees of freedom
      for (unsigned j = 0; j < no_int_dof_edge; j++) {
        dofs_[1].push_back(dof_idx);
        edge_dof_offset++;
        dof_entities_.push_back(edge);
        dof_idx++;
      }
    }
    // Set sentinel
    offsets_[1][no_edges] = edge_dof_offset;
 
    // Step III: Set indices for shape functions on cells (co-dimension = 0)
    const size_type no_cells = mesh_p_->NumEntities(0);
    // Set length of cell-related index vectors
    num_int_dofs_[0].resize(no_cells, 0);
    offsets_[0].resize(no_cells + 1, 0);
    // points to the position of the current dof index
    size_type cell_dof_offset = 0;
    for (glb_idx_t cell_idx = 0; cell_idx < no_cells; cell_idx++) {
      // Obtain pointer to current ell
      const mesh::Entity *cell{mesh_p_->EntityByIndex(0, cell_idx)};
      // Offset for indices of cell dof in index vector
      offsets_[0][cell_idx] = cell_dof_offset;
      const size_type no_int_dof_cell = locdof(*cell);
      num_int_dofs_[0][cell_idx] = no_int_dof_cell;
 
      // Obtain indices for basis functions in vertices
      for (const lf::mesh::Entity *vertex : cell->SubEntities(2)) {
        const glb_idx_t vt_idx(mesh_p_->Index(*vertex));
        const glb_idx_t vt_dof_offset = offsets_[2][vt_idx];
        const size_type no_int_dofs_vt = num_int_dofs_[2][vt_idx];
        // Copy indices of shape functions from nodes to cell
        for (unsigned j = 0; j < no_int_dofs_vt; j++) {
          dofs_[0].push_back(dofs_[2][vt_dof_offset + j]);
          cell_dof_offset++;
        }
      }
 
      // Collect indices of interior shape functions of edges
      // Internal ordering will depend on the orientation of the edge
      auto edge_orientations = cell->RelativeOrientations();
      auto edges = cell->SubEntities(1);
      // Loop over edges
      for (int ed_sub_idx = 0; ed_sub_idx < cell->RefEl().NumSubEntities(1);
           ed_sub_idx++) {
        const glb_idx_t edge_idx = mesh_p_->Index(*edges[ed_sub_idx]);
        const size_type no_int_dof_edge = num_int_dofs_[1][edge_idx];
        const glb_idx_t edge_int_dof_offset =
            offsets_[1][edge_idx + 1] - no_int_dof_edge;
 
        // Copy indices of shape functions from edges to cell
        // The order, in which they are copied depends on the relative
        // orientation of the edge w.r.t. the cell
        switch (edge_orientations[ed_sub_idx]) {
          case lf::mesh::Orientation::positive: {
            for (int j = 0; j < no_int_dof_edge; j++) {
              dofs_[0].push_back(dofs_[1][edge_int_dof_offset + j]);
              cell_dof_offset++;
            }
            break;
          }
          case lf::mesh::Orientation::negative: {
            for (int j = static_cast<int>(no_int_dof_edge - 1); j >= 0; j--) {
              dofs_[0].push_back(dofs_[1][edge_int_dof_offset + j]);
              cell_dof_offset++;
            }
            break;
          }
        }  // end switch
      }  // end loop over edges
 
      // Set indices for interior shape functions of the cell
      for (unsigned j = 0; j < no_int_dof_cell; j++) {
        dofs_[0].push_back(dof_idx);
        cell_dof_offset++;
        dof_entities_.push_back(cell);
        dof_idx++;
      }  // end loop over interior dof of cell
    }  // end loop over cells
    // Set sentinel
    offsets_[0][no_cells] = cell_dof_offset;
 
    // Finally set number of global shape functions
    num_dof_ = dof_idx;
  }  // end constructor

  [[nodiscard]] size_type NumDofs() const override { return num_dof_; }

  [[nodiscard]] size_type NumInteriorDofs(
      const lf::mesh::Entity &entity) const override;

  [[nodiscard]] size_type NumLocalDofs(
      const lf::mesh::Entity &entity) const override;

  [[nodiscard]] std::span<const gdof_idx_t> GlobalDofIndices(
      const lf::mesh::Entity &entity) const override;

  [[nodiscard]] std::span<const gdof_idx_t> InteriorGlobalDofIndices(
      const lf::mesh::Entity &entity) const override;

  [[nodiscard]] const lf::mesh::Entity &Entity(
      gdof_idx_t dofnum) const override {
    LF_VERIFY_MSG(dofnum < dof_entities_.size(),
                  "Illegal dof index " << dofnum << ", max = " << num_dof_);
    return *dof_entities_[dofnum];
  }

  [[nodiscard]] std::shared_ptr<const lf::mesh::Mesh> Mesh() const override {
    return mesh_p_;
  }
 
 private:
  std::shared_ptr<const lf::mesh::Mesh> mesh_p_;
  size_type num_dof_{0};
  std::vector<const lf::mesh::Entity *> dof_entities_;
  std::array<std::vector<size_type>, 3> num_int_dofs_;
  std::array<std::vector<size_type>, 3> offsets_;
  std::array<std::vector<gdof_idx_t>, 3> dofs_;
};
 
}  // namespace lf::assemble


template <>
struct fmt::formatter<lf::assemble::DofHandler> : ostream_formatter {};

#endif