DiscreteFunction.inc.hpp

#pragma once

#include <amdis/common/FieldMatVec.hpp>
#include <amdis/utility/LocalBasisCache.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);
    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
  friend GradientLocalFunction derivative(LocalFunction const& localFunction)
  {
    return GradientLocalFunction{localFunction.globalFunction_};
  }

  /// \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_;
  bool bound_ = false;
};


template <class GB, class VT, class TP>
class DiscreteFunction<GB,VT,TP>::GradientLocalFunction
{
  using R = typename DiscreteFunction::Range;
  using D = typename DiscreteFunction::Domain;
  using RawSignature = typename Dune::Functions::SignatureTraits<R(D)>::RawSignature;
  using DerivativeTraits = Dune::Functions::DefaultDerivativeTraits<RawSignature>;

public:
  using Domain = typename EntitySet::LocalCoordinate;
  using Range = typename DerivativeTraits::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.
  GradientLocalFunction(DiscreteFunction const& globalFunction)
    : globalFunction_(globalFunction)
    , localView_(globalFunction_.basis()->localView())
    , subTree_(&child(localView_.tree(), globalFunction_.treePath()))
  {}

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

  void bind(Element const& element)
  {
    localView_.bind(element);
    geometry_.emplace(element.geometry());
    bound_ = true;
  }

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

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

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

private:
  DiscreteFunction globalFunction_;
  LocalView localView_;
  SubTree const* subTree_;
  Dune::Std::optional<Geometry> geometry_;
  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&& coefficients = *globalFunction_.dofVector_;
  auto&& nodeToRangeEntry = globalFunction_.nodeToRangeEntry_;
  for_each_leaf_node(*subTree_, [&,this](auto const& node, auto const& tp)
  {
    auto&& fe = node.finiteElement();
    auto&& localBasis = fe.localBasis();

    NodeCache<TYPEOF(node)> localBasisCache(localBasis);
    auto const& shapeFunctionValues = localBasisCache.evaluateFunction(localView_.element().type(),x);

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

    for (std::size_t i = 0; i < localBasis.size(); ++i) {
      auto&& multiIndex = localView_.index(node.localIndex(i));

      // Get coefficient associated to i-th shape function
      auto c = Dune::Functions::flatVectorView(coefficients[multiIndex]);

      // 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>
typename DiscreteFunction<GB,VT,TP>::GradientLocalFunction::Range DiscreteFunction<GB,VT,TP>::
GradientLocalFunction::operator()(Domain const& x) const
{
  assert( bound_ );
  Range dy(0);

  auto&& coefficients = *globalFunction_.dofVector_;
  auto&& nodeToRangeEntry = globalFunction_.nodeToRangeEntry_;
  for_each_leaf_node(*subTree_, [&,this](auto const& node, auto const& tp)
  {
    // TODO: may DOFVectorView::Range to FieldVector type if necessary
    using LocalDerivativeTraits
      = Dune::Functions::DefaultDerivativeTraits<VT(Domain)>;
    using GradientBlock = typename LocalDerivativeTraits::Range;

    auto&& fe = node.finiteElement();
    auto&& localBasis = fe.localBasis();
    std::size_t size = localBasis.size();

    NodeCache<TYPEOF(node)> localBasisCache(localBasis);
    auto const& referenceGradients = localBasisCache.evaluateJacobian(localView_.element().type(),x);

    // The transposed inverse Jacobian of the map from the reference element to the element
    auto&& jacobian = geometry_.value().jacobianInverseTransposed(x);

    // Compute the shape function gradients on the real element
    std::vector<GradientBlock> gradients(referenceGradients.size());
    for (std::size_t i = 0; i < gradients.size(); ++i) // J^(-T) * D[phi] -> grad^T
      Dune::MatVec::as_vector(gradients[i]) = jacobian * Dune::MatVec::as_vector(referenceGradients[i]);

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

    for (std::size_t i = 0; i < size; ++i) {
      auto&& multiIndex = localView_.index(node.localIndex(i));

      // Get coefficient associated to i-th shape function
      auto c = Dune::Functions::flatVectorView(coefficients[multiIndex]);

      // 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 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