Assembler.inc.hpp

#pragma once

#include <dune/functions/functionspacebases/subspacebasis.hh>
#include <amdis/utility/TreePath.hpp>
#include <amdis/utility/Visitor.hpp>

#include <amdis/common/Math.hpp>

namespace AMDiS {

template <class Traits>
  template <class SystemMatrixType, class SystemVectorType>
void Assembler<Traits>::assemble(
    SystemMatrixType& matrix,
    SystemVectorType& solution,
    SystemVectorType& rhs,
    bool asmMatrix, bool asmVector)
{
  // 1. init matrix and rhs vector and initialize dirichlet boundary conditions
  initMatrixVector(matrix, solution, rhs, asmMatrix, asmVector);

  auto localView = globalBasis_.localView();

  // 2. create a local matrix and vector
  std::size_t localSize = localView.maxSize();

  mtl::mat::dense2D<double> elementMatrix(localSize, localSize);
  mtl::vec::dense_vector<double> elementVector(localSize);

  // 3. traverse grid and assemble operators on the elements
  for (auto const& element : elements(globalBasis_.gridView()))
  {
    set_to_zero(elementMatrix);
    set_to_zero(elementVector);

    localView.bind(element);
    auto geometry = element.geometry();

    // traverse type-tree of global-basis
    forEachNode_(localView.tree(), [&,this](auto const& rowNode, auto rowTreePath)
    {
      auto rowBasis = Dune::Functions::subspaceBasis(globalBasis_, rowTreePath);
      auto rowLocalView = rowBasis.localView();
      rowLocalView.bind(element);

      auto& rhsOp = rhsOperators_[rowNode];
      if (rhsOp.doAssemble(asmVector) && !rhsOp.empty()) {
        rhsOp.bind(element, geometry);

        auto vecAssembler = [&](auto const& context, auto& operator_list) {
          for (auto scaled : operator_list)
            scaled.op->assemble(context, rowLocalView.tree(), elementVector);
        };

        this->assembleElementOperators(element, rhsOp, vecAssembler);
      }

      forEachNode_(localView.tree(), [&,this](auto const& colNode, auto colTreePath)
      {
        auto& matOp = matrixOperators_[rowNode][colNode];
        if (matOp.doAssemble(asmMatrix) && !matOp.empty()) {
          auto colBasis = Dune::Functions::subspaceBasis(globalBasis_, colTreePath);
          auto colLocalView = colBasis.localView();
          colLocalView.bind(element);

          matOp.bind(element, geometry);

          auto matAssembler = [&](auto const& context, auto& operator_list) {
            for (auto scaled : operator_list)
              scaled.op->assemble(context, rowLocalView.tree(), colLocalView.tree(), elementMatrix);
          };

          this->assembleElementOperators(element, matOp, matAssembler);
        }
      });
    });

    // add element-matrix to system-matrix
    for (std::size_t i = 0; i < localView.size(); ++i) {
      auto const row = localView.index(i);
      for (std::size_t j = 0; j < localView.size(); ++j) {
        if (std::abs(elementMatrix(i,j)) > threshold<double>) {
          auto const col = localView.index(j);
          matrix(row,col) += elementMatrix(i,j);
        }
      }
    }

    // add element-vector to system-vector
    for (std::size_t i = 0; i < localView.size(); ++i) {
      if (std::abs(elementVector[i]) > threshold<double>) {
        auto const idx = localView.index(i);
        rhs[idx] += elementVector[i];
      }
    }

    // unbind all operators
    forEachNode_(localView.tree(), [&,this](auto const& rowNode, auto&&) {
      rhsOperators_[rowNode].unbind();
      forEachNode_(localView.tree(), [&,this](auto const& colNode, auto&&) {
        matrixOperators_[rowNode][colNode].unbind();
      });
    });

    localView.unbind();
  }

  // 4. finish matrix insertion and apply dirichlet boundary conditions
  std::size_t nnz = finishMatrixVector(matrix, solution, rhs, asmMatrix, asmVector);

  msg("fill-in of assembled matrix: ", nnz);
}


template <class Traits>
  template <class Element, class Operators, class ElementAssembler>
void Assembler<Traits>::assembleElementOperators(
    Element const& element,
    Operators& operators,
    ElementAssembler const& localAssembler) const
{
  // assemble element operators
  localAssembler(element, operators.element);

  // assemble intersection operators
  if (!operators.intersection.empty()
      || (!operators.boundary.empty() && element.hasBoundaryIntersections()))
  {
    for (auto const& intersection : intersections(globalBasis_.gridView(), element)) {
      if (intersection.boundary())
        localAssembler(intersection, operators.boundary);
      else
        localAssembler(intersection, operators.intersection);
    }
  }
}


template <class Traits>
  template <class SystemMatrixType, class SystemVectorType>
void Assembler<Traits>::initMatrixVector(
    SystemMatrixType& matrix,
    SystemVectorType& solution,
    SystemVectorType& rhs,
    bool asmMatrix, bool asmVector) const
{
  matrix.init(asmMatrix);
  solution.compress();
  rhs.compress();
  if (asmVector)
    rhs = 0;

  auto localView = globalBasis_.localView();
  forEachNode_(localView.tree(), [&,this](auto const& rowNode, auto rowTreePath)
  {
#ifdef HAVE_EXTENDED_DUNE_FUNCTIONS
    if (rowNode.isLeaf)
      msg(globalBasis_.dimension(rowTreePath), " DOFs for Basis[", to_string(rowTreePath), "]");
#endif

    auto rowBasis = Dune::Functions::subspaceBasis(globalBasis_, rowTreePath);

    forEachNode_(localView.tree(), [&,this](auto const& colNode, auto colTreePath)
    {
      auto colBasis = Dune::Functions::subspaceBasis(globalBasis_, colTreePath);

      for (auto bc : constraints_[rowNode][colNode])
        bc->init(matrix, solution, rhs, rowBasis, colBasis);
    });
  });
  msg(globalBasis_.dimension(), " total DOFs");
}


template <class Traits>
  template <class SystemMatrixType, class SystemVectorType>
std::size_t Assembler<Traits>::finishMatrixVector(
    SystemMatrixType& matrix,
    SystemVectorType& solution,
    SystemVectorType& rhs,
    bool asmMatrix, bool asmVector) const
{
  matrix.finish();

  auto localView = globalBasis_.localView();
  forEachNode_(localView.tree(), [&,this](auto const& rowNode, auto rowTreePath)
  {
    auto rowBasis = Dune::Functions::subspaceBasis(globalBasis_, rowTreePath);
    auto& rhsOp = rhsOperators_[rowNode];
    if (rhsOp.doAssemble(asmVector))
      rhsOp.assembled = true;

    forEachNode_(localView.tree(), [&,this](auto const& colNode, auto colTreePath)
    {
      auto colBasis = Dune::Functions::subspaceBasis(globalBasis_, colTreePath);
      auto& matOp = matrixOperators_[rowNode][colNode];
      if (matOp.doAssemble(asmMatrix))
        matOp.assembled = true;

      // finish boundary condition
      for (auto bc : constraints_[rowNode][colNode])
        bc->finish(matrix, solution, rhs, rowBasis, colBasis);
    });
  });

  return matrix.nnz();
}

} // end namespace AMDiS