assembler.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
 ***************************************************************************/
 
#ifndef _LF_ASSEMBLE_H
#define _LF_ASSEMBLE_H
 
#include <fmt/ranges.h>
#include <spdlog/fmt/ostr.h>
#include <spdlog/sinks/stdout_color_sinks.h>
#include <spdlog/spdlog.h>
 
#include <iostream>
 
#include "dofhandler.h"
 
namespace lf::assemble {
 
std::shared_ptr<spdlog::logger> &AssembleMatrixLogger();
 
template <typename TMPMATRIX, class ENTITY_MATRIX_PROVIDER>
void AssembleMatrixLocally(dim_t codim, const DofHandler &dof_handler_trial,
                           const DofHandler &dof_handler_test,
                           ENTITY_MATRIX_PROVIDER &entity_matrix_provider,
                           TMPMATRIX &matrix) {
  // Fetch pointer to underlying mesh
  auto mesh = dof_handler_trial.Mesh();
  LF_ASSERT_MSG(mesh == dof_handler_test.Mesh(),
                "Trial and test space must be defined on the same mesh");
  // Central assembly loop over entities of co-dimension specified by
  // the function argument codim
  for (const lf::mesh::Entity *entity : mesh->Entities(codim)) {
    // Some entities may be skipped
    if (entity_matrix_provider.isActive(*entity)) {
      // log the entity reference element + it's global index on level Debug
      SPDLOG_LOGGER_DEBUG(AssembleMatrixLogger(), "Entity {} ({})", *entity,
                          mesh->Index(*entity));
      // Size, aka number of rows and columns, of element matrix
      const size_type nrows_loc = dof_handler_test.NumLocalDofs(*entity);
      const size_type ncols_loc = dof_handler_trial.NumLocalDofs(*entity);
      // row indices of for contributions of cells
      nonstd::span<const gdof_idx_t> row_idx(
          dof_handler_test.GlobalDofIndices(*entity));
      // Column indices of for contributions of cells
      nonstd::span<const gdof_idx_t> col_idx(
          dof_handler_trial.GlobalDofIndices(*entity));
      // Request local matrix from entity_matrix_provider object. In the
      // case codim = 0, when `entity` is a cell, this is the element matrix
      const auto elem_mat{entity_matrix_provider.Eval(*entity)};
      LF_ASSERT_MSG(elem_mat.rows() >= nrows_loc,
                    "nrows mismatch " << elem_mat.rows() << " <-> " << nrows_loc
                                      << ", entity " << mesh->Index(*entity));
      LF_ASSERT_MSG(elem_mat.cols() >= ncols_loc,
                    "ncols mismatch " << elem_mat.cols() << " <-> " << nrows_loc
                                      << ", entity " << mesh->Index(*entity));
 
      // Log global row and column indices on level debug
      SPDLOG_LOGGER_DEBUG(AssembleMatrixLogger(), "row_idx = {}, col_idx = {}",
                          row_idx, col_idx);
 
      // Log shape of element matrix on level debug
      SPDLOG_LOGGER_DEBUG(AssembleMatrixLogger(), "{} x {} element matrix",
                          nrows_loc, ncols_loc);
 
      // Log element matrix itself on level trace
      // TODO(craffael): Doesn't work yet because of fmt issue:
      // https://github.com/fmtlib/fmt/issues/1885
      // SPDLOG_LOGGER_TRACE(AssembleMatrixLogger, elem_mat);
 
      // Assembly double loop
      std::stringstream ss;  // used to log all triplets on one line.
 
      for (int i = 0; i < nrows_loc; i++) {
        for (int j = 0; j < ncols_loc; j++) {
          // Add the element at position (i,j) of the local matrix
          // to the entry at (row_idx[i], col_idx[j]) of the global matrix
          matrix.AddToEntry(row_idx[i], col_idx[j], elem_mat(i, j));
 
          // log the added "triplet" on level trace:
          if (AssembleMatrixLogger()->should_log(spdlog::level::trace)) {
            // if we are on level trace, build string of all triplets:
            ss << "(" << row_idx[i] << ',' << col_idx[j]
               << ")+= " << elem_mat(i, j) << ", ";
          }
        }
      }  // end assembly local double loop
 
      // log all the triplets on one line on level trace:
      SPDLOG_LOGGER_TRACE(AssembleMatrixLogger(), ss.str());
 
    }  // end if(isActive() )
  }    // end main assembly loop
}  // end AssembleMatrixLocally
 
template <typename TMPMATRIX, class ENTITY_MATRIX_PROVIDER>
TMPMATRIX AssembleMatrixLocally(
    dim_t codim, const DofHandler &dof_handler_trial,
    const DofHandler &dof_handler_test,
    ENTITY_MATRIX_PROVIDER &entity_matrix_provider) {
  TMPMATRIX matrix{dof_handler_test.NumDofs(), dof_handler_trial.NumDofs()};
  matrix.setZero();
  AssembleMatrixLocally<TMPMATRIX, ENTITY_MATRIX_PROVIDER>(
      codim, dof_handler_trial, dof_handler_test, entity_matrix_provider,
      matrix);
  return matrix;
}
 
template <typename TMPMATRIX, class ENTITY_MATRIX_PROVIDER>
TMPMATRIX AssembleMatrixLocally(
    dim_t codim, const DofHandler &dof_handler,
    ENTITY_MATRIX_PROVIDER &entity_matrix_provider) {
  return AssembleMatrixLocally<TMPMATRIX, ENTITY_MATRIX_PROVIDER>(
      codim, dof_handler, dof_handler, entity_matrix_provider);
}  // end of group assemble_matrix_locally
 
template <typename VECTOR, class ENTITY_VECTOR_PROVIDER>
void AssembleVectorLocally(dim_t codim, const DofHandler &dof_handler,
                           ENTITY_VECTOR_PROVIDER &entity_vector_provider,
                           VECTOR &resultvector) {
  // Pointer to underlying mesh
  auto mesh = dof_handler.Mesh();
 
  // Central assembly loop over entities of the co-dimension specified via
  // the template argument CODIM
  for (const lf::mesh::Entity *entity : mesh->Entities(codim)) {
    // Some cells may be skipped
    if (entity_vector_provider.isActive(*entity)) {
      // Length of element vector
      const size_type veclen = dof_handler.NumLocalDofs(*entity);
      // global dof indices for contribution of the entity
      nonstd::span<const gdof_idx_t> dof_idx(
          dof_handler.GlobalDofIndices(*entity));
      // Request local vector from entity_vector_provider object. In the case
      // CODIM = 0, when `entity` is a cell, this is the element vector
      const auto elem_vec{entity_vector_provider.Eval(*entity)};
      LF_ASSERT_MSG(elem_vec.size() >= veclen,
                    "length mismatch " << elem_vec.size() << " <-> " << veclen
                                       << ", entity " << mesh->Index(*entity));
      // Assembly (single) loop
      for (int i = 0; i < veclen; i++) {
        resultvector[dof_idx[i]] += elem_vec[i];
      }  // end assembly localloop
    }    // end if(isActive() )
  }      // end main assembly loop
}  // end AssembleVectorLocally
 
template <typename VECTOR, class ENTITY_VECTOR_PROVIDER>
VECTOR AssembleVectorLocally(dim_t codim, const DofHandler &dof_handler,
                             ENTITY_VECTOR_PROVIDER &entity_vector_provider) {
  // Allocated vector holding r.h.s. vector to be assembled
  VECTOR resultvector{dof_handler.NumDofs()};
  // Initialize to zero: assembly of new vector
  resultvector.setZero();
  // Perform actual assembly
  AssembleVectorLocally<VECTOR, ENTITY_VECTOR_PROVIDER>(
      codim, dof_handler, entity_vector_provider, resultvector);
  return resultvector;
}  // end AssembleVectorLocally
  // end group assemble_vector_locally
 
}  // namespace lf::assemble
 
#endif