DiscreteLocalFunction.inc.hpp

#pragma once

#include <amdis/Output.hpp>
#include <amdis/algorithm/ForEach.hpp>
#include <amdis/common/DerivativeTraits.hpp>
#include <amdis/common/FieldMatVec.hpp>
#include <amdis/common/Index.hpp>
#include <amdis/common/Math.hpp>
#include <amdis/functions/HierarchicNodeToRangeMap.hpp>
#include <amdis/functions/NodeIndices.hpp>
#include <amdis/functions/Order.hpp>
#include <amdis/typetree/FiniteElementType.hpp>
#include <amdis/utility/LocalToGlobalAdapter.hpp>

#include <dune/common/ftraits.hh>

namespace AMDiS {
namespace Impl {

// specialization for containers with gather method
template <class Coeff, class LV, class LC>
auto gather(Coeff const& coeff, LV const& localView, LC& localCoeff, Dune::PriorityTag<2>)
  -> std::void_t<decltype(coeff.gather(localView, localCoeff))>
{
  coeff.gather(localView, localCoeff);
}

// implementation for containers with multi-index access
template <class Coeff, class LV, class LC,
  class MI = typename LV::MultiIndex>
auto gather(Coeff const& coeff, LV const& localView, LC& localCoeff, Dune::PriorityTag<1>)
  -> std::void_t<decltype(coeff[std::declval<MI>()])>
{
  localCoeff.resize(localView.size());
  auto it = localCoeff.begin();
  for (auto const& idx : nodeIndices(localView))
    *it++ = coeff[idx];
}

// fallback implementation
template <class Coeff, class LV, class LC>
auto gather(Coeff const& coeff, LV const& localView, LC& localCoeff, Dune::PriorityTag<0>)
  -> std::void_t<decltype(coeff[index_<0>])>
{
  auto backend = Dune::Functions::istlVectorBackend(coeff);
  localCoeff.resize(localView.size());
  auto it = localCoeff.begin();
  for (auto const& idx : nodeIndices(localView))
    *it++ = backend[idx];
}

} // end namespace Impl


template <class C, class GB, class TP, class R>
  template <class Type>
class DiscreteFunction<C const,GB,TP,R>::LocalFunction
{
  using DomainType = typename DiscreteFunction::Domain;
  using RangeType = typename DiscreteFunction::Range;

  template <class T>
  using DerivativeRange = typename DerivativeTraits<RangeType(DomainType), T>::Range;

public:
  using Domain = typename EntitySet::LocalCoordinate;
  using Range = DerivativeRange<Type>;

  enum { hasDerivative = std::is_same<Type, tag::value>::value };

private:
  using LocalView = typename GlobalBasis::LocalView;
  using Element = typename EntitySet::Element;
  using Geometry = typename Element::Geometry;
  using NodeToRangeMap = HierarchicNodeToRangeMap;

public:
  template <class LV>
  LocalFunction(std::shared_ptr<LV> localView, TP const& treePath, C const& coefficients, Type type)
    : localView_(std::move(localView))
    , treePath_(treePath)
    , coefficients_(coefficients)
    , type_(type)
  {}

  void bind(Element const& element)
  {
    localView_->bind(element);
    geometry_.emplace(element.geometry());
    Impl::gather(coefficients_, *localView_, localCoefficients_, Dune::PriorityTag<4>{});
    bound_ = true;
  }

  void unbind()
  {
    localView_->unbind();
    geometry_.reset();
    bound_ = false;
  }

  Range operator() (const Domain& local) const
  {
    assert(bound_);
    if constexpr (std::is_same_v<Type, tag::value>)
      return evaluateFunction(local);
    else if constexpr (std::is_same_v<Type, tag::gradient>)
      return evaluateJacobian(local);
    else if constexpr (std::is_same_v<Type, tag::partial>)
      return evaluatePartialDerivative(local);
    else if constexpr (std::is_same_v<Type, tag::divergence>) {
      static constexpr bool isVectorNode
        = SubTree::isPower && int(SubTree::degree()) == GridView::dimensionworld;
      static_assert(isVectorNode);
      if constexpr (isVectorNode)
        return evaluateDivergence(local);
    }

    return Range{};
  }

  template <class DerType>
  LocalFunction<DerType> makeDerivative(DerType type) const
  {
    static_assert(std::is_same_v<Type, tag::value>);
    return {localView_, treePath_, coefficients_, type};
  }

  auto order() const
    -> decltype(AMDiS::order(std::declval<SubTree>()))
  {
    assert(bound_);
    return AMDiS::order(subTree());
  }

  Element const& localContext() const
  {
    assert(bound_);
    return localView_->element();
  }

private:

  template <class LocalCoordinate>
  auto evaluateFunction (const LocalCoordinate& local) const;

  template <class LocalCoordinate>
  auto evaluateJacobian (const LocalCoordinate& local) const;

  template <class LocalCoordinate>
  auto evaluateDivergence (const LocalCoordinate& local) const;

  template <class LocalCoordinate>
  auto evaluatePartialDerivative (const LocalCoordinate& local) const;

  // get the geometry of the bound element, that is constructed in the bind() method
  Geometry const& geometry() const
  {
    assert(bound_);
    return *geometry_;
  }

  // return the sub-tree of the basis-tree associated to the tree-path
  SubTree const& subTree() const
  {
    return Dune::TypeTree::child(localView_->tree(), treePath_);
  }

  // return the sub-tree of the basis-tree-cache associated to the tree-path
  // NOTE: this is only return in case the local-view does provide a treeCache() function
  template <class LV,
    class = decltype(std::declval<LV>().treeCache())>
  auto const& subTreeCache(Dune::PriorityTag<1>) const
  {
    return Dune::TypeTree::child(localView_->treeCache(), treePath_);
  }

  // construct a new tree-cache that stores a reference to the sub-tree
  template <class>
  auto subTreeCache(Dune::PriorityTag<0>) const
  {
    return makeNodeCache(subTree());
  }

  decltype(auto) subTreeCache() const
  {
    return subTreeCache<LocalView>(Dune::PriorityTag<2>());
  }

private:
  std::shared_ptr<LocalView> localView_;
  TP treePath_;
  Coefficients const& coefficients_;
  Type type_;
  NodeToRangeMap nodeToRangeMap_;

  std::optional<Geometry> geometry_;
  std::vector<Coefficient> localCoefficients_;
  bool bound_ = false;
};


template <class C, class GB, class TP, class R>
  template <class Type>
    template <class LocalCoordinate>
auto DiscreteFunction<C const,GB,TP,R>::LocalFunction<Type>
::evaluateFunction(LocalCoordinate const& local) const
{
  Range y;
  Recursive::forEach(y, [](auto& v) { v = 0; });

  Traversal::forEachLeafNode(this->subTreeCache(), [&](auto const& node, auto const& tp)
  {
    auto const& shapeFunctionValues = node.localBasisValuesAt(local);
    std::size_t size = node.finiteElement().size();

    // Get range entry associated to this node
    auto re = Dune::Functions::flatVectorView(nodeToRangeMap_(node, tp, y));

    for (std::size_t i = 0; i < size; ++i) {
      // Get coefficient associated to i-th shape function
      auto c = Dune::Functions::flatVectorView(localCoefficients_[node.localIndex(i)]);

      // Get value of i-th shape function
      auto v = Dune::Functions::flatVectorView(shapeFunctionValues[i]);

      std::size_t dimC = c.size();
      std::size_t dimV = v.size();
      assert(dimC*dimV == std::size_t(re.size()));
      for(std::size_t j = 0; j < dimC; ++j) {
        auto&& c_j = c[j];
        for(std::size_t k = 0; k < dimV; ++k)
          re[j*dimV + k] += c_j*v[k];
      }
    }
  });

  return y;
}


template <class C, class GB, class TP, class R>
  template <class Type>
    template <class LocalCoordinate>
auto DiscreteFunction<C const,GB,TP,R>::LocalFunction<Type>
::evaluateJacobian(LocalCoordinate const& local) const
{
  Range dy;
  Recursive::forEach(dy, [](auto& v) { v = 0; });

  Traversal::forEachLeafNode(this->subTreeCache(), [&](auto const& node, auto const& tp)
  {
    LocalToGlobalBasisAdapter localBasis(node, this->geometry());
    auto const& gradients = localBasis.gradientsAt(local);

    // Get range entry associated to this node
    auto re = Dune::Functions::flatVectorView(nodeToRangeMap_(node, tp, dy));

    for (std::size_t i = 0; i < localBasis.size(); ++i) {
      // Get coefficient associated to i-th shape function
      auto c = Dune::Functions::flatVectorView(localCoefficients_[node.localIndex(i)]);

      // Get value of i-th transformed reference gradient
      auto grad = Dune::Functions::flatVectorView(gradients[i]);

      std::size_t dimC = c.size();
      std::size_t dimV = grad.size();
      assert(dimC*dimV == std::size_t(re.size()));
      for(std::size_t j = 0; j < dimC; ++j) {
        auto&& c_j = c[j];
        for(std::size_t k = 0; k < dimV; ++k)
          re[j*dimV + k] += c_j*grad[k];
      }
    }
  });

  return dy;
}


template <class C, class GB, class TP, class R>
  template <class Type>
    template <class LocalCoordinate>
auto DiscreteFunction<C const,GB,TP,R>::LocalFunction<Type>
::evaluateDivergence(LocalCoordinate const& local) const
{
  static_assert(static_size_v<Range> == 1, "Divergence implemented for scalar coefficients only.");

  Range dy;
  Recursive::forEach(dy, [](auto& v) { v = 0; });

  auto&& node = this->subTreeCache();

  LocalToGlobalBasisAdapter localBasis(node.child(0), this->geometry());
  auto const& gradients = localBasis.gradientsAt(local);

  auto re = Dune::Functions::flatVectorView(dy);
  assert(int(re.size()) == 1);
  for (std::size_t i = 0; i < localBasis.size(); ++i) {
    auto grad = Dune::Functions::flatVectorView(gradients[i]);

    assert(int(grad.size()) == GridView::dimensionworld);
    for (std::size_t j = 0; j < GridView::dimensionworld; ++j)
      re[0] += localCoefficients_[node.child(j).localIndex(i)] * grad[j];
  }

  return dy;
}


template <class C, class GB, class TP, class R>
  template <class Type>
    template <class LocalCoordinate>
auto DiscreteFunction<C const,GB,TP,R>::LocalFunction<Type>
::evaluatePartialDerivative(LocalCoordinate const& local) const
{
  std::size_t comp = this->type_.comp;

  Range dy;
  Recursive::forEach(dy, [](auto& v) { v = 0; });

  Traversal::forEachLeafNode(this->subTreeCache(), [&](auto const& node, auto const& tp)
  {
    LocalToGlobalBasisAdapter localBasis(node, this->geometry());
    auto const& partial = localBasis.partialsAt(local, comp);

    // Get range entry associated to this node
    auto re = Dune::Functions::flatVectorView(nodeToRangeMap_(node, tp, dy));

    for (std::size_t i = 0; i < localBasis.size(); ++i) {
      // Get coefficient associated to i-th shape function
      auto c = Dune::Functions::flatVectorView(localCoefficients_[node.localIndex(i)]);

      // Get value of i-th transformed reference partial_derivative
      auto d_comp = Dune::Functions::flatVectorView(partial[i]);

      std::size_t dimC = c.size();
      std::size_t dimV = d_comp.size();
      assert(dimC*dimV == std::size_t(re.size()));
      for(std::size_t j = 0; j < dimC; ++j) {
        auto&& c_j = c[j];
        for(std::size_t k = 0; k < dimV; ++k)
          re[j*dimV + k] += c_j*d_comp[k];
      }
    }
  });

  return dy;
}


} // end namespace AMDiS