DiscreteFunction.inc.hpp

#pragma once

#include <amdis/common/DerivativeTraits.hpp>
#include <amdis/common/FieldMatVec.hpp>
#include <amdis/utility/LocalBasisCache.hpp>
#include <amdis/utility/LocalToGlobalAdapter.hpp>

#include <dune/common/ftraits.hh>
#include <dune/grid/utility/hierarchicsearch.hh>

namespace AMDiS {

template <class GB, class VT, class TP>
class DiscreteFunction<GB,VT,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:
  /// Constructor. Stores a copy of the DiscreteFunction.
  LocalFunction(DiscreteFunction const& globalFunction)
    : globalFunction_(globalFunction)
    , localView_(globalFunction_.basis()->localView())
    , subTree_(&child(localView_.tree(), globalFunction_.treePath()))
  {}

  /// Copy constructor.
  LocalFunction(LocalFunction const& other)
    : globalFunction_(other.globalFunction_)
    , localView_(globalFunction_.basis()->localView())
    , subTree_(&child(localView_.tree(), globalFunction_.treePath()))
  {}

  /// \brief Bind the LocalView to the element
  void bind(Element const& element)
  {
    localView_.bind(element);

    globalFunction_.coefficients().gather(localView_, localCoefficients_);
    bound_ = true;
  }

  /// \brief Unbind the LocalView from the element
  void unbind()
  {
    localView_.unbind();
    bound_ = false;
  }

  /// \brief Evaluate LocalFunction at bound element in local coordinates
  Range operator()(Domain const& x) const;

  /// \brief Create a LocalFunction representing the gradient. \relates GradientLocalFunction
  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};
  }

  /// \brief The \ref polynomialDegree() of the LocalFunctions
  int order() const
  {
    assert( bound_ );
    return polynomialDegree(*subTree_);
  }

  /// \brief Return the bound element
  Element const& localContext() const
  {
    assert( bound_ );
    return localView_.element();
  }

private:
  DiscreteFunction globalFunction_;
  LocalView localView_;
  SubTree const* subTree_;

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


template <class GB, class VT, class TP>
typename DiscreteFunction<GB,VT,TP>::Range DiscreteFunction<GB,VT,TP>::
LocalFunction::operator()(Domain const& x) const
{
  assert( bound_ );
  Range y(0);

  auto&& nodeToRangeEntry = globalFunction_.nodeToRangeEntry_;
  for_each_leaf_node(*subTree_, [&,this](auto const& node, auto const& tp)
  {
    auto localBasisCache = makeNodeCache(node);
    auto const& shapeFunctionValues = localBasisCache.evaluateFunction(localView_.element().type(), 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;
}


template <class GB, class VT, class TP>
  template <class Type>
class DiscreteFunction<GB,VT,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:
  /// Constructor. Stores a copy of the DiscreteFunction.
  DerivativeLocalFunctionBase(DiscreteFunction const& globalFunction, Type const& type)
    : globalFunction_(globalFunction)
    , type_(type)
    , localView_(globalFunction_.basis()->localView())
    , subTree_(&child(localView_.tree(), globalFunction_.treePath()))
  {}

  /// Copy constructor.
  DerivativeLocalFunctionBase(DerivativeLocalFunctionBase const& other)
    : globalFunction_(other.globalFunction_)
    , type_(other.type_)
    , localView_(globalFunction_.basis()->localView())
    , subTree_(&child(localView_.tree(), globalFunction_.treePath()))
  {}

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

    globalFunction_.coefficients().gather(localView_, localCoefficients_);
    bound_ = true;
  }

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

  int order() const
  {
    assert( bound_ );
    return std::max(0, polynomialDegree(*subTree_)-1);
  }

  /// Return the bound element
  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_;
  SubTree const* subTree_;
  Dune::Std::optional<Geometry> geometry_;
  std::vector<VT> localCoefficients_;
  bool bound_ = false;
};


template <class GB, class VT, class TP>
class DiscreteFunction<GB,VT,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;

  /// Evaluate Gradient at bound element in local coordinates
  Range operator()(Domain const& x) const
  {
    assert( this->bound_ );
    Range dy(0);

    auto&& nodeToRangeEntry = this->globalFunction_.nodeToRangeEntry_;
    for_each_leaf_node(*this->subTree_, [&](auto const& node, auto const& tp)
    {
      auto localBasis = makeLocalToGlobalBasisAdapter(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 GB, class VT, class TP>
class DiscreteFunction<GB,VT,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;

  /// Evaluate divergence at bound element in local coordinates
  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->subTree_;

    auto localBasis = makeLocalToGlobalBasisAdapter(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 GB, class VT, class TP>
class DiscreteFunction<GB,VT,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;

  /// Evaluate partial derivative at bound element in local coordinates
  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->subTree_, [&](auto const& node, auto const& tp)
    {
      auto localBasis = makeLocalToGlobalBasisAdapter(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_;
};


template <class GB, class VT, class TP>
typename DiscreteFunction<GB,VT,TP>::Range DiscreteFunction<GB,VT,TP>::
operator()(Domain const& x) const
{
  using Grid = typename GlobalBasis::GridView::Grid;
  using IS = typename GlobalBasis::GridView::IndexSet;

  auto const& gv = this->basis()->gridView();
  Dune::HierarchicSearch<Grid,IS> hsearch{gv.grid(), gv.indexSet()};

  auto element = hsearch.findEntity(x);
  auto geometry = element.geometry();
  auto localFct = localFunction(*this);
  localFct.bind(element);
  return localFct(geometry.local(x));
}

} // end namespace AMDiS