limited-memory-qr.hpp

Go to the documentation of this file.

#include <Eigen/Jacobi>
#include <cstddef>
#include <panoc-alm/util/ringbuffer.hpp>
#include <panoc-alm/util/vec.hpp>
#include <type_traits>
 
namespace pa {
 
class LimitedMemoryQR {
  public:
    LimitedMemoryQR() = default;
 
    LimitedMemoryQR(size_t n, size_t m) : Q(n, m), R(m, m) {}
 
    size_t n() const { return Q.rows(); }
    size_t m() const { return Q.cols(); }
 
    template <class VecV>
    void add_column(const VecV &v) {
        assert(q_idx < m());
 
        auto q = Q.col(q_idx);
        auto r = R.col(r_idx_end);
 
        // Modified Gram-Schmidt to make q orthogonal to Q
        q = v;
        for (size_t i = 0; i < q_idx; ++i) {
            real_t s = Q.col(i).dot(q);
            r(i)     = s;
            q -= s * Q.col(i);
        }
 
        // Compute the norms of orthogonalized q and original v
        real_t norm_q = q.norm();
        real_t norm_v = v.norm();
 
        // If ‖q‖ is significantly smaller than ‖v‖, perform
        // reorthogonalization
        real_t η = 0.7;
        while (norm_q < η * norm_v) {
            ++reorth_count;
            for (size_t i = 0; i < q_idx; ++i) {
                real_t s = Q.col(i).dot(q);
                r(i) += s;
                q -= s * Q.col(i);
            }
            norm_v = norm_q;
            norm_q = q.norm();
        }
 
        // Normalize q such that new matrix (Q q) remains orthogonal (i.e. has
        // orthonormal columns)
        r(q_idx) = norm_q;
        q /= norm_q;
 
        // Increment indices, add a column to Q and R.
        ++q_idx;
        r_idx_end = r_succ(r_idx_end);
    }
 
    void remove_column() {
        assert(q_idx > 0);
 
        // After removing the first colomn of the upper triangular matrix R,
        // it becomes upper Hessenberg. Givens rotations are used to make it
        // triangular again.
        Eigen::JacobiRotation<real_t> G;
        size_t r = 0;                   // row index of R
        size_t c = r_succ(r_idx_start); // column index of R in storage
        while (r < q_idx - 1) {
            // Compute the Givens rotation that makes the subdiagonal element
            // of column c or R zero.
            G.makeGivens(R(r, c), R(r + 1, c), &R(r, c));
            // Apply it to the remaining columns of R.
            // Not the current column, because the diagonal element was updated
            // by the makeGivens function, and the subdiagonal element doesn't
            // have to be set to zero explicitly, it's implicit.
            // Also not the previous columns, because they are already
            // implicitly zero below the diagonal and this rotation wouldn't
            // have any effect there.
            // TODO: can this be sped up by applying it in two blocks instead
            //       of column by column?
            for (size_t cc = r_succ(c); cc != r_idx_end; cc = r_succ(cc))
                R.col(cc).applyOnTheLeft(r, r + 1, G.adjoint());
            // Apply the inverse of the Givens rotation to Q.
            Q.block(0, 0, Q.rows(), q_idx).applyOnTheRight(r, r + 1, G);
            // Advance indices to next diagonal element of R.
            ++r;
            c = r_succ(c);
        }
        // Remove rightmost column of Q, since it corresponds to the bottom row
        // of R, which was set to zero by the Givens rotations
        --q_idx;
        // Remove the first column of R.
        r_idx_start = r_succ(r_idx_start);
    }
 
    template <class VecB, class VecX>
    void solve_col(const VecB &b, VecX &x) const {
        // Iterate over the diagonal of R, starting at the bottom right,
        // this is standard back substitution
        // (recall that R is stored in a circular buffer, so R.col(i) is
        // not the mathematical i-th column)
        auto rev_bgn = ring_reverse_iter().begin();
        auto rev_end = ring_reverse_iter().end();
        auto fwd_end = ring_iter().end();
        for (auto it_d = rev_bgn; it_d != rev_end; ++it_d) {
            // Row/column index of diagonal element of R
            auto [rR, cR] = *it_d;
            // (r is the zero-based mathematical index, c is the index in
            // the circular buffer)
            x(rR) = Q.col(rR).transpose() * b; // Compute rhs Qᵀb
            // In the current row of R, iterate over the elements to the
            // right of the diagonal
            // Iterating from left to right seems to give better results
            for (auto it_c = it_d.forwardit; it_c != fwd_end; ++it_c) {
                auto [rX2, cR2] = *it_c;
                x(rR) -= R(rR, cR2) * x(rX2);
            }
            x(rR) /= R(rR, cR); // Divide by diagonal element
        }
    }
 
    template <class MatB, class MatX>
    void solve(const MatB &B, MatX &X) const {
        assert(B.cols() <= X.cols());
        assert(B.rows() >= Q.rows());
        assert(X.rows() >= Eigen::Index(num_columns()));
        // Each column of the right hand side is solved as an individual system
        for (Eigen::Index cB = 0; cB < B.cols(); ++cB) {
            auto b = B.col(cB);
            auto x = X.col(cB);
            solve_col(b, x);
        }
    }
 
    template <class Derived>
    using solve_ret_t = std::conditional_t<
        Eigen::internal::traits<Derived>::ColsAtCompileTime == 1, vec, mat>;
 
    template <class Derived>
    solve_ret_t<Derived> solve(const Eigen::DenseBase<Derived> &B) {
        solve_ret_t<Derived> X(m(), B.cols());
        solve(B, X);
        return X;
    }
 
    const mat &get_raw_Q() const { return Q; }
    const mat &get_raw_R() const { return R; }
 
    mat get_full_R() const {
        if (r_idx_start == 0)
            return R;
        // Using a permutation matrix here isn't as efficient as rotating the
        // matrix manually, but this function is only used in tests, so it
        // shouldn't matter.
        Eigen::PermutationMatrix<Eigen::Dynamic> P(R.cols());
        P.setIdentity();
        std::rotate(P.indices().data(), P.indices().data() + r_idx_start,
                    P.indices().data() + P.size());
        return R * P;
    }
    mat get_R() const {
        return get_full_R()
            .block(0, 0, q_idx, q_idx)
            .triangularView<Eigen::Upper>();
    }
    mat get_Q() const { return Q.block(0, 0, n(), q_idx); }
 
    void scale_R(real_t scal) {
        for (auto [i, r_idx] : ring_iter())
            R.col(r_idx).topRows(i + 1) *= scal;
    }
 
    size_t get_reorth_count() const { return reorth_count; }
    void clear_reorth_count() { reorth_count = 0; }
 
    void reset() {
        q_idx        = 0;
        r_idx_start  = 0;
        r_idx_end    = 0;
        reorth_count = 0;
    }
 
    void resize(size_t n, size_t m) {
        Q.resize(n, m);
        R.resize(m, m);
        reset();
    }
 
    size_t num_columns() const { return q_idx; }
    size_t ring_head() const { return r_idx_start; }
    size_t ring_tail() const { return r_idx_end; }
    size_t ring_next(size_t i) const { return r_succ(i); }
    size_t ring_prev(size_t i) const { return r_pred(i); }
 
    CircularRange<size_t> ring_iter() const {
        return {q_idx, r_idx_start, r_idx_end, m()};
    }
    ReverseCircularRange<size_t> ring_reverse_iter() const {
        return ring_iter();
    }
 
  private:
    mat Q; 
    mat R; 
 
    size_t q_idx       = 0; 
    size_t r_idx_start = 0; 
    size_t r_idx_end   = 0; 
 
    size_t reorth_count = 0; 
 
    size_t r_succ(size_t i) const { return i + 1 < m() ? i + 1 : 0; }
    size_t r_pred(size_t i) const { return i == 0 ? m() - 1 : i - 1; }
};
 
} // namespace pa