DOFVectorView.hpp

#pragma once

#include <vector>

#include <dune/common/std/optional.hh>
#include <dune/functions/common/defaultderivativetraits.hh>
#include <dune/functions/common/treedata.hh>
#include <dune/functions/functionspacebases/defaultnodetorangemap.hh>
#include <dune/functions/functionspacebases/flatvectorbackend.hh>
#include <dune/functions/gridfunctions/gridviewentityset.hh>
#include <dune/typetree/childextraction.hh>

#include <amdis/GridFunctions.hpp>
#include <amdis/utility/FiniteElementType.hpp>

namespace AMDiS
{
  /**
    * \addtogroup GridFunctions
    * @{
    **/

  template <class Traits, class TreePathType>
  class DOFVectorConstView
  {
  public:
    using GlobalBasis = typename Traits::GlobalBasis;
    using TreePath = TreePathType;
    using Vector = DOFVector<Traits>;

    using Tree = typename GlobalBasis::LocalView::Tree;
    using SubTree = typename Dune::TypeTree::ChildForTreePath<Tree, TreePath>;
    using NodeToRangeEntry = Dune::Functions::DefaultNodeToRangeMap<SubTree>;

    using GridView = typename GlobalBasis::GridView;
    using EntitySet = Dune::Functions::GridViewEntitySet<GridView, 0>;

    using Domain = typename EntitySet::GlobalCoordinate;
    using Range = RangeType_t<SubTree>;

    using RawSignature = typename Dune::Functions::SignatureTraits<Range(Domain)>::RawSignature;
    using DerivativeTraits = Dune::Functions::DefaultDerivativeTraits<RawSignature>;
    using DerivativeRange = typename DerivativeTraits::Range;

    using LocalDomain = typename EntitySet::LocalCoordinate;
    using Element = typename EntitySet::Element;
    using Geometry = typename Element::Geometry;

    template <class Block>
    using Flat = Dune::Functions::FlatVectorBackend<Block>;


  public: // a local view on the gradients

    /// A LocalFunction representing the derivative of the DOFVector
    class GradientLocalFunction
    {
    public:
      using Domain = LocalDomain;
      using Range = DerivativeRange;

    private:
      using LocalBasisView = typename GlobalBasis::LocalView;
      using LocalIndexSet = typename GlobalBasis::LocalIndexSet;

      template <class LeafNode>
      using LocalBasisJacobian = typename LeafNode::FiniteElement::Traits::LocalBasisType::Traits::JacobianType;

      template <class Node>
      using NodeData = typename std::vector<LocalBasisJacobian<Node>>;

      using ReferenceGradientContainer = Dune::Functions::TreeData<SubTree, NodeData, true>;

    public:
      GradientLocalFunction(DOFVectorConstView const& globalFunction)
        : globalFunction_(&globalFunction)
        , localBasisView_(globalFunction_->basis().localView())
        , localIndexSet_(globalFunction_->basis().localIndexSet())
        , subTree_(&Dune::TypeTree::child(localBasisView_.tree(), globalFunction_->treePath()))
      {
        referenceGradientContainer_.init(*subTree_);
      }

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

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

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

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

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

    private:
      DOFVectorConstView const* globalFunction_;
      LocalBasisView localBasisView_;
      LocalIndexSet localIndexSet_;

      SubTree const* subTree_;
      mutable ReferenceGradientContainer referenceGradientContainer_;

      Dune::Std::optional<Geometry> geometry_;
      bool bound_ = false;
    };


  public: // a local view on the values

    /// A LocalFunction, i.e., an element local view on the DOFVector
    class LocalFunction
    {
    public:
      using Domain = typename DOFVectorConstView::LocalDomain;
      using Range = typename DOFVectorConstView::Range;

    private:
      using LocalBasisView = typename GlobalBasis::LocalView;
      using LocalIndexSet = typename GlobalBasis::LocalIndexSet;

      template <class LeafNode>
      using LocalBasisRange = RangeType_t<LeafNode>;
        // = typename LeafNode::FiniteElement::Traits::LocalBasisType::Traits::RangeType;

      template <class Node>
      using NodeData = typename std::vector<LocalBasisRange<Node>>;

      using ShapeFunctionValueContainer = Dune::Functions::TreeData<SubTree, NodeData, true>;

    public:
      LocalFunction(DOFVectorConstView const& globalFunction)
        : globalFunction_(&globalFunction)
        , localBasisView_(globalFunction_->basis().localView())
        , localIndexSet_(globalFunction_->basis().localIndexSet())
        , subTree_(&Dune::TypeTree::child(localBasisView_.tree(), globalFunction_->treePath()))
      {
        shapeFunctionValueContainer_.init(*subTree_);
      }

      void bind(Element const& element)
      {
        localBasisView_.bind(element);
        localIndexSet_.bind(localBasisView_);
        bound_ = true;
      }

      void unbind()
      {
        localIndexSet_.unbind();
        localBasisView_.unbind();
        bound_ = false;
      }

      /// 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)
      {
        static_assert(isValidRange<DerivativeTraits>(),"Derivative of DOFVector not defined.");
        return GradientLocalFunction{*localFunction.globalFunction_};
      }

      friend int order(LocalFunction const& self)
      {
        assert( self.bound_ );
        return polynomialDegree(*self.subTree_);
      }

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

    private:
      DOFVectorConstView const* globalFunction_;
      LocalBasisView localBasisView_;
      LocalIndexSet localIndexSet_;
      SubTree const* subTree_;

      mutable ShapeFunctionValueContainer shapeFunctionValueContainer_;

      bool bound_ = false;
    };


  public:
    /// Constructor. Stores a pointer to the dofVector and a copy of the treePath.
    DOFVectorConstView(DOFVector<Traits> const& dofVector, TreePath const& treePath)
      : dofVector_(&dofVector)
      , treePath_(treePath)
      , entitySet_(dofVector.getFeSpace().gridView())
      , nodeToRangeEntry_(Dune::Functions::makeDefaultNodeToRangeMap(dofVector.getFeSpace(), treePath))
    {}

    /// Evaluate the view on this DOFVector in global coordinates
    Range operator()(Domain const& x) const
    {
      error_exit("Not implemented.");
      return Range(0);
    }

    /// \brief Create a local function for this view on the DOFVector. \relates LocalFunction
    friend LocalFunction localFunction(DOFVectorConstView const& self)
    {
      return LocalFunction{self};
    }

    EntitySet const& entitySet() const
    {
      return entitySet_;
    }

  public:
    /// Return global basis
    GlobalBasis const& basis() const
    {
      return dofVector_->getFeSpace();
    }

    /// Return treePath associated with this view
    TreePath const& treePath() const
    {
      return treePath_;
    }

    /// Return const coefficient vector
    DOFVector<Traits> const& coefficients() const
    {
      return *dofVector_;
    }

  protected:
    DOFVector<Traits> const* dofVector_;
    TreePath const treePath_;

    EntitySet entitySet_;
    NodeToRangeEntry nodeToRangeEntry_;
  };


  // A mutable version of DOFVectorView
  template <class Traits, class TreePathType>
  class DOFVectorMutableView
      : public DOFVectorConstView<Traits, TreePathType>
  {
    using Super = DOFVectorConstView<Traits, TreePathType>;

    using GlobalBasis = typename Traits::GlobalBasis;
    using TreePath = TreePathType;

  public:
    /// Constructor. Stores a pointer to the mutable `dofvector`.
    DOFVectorMutableView(DOFVector<Traits>& dofVector, TreePath const& treePath)
      : Super(dofVector, treePath)
      , mutableDofVector_(&dofVector)
    {}

  public:
    /// Interpolation of GridFunction to DOFVector
    template <class Expr>
    DOFVectorMutableView& interpolate(Expr&& expr)
    {
      auto const& basis = Super::basis();
      auto const& treePath = Super::treePath();

      auto&& gridFct = makeGridFunction(std::forward<Expr>(expr), basis.gridView());

      DOFVector<Traits> tmp(basis, "tmp");
      Dune::Functions::interpolate(basis, treePath, tmp, std::forward<decltype(gridFct)>(gridFct));

      // move data from temporary vector into stored DOFVector
      mutableDofVector_->getVector() = std::move(tmp.getVector());
      return *this;
    }

    /// Return the mutable DOFVector
    DOFVector<Traits>& coefficients() { return *mutableDofVector_; }

    /// Return the const DOFVector
    using Super::coefficients;

  protected:
    DOFVector<Traits>* mutableDofVector_;
  };

  /** @} **/


#ifndef DOXYGEN
  // A Generator for a const \ref DOFVectorView.
  template <class Traits, class TreePath>
  auto makeDOFVectorView(DOFVector<Traits> const& dofVector, TreePath const& treePath)
  {
    return DOFVectorConstView<Traits, TreePath>{dofVector, treePath};
  }

  // A Generator for a mutable \ref DOFVectorView.
  template <class Traits, class TreePath>
  auto makeDOFVectorView(DOFVector<Traits>& dofVector, TreePath const& treePath)
  {
    return DOFVectorMutableView<Traits, TreePath>{dofVector, treePath};
  }


  // A Generator for a const \ref DOFVectorView.
  template <class Traits>
  auto makeDOFVectorView(DOFVector<Traits> const& dofVector)
  {
    auto treePath = Dune::TypeTree::hybridTreePath();
    return DOFVectorConstView<Traits, decltype(treePath)>{dofVector, treePath};
  }

  // A Generator for a mutable \ref DOFVectorView.
  template <class Traits>
  auto makeDOFVectorView(DOFVector<Traits>& dofVector)
  {
    auto treePath = Dune::TypeTree::hybridTreePath();
    return DOFVectorMutableView<Traits, decltype(treePath)>{dofVector, treePath};
  }
#endif

} // end namespace AMDiS

#include "DOFVectorView.inc.hpp"