block_tensor_vector.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <uni20/krylov/matrix_free.hpp>
#include <uni20/symmetry/block_tensor_linear.hpp>
#include <uni20/symmetry/block_tensor_permute.hpp>
 
#include <algorithm>
#include <concepts>
#include <cstddef>
#include <functional>
#include <stdexcept>
#include <type_traits>
#include <utility>
#include <vector>
 
namespace uni20::krylov
{
namespace detail
{
 
template <class Tensor> struct IsBlockTensorValue : std::false_type
{};
 
template <typename Scalar, class DomainType, class CodomainType, BlockTensorStorage Storage>
struct IsBlockTensorValue<BlockTensor<Scalar, DomainType, CodomainType, Storage>> : std::true_type
{};
 
template <class View, class Tensor>
concept CompatibleBlockVector =
    BlockTensorView<View> && std::same_as<block_tensor_value_t<View>, typename Tensor::value_type> &&
    std::same_as<block_tensor_domain_t<View>, typename Tensor::domain_type> &&
    std::same_as<block_tensor_codomain_t<View>, typename Tensor::codomain_type> &&
    std::same_as<typename block_tensor_type_t<View>::storage_policy, typename Tensor::storage_policy> &&
    (DiagonalBlockTensorView<View> == DiagonalBlockTensorView<Tensor>);
 
} // namespace detail
 
template <class Tensor> class BlockTensorVectorOps {
  public:
    using tensor_type = std::remove_cvref_t<Tensor>;
    using scalar_type = typename tensor_type::value_type;
    using real_type = make_real_t<scalar_type>;
    using domain_type = typename tensor_type::domain_type;
    using codomain_type = typename tensor_type::codomain_type;
    using key_type = typename tensor_type::key_type;
    static_assert(detail::IsBlockTensorValue<tensor_type>::value,
                  "BlockTensorVectorOps requires an owning BlockTensor value");
    static_assert(RealOrComplex<scalar_type>, "BlockTensorVectorOps requires a real or complex scalar type");
 
    explicit BlockTensorVectorOps(tensor_type const& prototype)
        : symmetry_(prototype.symmetry()), domain_(prototype.domain()), codomain_(prototype.codomain()),
          dimension_(block_tensor_dimension(prototype))
    {
      auto const keys = prototype.stored_keys();
      stored_keys_.assign(keys.begin(), keys.end());
    }
 
    [[nodiscard]] auto problem_dimension() const noexcept -> std::size_t { return dimension_; }
 
    template <class Vector>
      requires detail::CompatibleBlockVector<Vector, tensor_type>
    [[nodiscard]] auto vector_dimension(Vector const& vector) const -> std::size_t
    {
      this->require_member(vector);
      return dimension_;
    }
 
    template <class Vector>
      requires detail::CompatibleBlockVector<Vector, tensor_type> &&
               (std::is_lvalue_reference_v<Vector&&> || BorrowedBlockTensorView<Vector>)
    [[nodiscard]] auto prototype_view(Vector&& vector) const
    {
      this->require_member(vector);
      return as_block_tensor_view(std::as_const(vector));
    }
 
    [[nodiscard]] auto allocate_like(tensor_type const& vector) const -> tensor_type
    {
      this->require_member(vector);
      if constexpr (requires { vector.allocate_like(uninitialized); })
      {
        return vector.allocate_like(uninitialized);
      }
      else if constexpr (CompleteBlockStorage<typename tensor_type::storage_policy>)
      {
        return tensor_type(uninitialized, symmetry_, domain_, codomain_);
      }
      else
      {
        return tensor_type(uninitialized, symmetry_, domain_, codomain_, stored_keys_);
      }
    }
 
    template <BorrowedBlockTensorView View>
      requires detail::CompatibleBlockVector<View, tensor_type>
    [[nodiscard]] auto allocate_like(View const& view) const -> tensor_type
    {
      this->require_member(view);
      if constexpr (CompleteBlockStorage<typename tensor_type::storage_policy>)
      {
        if constexpr (requires { view.allocation_context(); })
          return tensor_type(uninitialized, symmetry_, domain_, codomain_, view.allocation_context());
        else
          return tensor_type(uninitialized, symmetry_, domain_, codomain_);
      }
      else
      {
        if constexpr (requires { view.allocation_context(); })
          return tensor_type(uninitialized, symmetry_, domain_, codomain_, stored_keys_, view.allocation_context());
        else
          return tensor_type(uninitialized, symmetry_, domain_, codomain_, stored_keys_);
      }
    }
 
    void copy(tensor_type& output, tensor_type const& input) const
    {
      this->require_member(output);
      this->require_member(input);
      uni20::copy(output, input);
    }
 
    void axpy(tensor_type& output, scalar_type factor, tensor_type const& input) const
    {
      this->require_member(output);
      this->require_member(input);
      uni20::axpy(output, factor, input);
    }
 
    void scal(tensor_type& vector, scalar_type factor) const
    {
      this->require_member(vector);
      uni20::scale(vector, factor);
    }
 
    void set_zero(tensor_type& vector) const
    {
      this->require_member(vector);
      uni20::set_zero(vector);
    }
 
    [[nodiscard]] auto inner_product(tensor_type const& lhs, tensor_type const& rhs) const -> scalar_type
    {
      this->require_member(lhs);
      this->require_member(rhs);
      return uni20::inner_product_host(lhs, rhs);
    }
 
    [[nodiscard]] auto norm(tensor_type const& vector) const -> real_type
    {
      this->require_member(vector);
      return uni20::norm_host(vector);
    }
 
  private:
    [[nodiscard]] static auto block_tensor_dimension(tensor_type const& tensor) -> std::size_t
    {
      std::size_t dimension = 0;
      for (std::size_t ordinal = 0; ordinal < tensor.stored_block_count(); ++ordinal)
      {
        auto block = tensor.block_by_ordinal(ordinal);
        std::size_t block_dimension;
        if constexpr (DiagonalBlockTensorView<tensor_type>)
        {
          auto descriptor = mdspec_of(block);
          block_dimension = static_cast<std::size_t>(diagonal_components(descriptor).extent(0));
        }
        else
        {
          block_dimension = 1;
          for (std::size_t axis = 0; axis < tensor_type::dense_block_order(); ++axis)
            block_dimension *= static_cast<std::size_t>(block.extent(axis));
        }
        dimension += block_dimension;
      }
      return dimension;
    }
 
    template <class Vector>
      requires detail::CompatibleBlockVector<Vector, tensor_type>
    void require_member(Vector const& vector) const
    {
      if (vector.symmetry() != symmetry_ || vector.domain() != domain_ || vector.codomain() != codomain_)
      {
        throw std::invalid_argument("BlockTensor Krylov vector has incompatible symmetry or boundaries");
      }
      if (!std::ranges::equal(vector.stored_keys(), stored_keys_))
      {
        throw std::invalid_argument("BlockTensor Krylov vector has a different stored-key pattern");
      }
    }
 
    Symmetry symmetry_;
    domain_type domain_;
    codomain_type codomain_;
    std::vector<key_type> stored_keys_;
    std::size_t dimension_ = 0;
};
 
template <class Tensor, class Operator> class BlockTensorMatrixFreeOps : public BlockTensorVectorOps<Tensor> {
  public:
    using base_type = BlockTensorVectorOps<Tensor>;
    using tensor_type = typename base_type::tensor_type;
    using operator_type = std::remove_cvref_t<Operator>;
    static_assert(std::invocable<operator_type&, tensor_type&, tensor_type const&>,
                  "BlockTensor matrix-free operation must be callable as operation(output, input)");
    static_assert(std::same_as<std::invoke_result_t<operator_type&, tensor_type&, tensor_type const&>, void>,
                  "BlockTensor matrix-free operation must return void");
 
    BlockTensorMatrixFreeOps(tensor_type const& prototype, operator_type operation)
        : base_type(prototype), operation_(std::move(operation))
    {}
 
    void matvec(tensor_type& output, tensor_type const& input)
    {
      static_cast<void>(this->vector_dimension(output));
      static_cast<void>(this->vector_dimension(input));
      std::invoke(operation_, output, input);
    }
 
  private:
    operator_type operation_;
};
 
template <class Tensor, class Operator>
BlockTensorMatrixFreeOps(Tensor const&, Operator) -> BlockTensorMatrixFreeOps<std::remove_cvref_t<Tensor>, Operator>;
 
} // namespace uni20::krylov