MultiDiagonal.h

Go to the documentation of this file.

// Copyright András Vukics 2006–2026. Distributed under the Boost Software License, Version 1.0. (See accompanying file LICENSE.txt)
#pragma once
 
#include "MultiArrayComplex.h"
 
#include <algorithm>
#include <map>
 
 
namespace quantumoperator {
 
using namespace cppqedutils;
 
 
template <size_t RANK>
struct MultiDiagonal
{
  using Dimensions = Extents<RANK>;
  using Offsets = std::array<ptrdiff_t,RANK>;
  using Diagonal = MultiArray<dcomp,RANK>;
 
  using Diagonals = std::map<Offsets,Diagonal>;
 
  MultiDiagonal(const MultiDiagonal&) = delete; MultiDiagonal& operator=(const MultiDiagonal&) = delete;
  MultiDiagonal(MultiDiagonal&&) = default; MultiDiagonal& operator=(MultiDiagonal&&) = default;
 
  explicit MultiDiagonal(auto&&... args) requires std::constructible_from<Diagonals,decltype(args)...> : diagonals{FWD(args)...} {}
 
  friend MultiDiagonal copy(const MultiDiagonal& md)
  {
    MultiDiagonal res;
    for (const auto& [offsets,diag] : md.diagonals) res.diagonals.emplace(offsets,copy(diag));
    return res;
  }
 
  Diagonals diagonals;
 
  void operator () (double, MultiArrayConstView<dcomp,RANK> psi, MultiArrayView<dcomp,RANK> dpsidt) const
  {
    if (diagonals.empty()) return;
 
#ifndef NDEBUG
    if (psi.extents != calculateAndCheckDimensions(*this)) throw std::runtime_error("Mismatch between StateVector and MultiDiagonal dimensions");
    checkExtents(psi,MultiArrayConstView<dcomp,RANK>{dpsidt},"MultiDiagonal application");
#endif // NDEBUG
 
    const size_t strideRows0{dpsidt.strides[0]}, strideCols0{psi.strides[0]};
 
    for (const auto& [offsets,diag] : diagonals) {
 
      if (diag.dataView.empty()) continue; // |o|=D on some axis: a structural zero
 
      // first row o⁻ in dpsidt, first column o⁺ in psi
      dcomp* rows{dpsidt.dataView.data()+dpsidt.offset}; const dcomp* cols{psi.dataView.data()+psi.offset};
      for (auto&& [o,sR,sC] : std::views::zip(offsets,dpsidt.strides,psi.strides)) {
        if (o<0) rows+=sR*size_t(-o); else cols+=sC*size_t(o);
      }
 
      // The loop nest walks the diagonal d in storage order (an owning MultiArray, contiguous, axis 0 fastest), and returns
      // d advanced past the entries it used. len0 is the length of axis 0 if fixed at compile time, otherwise 0.
      // Bounds and strides are copied, and d is passed by value, so that the compiler sees that the stores into dpsidt
      // don't change them, also where it doesn't inline the lambda; otherwise it cannot vectorize.
      const auto loop=[&] (auto n, const auto& self, auto len0, dcomp* r, const dcomp* c, const dcomp* d) -> const dcomp*
      {
        using namespace hana::literals;
        if constexpr (n!=0_c) {
          const size_t extent{diag.extents[n]}, sR{dpsidt.strides[n]}, sC{psi.strides[n]};
          for (size_t k=0; k<extent; ++k) d=self(n-1_c,self,len0,r+k*sR,c+k*sC,d);
          return d;
        }
        else {
          const size_t sR{strideRows0}, sC{strideCols0};
          const auto multiplyAdd=[=] (size_t k) {r[k*sR]+=textbookProduct(d[k],c[k*sC]);};
          if constexpr (decltype(len0)::value) {
            hana::for_each(hana::make_range(hana::size_c<0>,len0),multiplyAdd);
            return d+decltype(len0)::value;
          }
          else {
            const size_t len{diag.extents[0]};
            if (sR==1 && sC==1) for (size_t k=0; k<len; ++k) r[k]+=textbookProduct(d[k],c[k]);
            else for (size_t k=0; k<len; ++k) multiplyAdd(k);
            return d+len;
          }
        }
      };
 
      const auto run=[&] (auto len0) {loop(hana::llong_c<RANK-1>,loop,len0,rows,cols,diag.dataView.data());};
 
      // A short axis 0 is unrolled at compile time: a vectorized loop over a few elements costs more than it saves.
      // Up to four is where the two cross over (benchmarks/kernels.cc, case E).
      switch (diag.extents[0]) {
        case 1:  run(hana::size_c<1>); break;
        case 2:  run(hana::size_c<2>); break;
        case 3:  run(hana::size_c<3>); break;
        case 4:  run(hana::size_c<4>); break;
        default: run(hana::size_c<0>);
      }
    }
  }
 
 
  template <size_t RANK2>
  friend auto operator* (const MultiDiagonal<RANK>& md1, const MultiDiagonal<RANK2>& md2)
  {
    MultiDiagonal<RANK+RANK2> res;
    for (const auto& [offsets1,diag1] : md1.diagonals) for (const auto& [offsets2,diag2] : md2.diagonals)
      res.diagonals.emplace(concatenate(offsets1,offsets2),directProduct(diag1,diag2));
    return res;
  }
 
 
 
 
  MultiDiagonal& hermitianConjugate()
  {
    Diagonals newDiagonals;
    for (auto& [offsets,diag] : diagonals) {
      Offsets newOffsets;
      for (size_t i=0; i<RANK; ++i) newOffsets[i]=-offsets[i];
      conj(diag);
      newDiagonals.emplace(newOffsets,std::move(diag));
    }
    diagonals.swap(newDiagonals);
    return *this;
  }
 
  MultiDiagonal& dagger() {return hermitianConjugate();}
 
  friend MultiDiagonal hermitianConjugateOf(const MultiDiagonal& md) {MultiDiagonal res(copy(md)); res.hermitianConjugate(); return res;}
 
  friend MultiDiagonal twoTimesRealPartOf(const MultiDiagonal& md) {return md+hermitianConjugateOf(md);}
 
  friend MultiDiagonal twoITimesImagPartOf(const MultiDiagonal& md) {return md-hermitianConjugateOf(md);}
 
 
 
  MultiDiagonal operator-() const {MultiDiagonal res{copy(*this)}; res*=-1.; return res;}
 
  MultiDiagonal operator+() const {return copy(*this);}
 
  MultiDiagonal& operator+=(const MultiDiagonal& md)
  {
#ifndef NDEBUG
    if (auto dim{calculateAndCheckDimensions(*this)}, mdDim{calculateAndCheckDimensions(md)};
      dim!=Dimensions{} && mdDim!=Dimensions{} && dim!=mdDim)
        throw std::runtime_error("Dimension mismatch in addition of MultiDiagonals");
#endif // NDEBUG
    for (const auto& [offsets,diag] : md.diagonals)
      if (auto it=diagonals.find(offsets); it!=diagonals.end())
        for (auto&& [to,from] : std::views::zip(it->second.dataStorage(),diag.dataStorage())) to+=from;
      else diagonals.emplace(offsets,copy(diag));
    return *this;
  }
 
  MultiDiagonal& operator-=(const MultiDiagonal& md) {return operator+=(-md);}
 
  MultiDiagonal& operator*=(scalar auto d)
  {
    for (auto& [offsets,diag] : diagonals) for (dcomp& v : diag.dataStorage()) v*=d;
    return *this;
  }
 
  MultiDiagonal& operator/=(scalar auto d) {(*this)*=1./d; return *this;}
 
  friend auto operator+(const MultiDiagonal& md1, const MultiDiagonal& md2) {MultiDiagonal res{copy(md1)}; res+=md2; return res;}
  friend auto operator-(const MultiDiagonal& md1, const MultiDiagonal& md2) {MultiDiagonal res{copy(md1)}; res-=md2; return res;}
 
  friend auto operator*(scalar auto v, const MultiDiagonal& md) {MultiDiagonal res{copy(md)}; res*=v; return res;}
  friend auto operator/(const MultiDiagonal& md, scalar auto v) {MultiDiagonal res{copy(md)}; res/=v; return res;}
 
 
  friend Dimensions calculateAndCheckDimensions(const MultiDiagonal& md)
  {
    if (md.diagonals.empty()) return {};
 
    auto impliedDimensions = [] (const auto& kv) -> Dimensions {
      auto res{kv.second.extents};
      for (auto&& [v,o] : std::views::zip(res,kv.first)) v += size_t(o<0 ? -o : o);
      return res;
    };
 
    const auto res{impliedDimensions(*md.diagonals.cbegin())};
 
    for (const auto& kv : md.diagonals)
      if (impliedDimensions(kv)!=res) throw std::runtime_error("Dimensions mismatch in MultiDiagonal");
 
    return res;
  }
 
  friend void tag_invoke( const json::value_from_tag&, json::value& jv, const MultiDiagonal& md )
  {
    json::array& arr = jv.emplace_array();
    for (const auto& [offsets,diag] : md.diagonals)
      arr.push_back( json::object{ {"offsets",json::value_from(offsets)}, {"diagonal",json::value_from(diag)} } );
  }
 
};
 
 
template <size_t RANK>
MultiDiagonal<RANK> operator|(const MultiDiagonal<RANK>& a, const MultiDiagonal<RANK>& b)
{
  using MD = MultiDiagonal<RANK>;
 
  if (a.diagonals.empty() || b.diagonals.empty()) return MD{}; // composition with the zero operator
 
  const auto dim = calculateAndCheckDimensions(a);
  if (dim != calculateAndCheckDimensions(b)) throw std::runtime_error("Mismatch in MultiDiagonal composition dimensions");
 
  MD res;
 
  for (const auto& [p,diagA] : a.diagonals) for (const auto& [q,diagB] : b.diagonals) {
 
    typename MD::Offsets s;                     // s = p + q
    Extents<RANK> resExtents, start, box;       // result-diagonal extents; valid window [start, start+box)
    std::array<ptrdiff_t,RANK> iA, iB;          // index shifts into the operand diagonals
 
    bool nonzero=true;
 
    for (size_t i=0; i<RANK; ++i) {
      const ptrdiff_t D=ptrdiff_t(dim[i]), pi=p[i], qi=q[i], si=pi+qi,
        sm = si<0 ? -si : 0, pm = pi<0 ? -pi : 0, qm = qi<0 ? -qi : 0,      // negative parts o⁻
        sa = si<0 ? -si : si, pa = pi<0 ? -pi : pi, qa = qi<0 ? -qi : qi;   // magnitudes |o|
 
      if (sa >= D) {nonzero=false; break;}
 
      const ptrdiff_t ia = sm - pm, ib = sm + pi - qm,
        st = std::max({ptrdiff_t{0}, -ia, -ib}),
        en = std::min({D - sa, D - pa - ia, D - qa - ib});
 
      if (en <= st) {nonzero=false; break;}
 
      s[i]=si; resExtents[i]=size_t(D-sa); start[i]=size_t(st); box[i]=size_t(en-st); iA[i]=ia; iB[i]=ib;
    }
    if (!nonzero) continue;
 
    // find-or-create the result diagonal (zero-initialized); accumulate the windowed products
    auto& resDiag = res.diagonals.try_emplace(s, resExtents, multiarray::zeroInit<dcomp>).first->second;
 
    for (Extents<RANK> t{}; t[0]!=box[0]; incrementMultiIndex(t,box)) {
      Extents<RANK> kRes, kA, kB;
      for (size_t i=0; i<RANK; ++i) {
        kRes[i]=start[i]+t[i];
        kA[i]=size_t(ptrdiff_t(kRes[i])+iA[i]);
        kB[i]=size_t(ptrdiff_t(kRes[i])+iB[i]);
      }
      resDiag(kRes) += diagA(kA)*diagB(kB);
    }
  }
 
  return res;
}
 
 
namespace multidiagonal {
 
MultiDiagonal<1> identity(size_t dim);
 
} // multidiagonal
 
 
} // quantumoperator
 
 
 
inline quantumoperator::MultiDiagonal<1> quantumoperator::multidiagonal::identity(size_t dim)
{
  MultiDiagonal<1> res;
  res.diagonals.emplace(
    MultiDiagonal<1>::Offsets{0},
    MultiDiagonal<1>::Diagonal{ {dim}, [] (size_t e) {return multiarray::StorageType<dcomp>(e,1.);} } );
  return res;
}