tensors.hh

Go to the documentation of this file.

// -*- tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 2 -*-
// vi: set et ts=4 sw=2 sts=2:
 
// SPDX-FileCopyrightText: Copyright © DUNE-FUFEM Project contributors, see file AUTHORS.md
// SPDX-License-Identifier: LicenseRef-GPL-2.0-only-with-DUNE-exception OR LGPL-3.0-or-later
 
#ifndef DUNE_FUFEM_FORMS_TENSORS_HH
#define DUNE_FUFEM_FORMS_TENSORS_HH
 
#include <array>
#include <cstddef>
#include <type_traits>
#include <utility>
 
#include <dune/common/indices.hh>
#include <dune/common/rangeutilities.hh>
#include <dune/common/typelist.hh>
 
 
 
namespace Dune::Fufem::Forms::Impl::Tensor {
 
 
 
  // Class to tag tensor implementations
  template<class T>
  struct IsTensor : public std::false_type
  {};
 
  // Class to tag tensor with lazy evaluation
  template<class T>
  struct IsLazyTensor : public std::false_type
  {};
 
 
 
  // This forwards operator()(i,...) to subsequent operator[](i) calls
  template<class T, std::size_t k>
  class TensorView
  {
    T& t_;
 
    static decltype(auto) resolveIndices(auto&& t)
    {
      return t;
    }
 
    static decltype(auto) resolveIndices(auto&& t, const auto& i0, const auto&... i)
    {
      return resolveIndices(t[i0], i...);
    }
 
    template<class I, I... i>
    static decltype(auto) resolveIndexSequence(auto&&t, std::integer_sequence<I, i...>)
    {
      return resolveIndices(t, i...);
    }
 
  public:
 
    using value_type = std::decay_t<decltype(resolveIndexSequence(std::declval<T&>(), std::make_index_sequence<k>()))>;
 
    static constexpr Dune::index_constant<k> rank()
    {
      return {};
    }
 
    TensorView(T& t, Dune::index_constant<k>)
      : t_(t)
    {}
 
    decltype(auto) operator()(const auto&... i) const
    {
      return resolveIndices(t_, i...);
    }
  };
 
  template<class T, std::size_t k>
  struct IsTensor<TensorView<T, k>> : public std::true_type
  {};
 
 
 
  // This wraps a T and makes it available using k calls to operator[]
  template<class T, std::size_t k>
  requires(k>0)
  class LiftRank
  {
    using RawT = Dune::ResolveRef_t<T>;
    T t_;
  public:
 
    using value_type = typename T::value_type;
 
    static constexpr Dune::index_constant<RawT::rank()+k> rank()
    {
      return {};
    }
 
    LiftRank(T t, Dune::index_constant<k>)
      : t_(t)
    {}
 
    template<class I0, class... I>
    requires (k==1)
    decltype(auto) operator()(const I0& i0, const I&... i) const
    {
      return Dune::resolveRef(t_)(i...);
    }
 
    template<class I0, class I1, class... I>
    requires (k==2)
    decltype(auto) operator()(const I0& i0, const I1& i1, const I&... i) const
    {
      return Dune::resolveRef(t_)(i...);
    }
  };
 
  template<class T, std::size_t k>
  struct IsTensor<LiftRank<T, k>> : public std::true_type
  {};
 
 
 
  // This wraps flips the first k indices between 0 and 1
  template<class T, std::size_t k>
  class FlipIndices
  {
    using RawT = Dune::ResolveRef_t<T>;
    T t_;
  public:
 
    using value_type = typename RawT::value_type;
 
    static constexpr Dune::index_constant<RawT::rank()> rank()
    {
      return {};
    }
 
    FlipIndices(T t, Dune::index_constant<k>)
      : t_(t)
    {}
 
    template<class I0, class... I>
    requires (k==1)
    decltype(auto) operator()(const I0& i0, const I&... i) const
    {
      return Dune::resolveRef(t_)(not i0, i...);
    }
 
    template<class I0, class I1, class... I>
    requires (k==2)
    decltype(auto) operator()(const I0& i0, const I1& i1, const I&... i) const
    {
      return Dune::resolveRef(t_)(not i0, not i1, i...);
    }
  };
 
  template<class T, std::size_t k>
  struct IsTensor<FlipIndices<T, k>> : public std::true_type
  {};
 
 
 
  template<class Value>
  class RankZeroTensor
  {
  public:
 
    using value_type = Value;
 
    static constexpr Dune::index_constant<0> rank()
    {
      return {};
    }
 
    RankZeroTensor(Value value)
      : value_(std::move(value))
    {}
 
    template<class F>
    friend void sparseForEach(const RankZeroTensor& tensor, F&& f)
    {
      f(tensor.value_);
    }
 
    template<class Outer>
    friend auto compose(const Outer& outer, const RankZeroTensor& inner)
    {
      return Dune::Fufem::Forms::Impl::Tensor::RankZeroTensor(outer(inner()));
    }
 
    static constexpr std::size_t nnz()
    {
      return 1;
    }
 
    // Special interface for rank=0
    const Value& operator()() const
    {
      return value_;
    }
 
  protected:
    Value value_;
  };
 
  template<class Value>
  struct IsTensor<RankZeroTensor<Value>> : public std::true_type
  {};
 
 
 
  template<std::size_t r, class Value, class ForEach>
  class SparseTensor
  {
  public:
 
    using value_type = std::decay_t<Value>;
 
    static constexpr Dune::index_constant<r> rank()
    {
      return {};
    }
 
    SparseTensor(Dune::index_constant<r>, ForEach forEach, std::size_t nnz, Dune::MetaType<Value>)
      : forEach_(std::move(forEach))
      , nnz_(nnz)
    {}
 
    template<class F>
    friend void sparseForEach(const SparseTensor& tensor, F&& f)
    {
      tensor.forEach_(f);
    }
 
    template<class Outer>
    friend auto compose(const Outer& outer, const SparseTensor& inner)
    {
      using ComposedRange = std::decay_t<std::invoke_result_t<Outer, typename SparseTensor::value_type>>;
      return Dune::Fufem::Forms::Impl::Tensor::SparseTensor(
        inner.rank(),
        [outer, innerForEach=inner.forEach_](auto&& f) {
          innerForEach([&](const auto& y_i, auto... i) {
            f(outer(y_i), i...);
          });
        },
        inner.nnz(),
        Dune::MetaType<ComposedRange&&>{}
      );
    }
 
    std::size_t nnz() const
    {
      return nnz_;
    }
 
  protected:
    ForEach forEach_;
    std::size_t nnz_;
  };
 
  template<std::size_t r, class Value, class ForEach>
  struct IsTensor<SparseTensor<r, Value, ForEach>> : public std::true_type
  {};
 
  template<std::size_t r, class Value, class ForEach>
  struct IsLazyTensor<SparseTensor<r, Value, ForEach>> : public std::is_rvalue_reference<Value>
  {};
 
 
 
  template<class Product, class T0, class T1>
  class SparseProductTensor
  {
  public:
 
    using value_type = std::decay_t<std::invoke_result_t<Product, typename T0::value_type, typename T1::value_type>>;
 
    static constexpr Dune::index_constant<T0::rank() + T1::rank()> rank()
    {
      return {};
    }
 
    SparseProductTensor(Product product, T0 t0, T1 t1)
      : product_(std::move(product))
      , t0_(std::move(t0))
      , t1_(std::move(t1))
    {}
 
    template<class F>
    friend void sparseForEach(const SparseProductTensor& tensor, F&& f)
    {
      using namespace Dune::Indices;
      if constexpr(T0::rank() == 0)
        sparseForEach(tensor.t1_, [&](const auto& x1_j, auto... j) {
          f(tensor.product_(tensor.t0_(), x1_j), j...);
        });
      else if constexpr(T1::rank() == 0)
        sparseForEach(tensor.t0_, [&](const auto& x0_i, auto... i) {
          f(tensor.product_(x0_i, tensor.t1_()), i...);
        });
      else if constexpr((T0::rank() == 2) and (T1::rank() == 2))
        sparseForEach(tensor.t0_, [&](const auto& x0_i, auto i0, auto i1) {
          sparseForEach(tensor.t1_, [&](const auto& x1_j, auto j0, auto j1) {
            f(tensor.product_(x0_i, x1_j), i0, j0, i1, j1);
          });
        });
    }
 
    template<class Outer>
    friend auto compose(const Outer& outer, const SparseProductTensor& inner)
    {
      using namespace Dune::Indices;
      return Dune::Fufem::Forms::Impl::Tensor::SparseProductTensor(
        [outer, innerProduct=inner.product()](const auto&... args) { return outer(innerProduct(args...)); },
        inner.factor(_0),
        inner.factor(_1)
      );
    }
 
    std::size_t nnz() const
    {
      return t0_.nnz() * t1_.nnz();
    }
 
    // Special interface for product tensor
    const Product& product() const
    {
      return product_;
    }
 
    template<std::size_t k>
    const auto& factor(Dune::index_constant<k> = {}) const
    {
      if constexpr(k==0)
        return t0_;
      else if constexpr(k==1)
        return t1_;
    }
 
  protected:
    Product product_;
    T0 t0_;
    T1 t1_;
  };
 
  template<class Product, class T0, class T1>
  struct IsTensor<SparseProductTensor<Product, T0, T1>> : public std::true_type
  {};
 
  template<class Product, class T0, class T1>
  struct IsLazyTensor<SparseProductTensor<Product, T0, T1>> : public std::true_type
  {};
 
 
 
  template<class Product, class X, class Y>
  requires (IsTensor<X>::value and IsTensor<Y>::value)
  auto interleavedOuterProduct(Product product, const X& x, const Y& y)
  {
    using namespace Dune::Indices;
    if constexpr((X::rank() == 0) and (Y::rank() == 0))
      return RankZeroTensor(product(x(), y()));
    else
      return SparseProductTensor(product, x, y);
  }
 
 
 
  template<class K, std::size_t rank, class Value, class ForEach, class Y>
  requires(
    Dune::IsNumber<K>::value
    and IsTensor<Y>::value
    and (rank == Y::rank())
  )
  void axpy(const K& alpha, const SparseTensor<rank, Value, ForEach>& x, Y& y)
  {
    sparseForEach(x, [&](const auto& xi, auto...i) {
      y(i...) += alpha*xi;
    });
  }
 
  template<class K, class Product, class T0, class T1, class Y>
  requires(
    Dune::IsNumber<K>::value
    and IsTensor<Y>::value
    and (SparseProductTensor<Product, T0, T1>::rank() == Y::rank())
  )
  void axpy(const K& alpha, const SparseProductTensor<Product, T0, T1>& x, Y& y)
  {
    using namespace Dune::Indices;
    // Since x.product() is a bilinear we can factor out
    // the multiplication with alpha.
    if constexpr(T0::rank() == 0)
      sparseForEach(x.factor(_1), [&](const auto& x1_j, auto... j) {
        y(j...) += x.product()(alpha*x.factor(_0)(), x1_j);
      });
    else if constexpr(T1::rank() == 0)
      sparseForEach(x.factor(_0), [&](const auto& x0_i, auto... i) {
        y(i...) += x.product()(x0_i, alpha*x.factor(_1)());
      });
    else if constexpr((T0::rank() == 2) and (T1::rank() == 2))
      sparseForEach(x.factor(_0), [&](const auto& x0_i, auto i0, auto i1) {
        auto x0_i_alpha = x0_i*alpha;
        sparseForEach(x.factor(_1), [&](const auto& x1_j, auto j0, auto j1) {
          y(i0, j0, i1, j1) += x.product()(x0_i_alpha, x1_j);
        });
      });
    else
      static_assert((T0::rank() == 2) and (T1::rank() == 2));
  }
 
} // namespace Dune::Fufem::Forms::Impl::Tensor
 
 
 
#endif // DUNE_FUFEM_FORMS_TENSORS_HH