environment_cache.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <uni20/tensor_network/environment.hpp>
#include <uni20/tensor_network/finite_chain.hpp>
 
#include <concepts>
#include <cstddef>
#include <memory>
#include <optional>
#include <stdexcept>
#include <type_traits>
#include <vector>
 
namespace uni20::tensor_network
{
 
template <class MpsChain, class MpoChain, SparseBlockStorage EnvironmentStorage = SeparateSparseBlockStorage<>>
class MpoEnvironmentCache {
  public:
    using mps_type = std::remove_cvref_t<MpsChain>;
    using mpo_type = std::remove_cvref_t<MpoChain>;
    using value_type = typename mps_type::value_type;
    using bond_space_type = typename mps_type::bond_space_type;
    using auxiliary_space_type = typename mpo_type::auxiliary_space_type;
    using environment_type =
        MpoEnvironment<value_type, bond_space_type, auxiliary_space_type, bond_space_type, EnvironmentStorage>;
    using mpo_block_type = block_tensor_const_block_t<typename mpo_type::site_type>;
 
    static_assert(std::same_as<value_type, typename mpo_type::value_type>,
                  "MPS and MPO chains must have the same scalar type");
    static_assert(std::same_as<typename mps_type::physical_space_type, typename mpo_type::physical_space_type>,
                  "MPS and MPO chains must use the same physical-space type");
    static_assert(std::same_as<bond_space_type, BlockSpace>,
                  "the first finite environment cache requires BlockSpace bonds");
    static_assert(std::same_as<auxiliary_space_type, LocalSpace>,
                  "the first finite environment cache requires LocalSpace MPO auxiliaries");
    static_assert(BlockTensorView<typename mps_type::site_type>);
    static_assert(ImmediateBlockTensorView<typename mpo_type::site_type>);
    static_assert(HostAccessibleMdspan<immediate_tensor_mdspan_t<mpo_block_type>>);
    static_assert(detail::EnvironmentStorageFor<EnvironmentStorage, value_type>);
    static_assert(std::same_as<typename mps_type::storage_policy::leaf_storage_policy,
                               typename EnvironmentStorage::leaf_storage_policy>,
                  "MPS sites and cached environments must use one leaf memory domain");
 
    MpoEnvironmentCache(mps_type const& mps, mpo_type const& mpo, std::size_t left_auxiliary_index,
                        std::size_t right_auxiliary_index)
        : mps_(&mps), mpo_(&mpo), left_(mps.size() + 1), right_(mps.size() + 1), mps_revisions_(mps.size()),
          mpo_revisions_(mpo.size())
    {
      this->validate();
      for (std::size_t index = 0; index < this->size(); ++index)
      {
        mps_revisions_[index] = mps_->revision(index);
        mpo_revisions_[index] = mpo_->revision(index);
      }
 
      left_[0].emplace([&] {
        auto const& site = mps_->site(0);
        if constexpr (requires { site.allocation_context(); })
          return make_identity_mpo_environment<value_type, EnvironmentStorage>(
              site.domain().template space<0>(), mpo_->site(0).domain().template space<0>(), left_auxiliary_index,
              site.allocation_context());
        else
          return make_identity_mpo_environment<value_type, EnvironmentStorage>(
              site.domain().template space<0>(), mpo_->site(0).domain().template space<0>(), left_auxiliary_index);
      }());
      right_[this->size()].emplace([&] {
        auto const& site = mps_->site(this->size() - 1);
        if constexpr (requires { site.allocation_context(); })
          return make_identity_mpo_environment<value_type, EnvironmentStorage>(
              site.codomain().template space<0>(), mpo_->site(this->size() - 1).codomain().template space<0>(),
              right_auxiliary_index, site.allocation_context());
        else
          return make_identity_mpo_environment<value_type, EnvironmentStorage>(
              site.codomain().template space<0>(), mpo_->site(this->size() - 1).codomain().template space<0>(),
              right_auxiliary_index);
      }());
    }
 
    auto size() const noexcept -> std::size_t { return mps_->size(); }
 
    auto is_attached_to(mps_type const& mps, mpo_type const& mpo) const noexcept -> bool
    {
      return mps_ == std::addressof(mps) && mpo_ == std::addressof(mpo);
    }
 
    auto left_cached(std::size_t bond) -> bool
    {
      this->validate_bond(bond);
      this->synchronize_revisions();
      return left_[bond].has_value();
    }
 
    auto right_cached(std::size_t bond) -> bool
    {
      this->validate_bond(bond);
      this->synchronize_revisions();
      return right_[bond].has_value();
    }
 
    auto left_environment(std::size_t bond) -> environment_type const&
    {
      this->validate_bond(bond);
      this->synchronize_revisions();
      if (!left_[bond])
      {
        std::size_t first_missing = bond;
        while (!left_[first_missing])
          --first_missing;
        for (std::size_t site = first_missing; site < bond; ++site)
        {
          left_[site + 1].emplace(
              extend_left_environment<EnvironmentStorage>(*left_[site], mps_->site(site), mpo_->site(site)));
        }
      }
      return *left_[bond];
    }
 
    auto right_environment(std::size_t bond) -> environment_type const&
    {
      this->validate_bond(bond);
      this->synchronize_revisions();
      if (!right_[bond])
      {
        std::size_t first_available = bond;
        while (!right_[first_available])
          ++first_available;
        for (std::size_t site = first_available; site > bond; --site)
        {
          std::size_t const site_index = site - 1;
          right_[site_index].emplace(extend_right_environment<EnvironmentStorage>(*right_[site], mps_->site(site_index),
                                                                                  mpo_->site(site_index)));
        }
      }
      return *right_[bond];
    }
 
    void build_all()
    {
      static_cast<void>(this->left_environment(this->size()));
      static_cast<void>(this->right_environment(0));
    }
 
  private:
    void validate() const
    {
      if (mps_->size() != mpo_->size())
        throw std::invalid_argument("MPO environment cache requires equal nonzero chain lengths");
      if (mps_->symmetry() != mpo_->symmetry())
        throw std::invalid_argument("MPO environment cache requires one symmetry");
      for (std::size_t index = 0; index < mps_->size(); ++index)
      {
        auto const& physical = mps_->site(index).domain().template space<1>();
        if (physical != mpo_->site(index).domain().template space<1>() ||
            physical != mpo_->site(index).codomain().template space<1>())
        {
          throw std::invalid_argument("MPO environment cache physical spaces must compare exactly equal");
        }
      }
    }
 
    void validate_bond(std::size_t bond) const
    {
      if (bond > this->size()) throw std::out_of_range("MPO environment cache bond index is out of range");
    }
 
    void synchronize_revisions()
    {
      for (std::size_t site = 0; site < this->size(); ++site)
      {
        std::size_t const mps_revision = mps_->revision(site);
        std::size_t const mpo_revision = mpo_->revision(site);
        if (mps_revision == mps_revisions_[site] && mpo_revision == mpo_revisions_[site]) continue;
 
        for (std::size_t bond = site + 1; bond <= this->size(); ++bond)
          left_[bond].reset();
        for (std::size_t bond = 0; bond <= site; ++bond)
          right_[bond].reset();
        mps_revisions_[site] = mps_revision;
        mpo_revisions_[site] = mpo_revision;
      }
    }
 
    mps_type const* mps_;
    mpo_type const* mpo_;
    std::vector<std::optional<environment_type>> left_;
    std::vector<std::optional<environment_type>> right_;
    std::vector<std::size_t> mps_revisions_;
    std::vector<std::size_t> mpo_revisions_;
};
 
template <class MpsChain, class MpoChain>
MpoEnvironmentCache(MpsChain const&, MpoChain const&, std::size_t,
                    std::size_t) -> MpoEnvironmentCache<MpsChain, MpoChain>;
 
} // namespace uni20::tensor_network