DiscreteLocalFunction.inc.hpp

#pragma once

#include <amdis/Output.hpp>
#include <amdis/common/DerivativeTraits.hpp>
#include <amdis/common/FieldMatVec.hpp>
#include <amdis/common/Math.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 of Coeff has gather method
template <class Coeff, class LocalView, class LocalCoeff,
  class = decltype(std::declval<Coeff>().gather(std::declval<LocalView>(), std::declval<LocalCoeff&>()))>
void gather(Coeff const& coeff, LocalView const& localView, LocalCoeff& localCoeff, Dune::PriorityTag<2>)
{
  coeff.gather(localView, localCoeff);
}

// fallback implementation
template <class Coeff, class LocalView, class LocalCoeff>
void gather(Coeff const& coeff, LocalView const& localView, LocalCoeff& localCoeff, Dune::PriorityTag<1>)
{
  localCoeff.resize(localView.size());
  auto it = localCoeff.begin();
  for (auto const& idx : nodeIndices(localView))
    *it++ = coeff[idx];
}

} // end namespace Impl


template <class Coeff, class GB, class TP>
class DiscreteFunction<Coeff const,GB,TP>::LocalFunction
{
public:
  using Domain = typename EntitySet::LocalCoordinate;
  using Range = typename DiscreteFunction::Range;

  enum { hasDerivative = true };

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

public:
  LocalFunction(DiscreteFunction const& globalFunction)
    : globalFunction_(globalFunction)
    , localView_(globalFunction_.basis().localView())
    , subTreeCache_(&Dune::TypeTree::child(localView_.treeCache(), globalFunction_.treePath()))
  {}

  LocalFunction(LocalFunction const& other)
    : globalFunction_(other.globalFunction_)
    , localView_(globalFunction_.basis().localView())
    , subTreeCache_(&Dune::TypeTree::child(localView_.treeCache(), globalFunction_.treePath()))
  {}

  void bind(Element const& element)
  {
    localView_.bind(element);
    Impl::gather(globalFunction_.coefficients(), localView_, localCoefficients_, Dune::PriorityTag<4>{});
    bound_ = true;
  }

  void unbind()
  {
    localView_.unbind();
    bound_ = false;
  }

  Range operator()(Domain const& x) const
  {
    assert( bound_ );
    Range y(0);

    auto&& nodeToRangeEntry = globalFunction_.nodeToRangeEntry_;
    for_each_leaf_node(*subTreeCache_, [&,this](auto const& node, auto const& tp)
    {
      auto const& shapeFunctionValues = node.localBasisValuesAt(x);
      std::size_t size = node.finiteElement().size();

      // Get range entry associated to this node
      auto re = Dune::Functions::flatVectorView(nodeToRangeEntry(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;
  }

  GradientLocalFunction makeDerivative(tag::gradient type) const
  {
    return GradientLocalFunction{globalFunction_, type};
  }

  DivergenceLocalFunction makeDerivative(tag::divergence type) const
  {
    return DivergenceLocalFunction{globalFunction_, type};
  }

  PartialLocalFunction makeDerivative(tag::partial type) const
  {
    return PartialLocalFunction{globalFunction_, type};
  }

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

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

private:
  DiscreteFunction globalFunction_;
  LocalView localView_;
  SubTreeCache const* subTreeCache_;

  std::vector<ValueType> localCoefficients_;
  bool bound_ = false;
};


template <class Coeff, class GB, class TP>
  template <class Type>
class DiscreteFunction<Coeff const,GB,TP>::DerivativeLocalFunctionBase
{
  using R = typename DiscreteFunction::Range;
  using D = typename DiscreteFunction::Domain;
  using RawSignature = typename Dune::Functions::SignatureTraits<R(D)>::RawSignature;
  using Traits = DerivativeTraits<RawSignature, Type>;

public:
  using Domain = typename EntitySet::LocalCoordinate;
  using Range = typename Traits::Range;

  enum { hasDerivative = false };

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

public:
  DerivativeLocalFunctionBase(DiscreteFunction const& globalFunction, Type const& type)
    : globalFunction_(globalFunction)
    , type_(type)
    , localView_(globalFunction_.basis().localView())
    , subTreeCache_(&Dune::TypeTree::child(localView_.treeCache(), globalFunction_.treePath()))
  {}

  DerivativeLocalFunctionBase(DerivativeLocalFunctionBase const& other)
    : globalFunction_(other.globalFunction_)
    , type_(other.type_)
    , localView_(globalFunction_.basis().localView())
    , subTreeCache_(&Dune::TypeTree::child(localView_.treeCache(), globalFunction_.treePath()))
  {}

  void bind(Element const& element)
  {
    localView_.bind(element);
    geometry_.emplace(element.geometry());

    Impl::gather(globalFunction_.coefficients(), localView_, localCoefficients_, Dune::PriorityTag<4>{});
    bound_ = true;
  }

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

  auto order() const
    -> decltype(AMDiS::order(std::declval<SubTree>()))
  {
    assert( bound_ );
    return std::max(0, AMDiS::order(*subTreeCache_)-1);
  }

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

  Geometry const& geometry() const
  {
    assert( bound_ );
    return geometry_.value();
  }

  LocalView const& localView() const
  {
    return localView_;
  }

protected:
  DiscreteFunction globalFunction_;
  Type type_;
  LocalView localView_;
  SubTreeCache const* subTreeCache_;
  std::optional<Geometry> geometry_;
  std::vector<ValueType> localCoefficients_;
  bool bound_ = false;
};


template <class Coeff, class GB, class TP>
class DiscreteFunction<Coeff const,GB,TP>::GradientLocalFunction
    : public DerivativeLocalFunctionBase<tag::gradient>
{
  using Super = DerivativeLocalFunctionBase<tag::gradient>;
public:
  using Range = typename Super::Range;
  using Domain = typename Super::Domain;

  // adopt constructor from base class
  using DerivativeLocalFunctionBase<tag::gradient>::DerivativeLocalFunctionBase;

  Range operator()(Domain const& x) const
  {
    assert( this->bound_ );
    Range dy(0);

    auto&& nodeToRangeEntry = this->globalFunction_.nodeToRangeEntry_;
    for_each_leaf_node(*this->subTreeCache_, [&](auto const& node, auto const& tp)
    {
      LocalToGlobalBasisAdapter localBasis(node, this->geometry());
      auto const& gradients = localBasis.gradientsAt(x);

      // Get range entry associated to this node
      auto re = Dune::Functions::flatVectorView(nodeToRangeEntry(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(this->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;
  }

  using Super::localCoefficients_;
};


template <class Coeff, class GB, class TP>
class DiscreteFunction<Coeff const,GB,TP>::DivergenceLocalFunction
    : public DerivativeLocalFunctionBase<tag::divergence>
{
  using Super = DerivativeLocalFunctionBase<tag::divergence>;

public:
  using Range = typename Super::Range;
  using Domain = typename Super::Domain;

  // adopt constructor from base class
  using DerivativeLocalFunctionBase<tag::divergence>::DerivativeLocalFunctionBase;

  Range operator()(Domain const& x) const
  {
    return evaluate_node(x, bool_t<SubTree::isPower && SubTree::degree() == GridView::dimensionworld>{});
  }

private:
  Range evaluate_node(Domain const& x, std::false_type) const
  {
    error_exit("Cannot calculate divergence(node) where node is not a power node.");
    return Range(0);
  }

  Range evaluate_node(Domain const& x, std::true_type) const
  {
    static_assert(static_size_v<Range> == 1, "Implemented for scalar coefficients only.");

    assert( this->bound_ );
    Range dy(0);

    auto&& node = *this->subTreeCache_;

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

    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;
  }

  using Super::localCoefficients_;
};


template <class Coeff, class GB, class TP>
class DiscreteFunction<Coeff const,GB,TP>::PartialLocalFunction
    : public DerivativeLocalFunctionBase<tag::partial>
{
  using Super = DerivativeLocalFunctionBase<tag::partial>;

public:
  using Range = typename Super::Range;
  using Domain = typename Super::Domain;

  using DerivativeLocalFunctionBase<tag::partial>::DerivativeLocalFunctionBase;

  Range operator()(Domain const& x) const
  {
    assert( this->bound_ );
    Range dy(0);

    std::size_t comp = this->type_.comp;

    auto&& nodeToRangeEntry = this->globalFunction_.nodeToRangeEntry_;
    for_each_leaf_node(*this->subTreeCache_, [&](auto const& node, auto const& tp)
    {
      LocalToGlobalBasisAdapter localBasis(node, this->geometry());
      auto const& partial = localBasis.partialsAt(x, comp);

      // Get range entry associated to this node
      auto re = Dune::Functions::flatVectorView(nodeToRangeEntry(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(this->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;
  }

  using Super::localCoefficients_;
};


} // end namespace AMDiS