Bccb.h

// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
// SPDX-FileCopyrightText: The Eigen Authors
// SPDX-License-Identifier: MPL-2.0
 
// References:
//  [1] P. J. Davis, "Circulant Matrices", Wiley, 1979. Diagonalization of
//      circulant and block circulant matrices by the DFT; the closed-form
//      eigenstructure used by eigenvalues()/eigenvectors(), and the SVD and
//      pseudo-inverse below, follow from it by taking moduli/phases of the
//      eigenvalues.
//  [2] R. H. Chan and X.-Q. Jin, "An Introduction to Iterative Toeplitz
//      Solvers", SIAM, 2007. BCCB matrices are diagonalized by the 2-D DFT
//      F_{n1} (x) F_{n2}; the FFT-based products and solves below, and the use
//      of BCCB operators as preconditioners for two-level Toeplitz (BTTB)
//      systems, follow this reference.
//  [3] G. H. Golub and C. F. Van Loan, "Matrix Computations", 4th ed., Johns
//      Hopkins University Press, 2013, chapter 5.4 (numerical rank conventions).
//  [4] J. J. Dongarra, J. R. Bunch, C. B. Moler and G. W. Stewart, "LINPACK
//      Users' Guide", SIAM, 1979. determinant()'s balanced accumulation follows
//      the convention of its xGEDI routines, which return determinants as a
//      (fraction, exponent) pair to avoid spurious overflow/underflow.
//  [5] P. H. Sterbenz, "Floating-Point Computation", Prentice-Hall, 1974.
//      Scaling by a power of two is exact, the property the balanced
//      accumulation relies on.
 
#ifndef EIGEN_STRUCTURED_BCCB_H
#define EIGEN_STRUCTURED_BCCB_H
 
// IWYU pragma: private
#include "./InternalHeaderCheck.h"
 
namespace Eigen {
 
template <typename Scalar_, int BlockSize_ = Dynamic, int NumBlocks_ = Dynamic>
class Bccb;
 
namespace internal {

constexpr int bccb_dim(int a, int b) { return (a == Dynamic || b == Dynamic) ? Dynamic : a * b; }
 
template <typename Scalar_, int BlockSize_, int NumBlocks_>
struct traits<Bccb<Scalar_, BlockSize_, NumBlocks_>> {
  using Scalar = Scalar_;
  using StorageKind = Dense;
  using XprKind = MatrixXpr;
  using StorageIndex = int;
  static constexpr int RowsAtCompileTime = bccb_dim(BlockSize_, NumBlocks_);
  static constexpr int ColsAtCompileTime = RowsAtCompileTime;
  static constexpr int MaxRowsAtCompileTime = RowsAtCompileTime;
  static constexpr int MaxColsAtCompileTime = RowsAtCompileTime;
  // Deliberately no NestByRefBit: transpose(), conjugate() and adjoint() return
  // owning temporaries, so Product must nest the operator by value for a
  // delayed-evaluated product expression to keep its left factor alive. The copy
  // is O(N), negligible against the O(N log N) product evaluation.
  static constexpr unsigned int Flags = 0;
};
 
template <typename Scalar_, int BlockSize_, int NumBlocks_>
struct evaluator_traits<Bccb<Scalar_, BlockSize_, NumBlocks_>> {
  using Kind = IndexBased;
  using Shape = StructuredShape;
};
 
}  // namespace internal

template <typename Scalar_, int BlockSize_, int NumBlocks_>
class Bccb : public EigenBase<Bccb<Scalar_, BlockSize_, NumBlocks_>> {
  enum DenseAssignment { SetAssignment, AddAssignment, SubAssignment };
 
 public:
  using Scalar = Scalar_;
  using RealScalar = typename NumTraits<Scalar>::Real;
  using StorageIndex = int;
  using Complex = std::complex<RealScalar>;
  // The two-level structure keys on the *column-major* flattening of the n2 x n1
  // generating array and symbol (entry (k2, k1) is flat index k1*n2 + k2), so
  // every internal 2-D type is pinned to ColMajor explicitly: the semantics must
  // not change under EIGEN_DEFAULT_TO_ROW_MAJOR. The single-row special case
  // only satisfies Eigen's rule that 1 x n matrices be row-major; with one row
  // the two orders coincide.
  using GeneratorType =
      Matrix<Scalar, BlockSize_, NumBlocks_, (BlockSize_ == 1 && NumBlocks_ != 1) ? int(RowMajor) : int(ColMajor)>;
  using ComplexArray = Matrix<Complex, Dynamic, Dynamic, ColMajor>;
  using ComplexVector = Matrix<Complex, Dynamic, 1>;
  using RealVector = Matrix<RealScalar, Dynamic, 1>;
  using RealArray = Matrix<RealScalar, Dynamic, Dynamic, ColMajor>;
  using ComplexMatrix = Matrix<Complex, Dynamic, Dynamic, ColMajor>;
 
  static constexpr int RowsAtCompileTime = internal::bccb_dim(BlockSize_, NumBlocks_);
  static constexpr int ColsAtCompileTime = RowsAtCompileTime;
  static constexpr int MaxRowsAtCompileTime = RowsAtCompileTime;
  static constexpr int MaxColsAtCompileTime = RowsAtCompileTime;
  static constexpr int SizeAtCompileTime = internal::size_at_compile_time(RowsAtCompileTime, ColsAtCompileTime);
  static constexpr int MaxSizeAtCompileTime = SizeAtCompileTime;
  static constexpr bool IsRowMajor = false;
  // Deliberately no IsVectorAtCompileTime: Ref<const Bccb>'s default StrideType
  // argument reads it, so its absence makes internal::is_ref_compatible SFINAE to
  // false and keeps the iterative solvers on their matrix-free path.

  template <typename Derived>
  explicit Bccb(const MatrixBase<Derived>& generator) : m_g(generator) {
    eigen_assert(m_g.rows() > 0 && m_g.cols() > 0 && "Bccb generator must be non-empty");
    if (m_g.size() > internal::structured_direct_threshold()) {
      m_symbol = computeSymbol();
      m_prodSymbol = computeProdSymbol(m_g);
    }
  }
 
  EIGEN_DEVICE_FUNC Index rows() const { return m_g.size(); }
  EIGEN_DEVICE_FUNC Index cols() const { return m_g.size(); }

  Index blockSize() const { return m_g.rows(); }
  Index numBlocks() const { return m_g.cols(); }
  const GeneratorType& generator() const { return m_g; }

  ComplexArray symbol() const { return m_symbol.size() > 0 ? m_symbol : computeSymbol(); }

  Scalar coeff(Index row, Index col) const {
    const Index n2 = blockSize();
    Index k2 = row % n2 - col % n2;
    if (k2 < 0) k2 += n2;
    Index k1 = row / n2 - col / n2;
    if (k1 < 0) k1 += numBlocks();
    return m_g.coeff(k2, k1);
  }

  Bccb transpose() const { return Bccb(reverse2(m_g), reverse2(m_symbol), reverse2(m_prodSymbol)); }

  Bccb conjugate() const {
    return Bccb(m_g.conjugate(), reverse2(m_symbol).conjugate(), reverse2(m_prodSymbol).conjugate());
  }

  Bccb adjoint() const { return Bccb(reverse2(m_g).conjugate(), m_symbol.conjugate(), m_prodSymbol.conjugate()); }

  template <typename Rhs>
  Matrix<Scalar, RowsAtCompileTime, Rhs::ColsAtCompileTime> solve(const MatrixBase<Rhs>& b) const {
    EIGEN_STATIC_ASSERT(RowsAtCompileTime == Dynamic || Rhs::RowsAtCompileTime == Dynamic ||
                            int(RowsAtCompileTime) == int(Rhs::RowsAtCompileTime),
                        YOU_MIXED_MATRICES_OF_DIFFERENT_SIZES)
    const Index N = rows();
    eigen_assert(b.rows() == N && "right-hand side has the wrong number of rows");
    const ComplexArray s = symbol();
    const RealArray mods = s.cwiseAbs();
    const ComplexArray sinv = internal::structured_pinv_symbol(s, mods, internal::structured_rank_threshold(s, mods));
    Matrix<Scalar, RowsAtCompileTime, Rhs::ColsAtCompileTime> x(N, b.cols());
    applySymbol(x, sinv, b.derived(), Scalar(1), /*accumulate=*/false);
    return x;
  }

  Index rank() const {
    const ComplexArray s = symbol();
    const RealArray mods = s.cwiseAbs();
    return (!(mods.array() < internal::structured_rank_threshold(s, mods))).count();  // NaN entries count as non-zero
  }

  Bccb inverse() const {
    ComplexArray sinv = symbol().cwiseInverse();
    const GeneratorType ginv = generatorFromSymbol(sinv);
    const bool cache = rows() > internal::structured_direct_threshold();
    return Bccb(ginv, cache ? sinv : ComplexArray(), cache ? computeProdSymbol(ginv) : ComplexArray());
  }

  Scalar determinant() const {
    const ComplexArray s = symbol();
    Complex det(1);
    Index exponent = 0;
    for (Index k1 = 0; k1 < s.cols(); ++k1)
      for (Index k2 = 0; k2 < s.rows(); ++k2)
        det = internal::structured_balance(Complex(det * internal::structured_balance(s(k2, k1), exponent)), exponent);
    return toScalar(internal::structured_ldexp_clamped(det, exponent), std::is_same<RealScalar, Scalar>());
  }

  ComplexVector eigenvalues() const { return symbol().reshaped(); }

  ComplexMatrix eigenvectors() const {
    const Index n2 = blockSize(), n1 = numBlocks(), N = rows();
    ComplexMatrix F(N, N);
    for (Index f1 = 0; f1 < n1; ++f1)
      for (Index f2 = 0; f2 < n2; ++f2) fourierColumn(F, f1, f2, f1 * n2 + f2);
    return F;
  }

  RealVector singularValues() const {
    const ComplexVector s = eigenvalues();
    const RealVector mods = s.cwiseAbs();
    const std::vector<Index> perm = internal::structured_svd_permutation(mods);
    RealVector sv(s.size());
    for (Index t = 0; t < s.size(); ++t) sv[t] = mods[perm[t]];
    return sv;
  }

  ComplexMatrix matrixU() const {
    const Index n2 = blockSize(), N = rows();
    const ComplexVector s = eigenvalues();
    const RealVector mods = s.cwiseAbs();
    const std::vector<Index> perm = internal::structured_svd_permutation(mods);
    ComplexMatrix U(N, N);
    for (Index t = 0; t < N; ++t) {
      fourierColumn(U, perm[t] / n2, perm[t] % n2, t);
      const RealScalar a = mods[perm[t]];
      if (a > RealScalar(0)) U.col(t) *= s[perm[t]] / a;
    }
    return U;
  }

  ComplexMatrix matrixV() const {
    const Index n2 = blockSize(), N = rows();
    const ComplexVector s = eigenvalues();
    const std::vector<Index> perm = internal::structured_svd_permutation(RealVector(s.cwiseAbs()));
    ComplexMatrix V(N, N);
    for (Index t = 0; t < N; ++t) fourierColumn(V, perm[t] / n2, perm[t] % n2, t);
    return V;
  }

  template <typename Dest>
  void evalTo(Dest& dst) const {
    assignTo<SetAssignment>(dst);
  }

  template <typename Dest>
  void addTo(Dest& dst) const {
    assignTo<AddAssignment>(dst);
  }

  template <typename Dest>
  void subTo(Dest& dst) const {
    assignTo<SubAssignment>(dst);
  }

  template <typename Rhs>
  Product<Bccb, Rhs> operator*(const MatrixBase<Rhs>& x) const {
    EIGEN_STATIC_ASSERT(ColsAtCompileTime == Dynamic || Rhs::RowsAtCompileTime == Dynamic ||
                            int(ColsAtCompileTime) == int(Rhs::RowsAtCompileTime),
                        INVALID_MATRIX_PRODUCT)
    eigen_assert(x.rows() == cols() && "invalid product: dimensions do not match");
    return Product<Bccb, Rhs>(*this, x.derived());
  }

  template <typename Dest, typename Rhs, typename ProductScalar>
  void addProduct(Dest& dst, const Rhs& rhs, const ProductScalar& alpha) const {
    const Index N = rows();
    eigen_assert(rhs.rows() == N && "invalid product: dimensions do not match");
    if (N <= internal::structured_direct_threshold()) {
      directProduct(dst, rhs, alpha);
      return;
    }
    // Products use the padded embedding symbol when a block dimension is not
    // 5-smooth (see computeProdSymbol()); the exact-size transform of such a
    // dimension runs through kissfft's quadratic generic butterfly.
    const ComplexArray& s = m_prodSymbol.size() > 0 ? m_prodSymbol : m_symbol;
    applySymbol(dst, s, rhs, alpha, /*accumulate=*/true);
  }
 
 private:
  template <DenseAssignment Assignment, typename Dest>
  void assignTo(Dest& dst) const {
    const Index n2 = blockSize(), n1 = numBlocks();
    for (Index c1 = 0; c1 < n1; ++c1)
      for (Index j2 = 0; j2 < n2; ++j2) {
        auto col = dst.col(c1 * n2 + j2);
        for (Index b1 = 0; b1 < n1; ++b1) {
          Index k1 = b1 - c1;
          if (k1 < 0) k1 += n1;
          EIGEN_IF_CONSTEXPR (Assignment == SetAssignment) {
            col.segment(b1 * n2, j2) = m_g.col(k1).tail(j2);
            col.segment(b1 * n2 + j2, n2 - j2) = m_g.col(k1).head(n2 - j2);
          } else EIGEN_IF_CONSTEXPR (Assignment == AddAssignment) {
            col.segment(b1 * n2, j2) += m_g.col(k1).tail(j2);
            col.segment(b1 * n2 + j2, n2 - j2) += m_g.col(k1).head(n2 - j2);
          } else {
            col.segment(b1 * n2, j2) -= m_g.col(k1).tail(j2);
            col.segment(b1 * n2 + j2, n2 - j2) -= m_g.col(k1).head(n2 - j2);
          }
        }
      }
  }

  Bccb(const GeneratorType& g, const ComplexArray& symbol, const ComplexArray& prodSymbol)
      : m_g(g), m_symbol(symbol), m_prodSymbol(prodSymbol) {}

  ComplexArray computeProdSymbol(const GeneratorType& g) const {
    const Index n2 = g.rows(), n1 = g.cols();
    const Index p2 = internal::fft_next_good_size(n2) == n2 ? n2 : internal::fft_next_good_size(2 * n2 - 1);
    const Index p1 = internal::fft_next_good_size(n1) == n1 ? n1 : internal::fft_next_good_size(2 * n1 - 1);
    if (p2 == n2 && p1 == n1) return ComplexArray();
    ComplexArray Ge = ComplexArray::Zero(p2, p1);
    Ge.topLeftCorner(n2, n1) = g.template cast<Complex>();
    if (p2 != n2) Ge.bottomLeftCorner(n2 - 1, n1) = g.bottomRows(n2 - 1).template cast<Complex>();
    if (p1 != n1) Ge.topRightCorner(n2, n1 - 1) = g.rightCols(n1 - 1).template cast<Complex>();
    if (p2 != n2 && p1 != n1)
      Ge.bottomRightCorner(n2 - 1, n1 - 1) = g.bottomRightCorner(n2 - 1, n1 - 1).template cast<Complex>();
    fft2(Ge);
    return Ge;
  }

  template <typename Dest, typename Rhs, typename ProductScalar>
  void directProductColumn(Dest& dst, const Rhs& rhs, Index k, const ProductScalar& alpha) const {
    const Index N = rows();
    // A unit alpha must not multiply: even the identity complex scalar (1,0)
    // pollutes an (Inf,0) value with NaN through the 0*Inf cross term.
    const bool unitAlpha = alpha == ProductScalar(1);
    for (Index i = 0; i < N; ++i) {
      ProductScalar acc(0);
      for (Index j = 0; j < N; ++j) acc += coeff(i, j) * rhs.coeff(j, k);
      dst.coeffRef(i, k) += unitAlpha ? acc : ProductScalar(alpha * acc);
    }
  }

  template <typename Dest, typename Rhs, typename ProductScalar>
  void directProduct(Dest& dst, const Rhs& rhs, const ProductScalar& alpha) const {
    for (Index k = 0; k < rhs.cols(); ++k) directProductColumn(dst, rhs, k, alpha);
  }

  template <typename Dest, typename Rhs, typename ProductScalar>
  void applySymbol(Dest& dst, const ComplexArray& s, const Rhs& rhs, const ProductScalar& alpha,
                   bool accumulate) const {
    const Index n2 = blockSize(), n1 = numBlocks();
    const Index p2 = s.rows(), p1 = s.cols();  // the symbol's transform grid: (n2, n1) or the padded product grid
    const bool padded = p2 != n2 || p1 != n1;
    ComplexVector xc(rows());
    ComplexArray X(p2, p1), Xn;  // Xn: the leading-block extraction of a padded grid
    for (Index k = 0; k < rhs.cols(); ++k) {
      // reshaped() defaults to column-major traversal, the flattening the
      // two-level structure keys on, regardless of EIGEN_DEFAULT_TO_ROW_MAJOR.
      // Plain objects keep both reshapes linear.
      xc = rhs.col(k).template cast<Complex>();
      if (padded) {
        X.setZero();
        X.topLeftCorner(n2, n1) = xc.reshaped(n2, n1);
      } else {
        X = xc.reshaped(n2, n1);
      }
      fft2(X);
      X.array() *= s.array();
      ifft2(X);
      if (padded) Xn = X.topLeftCorner(n2, n1);
      const auto out = internal::structured_scalar_part_impl<ProductScalar>::run((padded ? Xn : X).reshaped());
      if (accumulate)
        dst.col(k) += alpha * out;
      else
        dst.col(k) = alpha * out;
    }
  }

  template <bool Inverse>
  void transform2(ComplexArray& X) const {
    auto&& fft = internal::structured_fft_engine<RealScalar>();
    const Index n2 = X.rows(), n1 = X.cols();
    // Both passes share one output buffer (resized by the engine) and one packed
    // input buffer, allocated at most twice for the whole transform. A column of
    // the column-major grid is already packed and goes to the engine as is; a row
    // is strided, and the engine would pack it into a fresh temporary per call.
    ComplexVector tmp, rowv;
    if (n2 > 1) {
      for (Index k1 = 0; k1 < n1; ++k1) {
        EIGEN_IF_CONSTEXPR (Inverse) {
          fft.inv(tmp, X.col(k1), n2);
        } else {
          fft.fwd(tmp, X.col(k1), n2);
        }
        X.col(k1) = tmp;
      }
    }
    if (n1 > 1) {
      for (Index k2 = 0; k2 < n2; ++k2) {
        rowv = X.row(k2).transpose();
        EIGEN_IF_CONSTEXPR (Inverse) {
          fft.inv(tmp, rowv, n1);
        } else {
          fft.fwd(tmp, rowv, n1);
        }
        X.row(k2) = tmp.transpose();
      }
    }
  }

  void fft2(ComplexArray& X) const { transform2<false>(X); }

  void ifft2(ComplexArray& X) const { transform2<true>(X); }

  GeneratorType generatorFromSymbol(ComplexArray symbol) const {
    ifft2(symbol);
    GeneratorType generator = internal::structured_scalar_part_impl<Scalar>::run(symbol);
    return generator;
  }

  ComplexArray computeSymbol() const {
    ComplexArray s = m_g.template cast<Complex>();
    fft2(s);
    return s;
  }

  template <typename MatType>
  static MatType reverse2(const MatType& M) {
    const Index r = M.rows(), c = M.cols();
    MatType R(r, c);
    if (M.size() == 0) return R;
    R(0, 0) = M(0, 0);
    if (c > 1) R.row(0).tail(c - 1) = M.row(0).tail(c - 1).reverse();
    if (r > 1) R.col(0).tail(r - 1) = M.col(0).tail(r - 1).reverse();
    if (r > 1 && c > 1) R.bottomRightCorner(r - 1, c - 1) = M.bottomRightCorner(r - 1, c - 1).reverse();
    return R;
  }

  void fourierColumn(ComplexMatrix& F, Index f1, Index f2, Index dstCol) const {
    const Index n2 = blockSize(), n1 = numBlocks();
    const RealScalar scale = RealScalar(1) / numext::sqrt(RealScalar(rows()));
    ComplexVector w2(n2);
    Index jf = 0;  // i2 * f2 mod n2
    for (Index i2 = 0; i2 < n2; ++i2) {
      w2[i2] = std::polar(scale, RealScalar(2 * EIGEN_PI) * RealScalar(jf) / RealScalar(n2));
      jf += f2;
      if (jf >= n2) jf -= n2;
    }
    Index bf = 0;  // b1 * f1 mod n1
    for (Index b1 = 0; b1 < n1; ++b1) {
      const Complex w1 = std::polar(RealScalar(1), RealScalar(2 * EIGEN_PI) * RealScalar(bf) / RealScalar(n1));
      F.col(dstCol).segment(b1 * n2, n2) = w1 * w2;
      bf += f1;
      if (bf >= n1) bf -= n1;
    }
  }

  static Scalar toScalar(const Complex& z, std::true_type /*scalar_is_real*/) { return numext::real(z); }
  static Scalar toScalar(const Complex& z, std::false_type /*scalar_is_real*/) { return z; }
 
  GeneratorType m_g;
  ComplexArray m_symbol;
  // The padded embedding symbol products use when a block dimension is not
  // 5-smooth; empty otherwise. See computeProdSymbol().
  ComplexArray m_prodSymbol;
};

template <typename Derived>
Bccb<typename Derived::Scalar, Derived::RowsAtCompileTime, Derived::ColsAtCompileTime> makeBccb(
    const MatrixBase<Derived>& generator) {
  return Bccb<typename Derived::Scalar, Derived::RowsAtCompileTime, Derived::ColsAtCompileTime>(generator);
}
 
namespace internal {
 
template <typename Scalar_, int BlockSize_, int NumBlocks_, typename Rhs, int ProductTag>
struct generic_product_impl<Bccb<Scalar_, BlockSize_, NumBlocks_>, Rhs, StructuredShape, DenseShape, ProductTag>
    : structured_product_impl<Bccb<Scalar_, BlockSize_, NumBlocks_>, Rhs> {};
 
}  // namespace internal
 
}  // namespace Eigen
 
#endif  // EIGEN_STRUCTURED_BCCB_H