Line data Source code
1 : #include <Eigen/Jacobi>
2 : #include <cstddef>
3 : #include <panoc-alm/util/ringbuffer.hpp>
4 : #include <panoc-alm/util/vec.hpp>
5 : #include <type_traits>
6 :
7 : namespace pa {
8 :
9 : /// Incremental QR factorization using modified Gram-Schmidt with
10 : /// reorthogonalization.
11 : ///
12 : /// Computes A = QR while allowing efficient removal of the first
13 : /// column of A or adding new columns at the end of A.
14 6 : class LimitedMemoryQR {
15 : public:
16 : LimitedMemoryQR() = default;
17 :
18 : /// @param n
19 : /// The size of the vectors, the number of rows of A.
20 : /// @param m
21 : /// The maximum number of columns of A.
22 : ///
23 : /// The maximum dimensions of Q are n×m and the maximum dimensions of R are
24 : /// m×m.
25 12 : LimitedMemoryQR(size_t n, size_t m) : Q(n, m), R(m, m) {}
26 :
27 84 : size_t n() const { return Q.rows(); }
28 404 : size_t m() const { return Q.cols(); }
29 :
30 : /// Add the given column to the right.
31 : template <class VecV>
32 31 : void add_column(const VecV &v) {
33 31 : assert(q_idx < m());
34 :
35 31 : auto q = Q.col(q_idx);
36 31 : auto r = R.col(r_idx_end);
37 :
38 : // Modified Gram-Schmidt to make q orthogonal to Q
39 31 : q = v;
40 75 : for (size_t i = 0; i < q_idx; ++i) {
41 44 : real_t s = Q.col(i).dot(q);
42 44 : r(i) = s;
43 44 : q -= s * Q.col(i);
44 44 : }
45 :
46 : // Compute the norms of orthogonalized q and original v
47 31 : real_t norm_q = q.norm();
48 31 : real_t norm_v = v.norm();
49 :
50 : // If ‖q‖ is significantly smaller than ‖v‖, perform
51 : // reorthogonalization
52 31 : real_t η = 0.7;
53 46 : while (norm_q < η * norm_v) {
54 15 : ++reorth_count;
55 44 : for (size_t i = 0; i < q_idx; ++i) {
56 29 : real_t s = Q.col(i).dot(q);
57 29 : r(i) += s;
58 29 : q -= s * Q.col(i);
59 29 : }
60 15 : norm_v = norm_q;
61 15 : norm_q = q.norm();
62 : }
63 :
64 : // Normalize q such that new matrix (Q q) remains orthogonal (i.e. has
65 : // orthonormal columns)
66 31 : r(q_idx) = norm_q;
67 31 : q /= norm_q;
68 :
69 : // Increment indices, add a column to Q and R.
70 31 : ++q_idx;
71 31 : r_idx_end = r_succ(r_idx_end);
72 31 : }
73 :
74 : /// Remove the leftmost column.
75 15 : void remove_column() {
76 15 : assert(q_idx > 0);
77 :
78 : // After removing the first colomn of the upper triangular matrix R,
79 : // it becomes upper Hessenberg. Givens rotations are used to make it
80 : // triangular again.
81 15 : Eigen::JacobiRotation<real_t> G;
82 15 : size_t r = 0; // row index of R
83 15 : size_t c = r_succ(r_idx_start); // column index of R in storage
84 44 : while (r < q_idx - 1) {
85 : // Compute the Givens rotation that makes the subdiagonal element
86 : // of column c or R zero.
87 29 : G.makeGivens(R(r, c), R(r + 1, c), &R(r, c));
88 : // Apply it to the remaining columns of R.
89 : // Not the current column, because the diagonal element was updated
90 : // by the makeGivens function, and the subdiagonal element doesn't
91 : // have to be set to zero explicitly, it's implicit.
92 : // Also not the previous columns, because they are already
93 : // implicitly zero below the diagonal and this rotation wouldn't
94 : // have any effect there.
95 : // TODO: can this be sped up by applying it in two blocks instead
96 : // of column by column?
97 43 : for (size_t cc = r_succ(c); cc != r_idx_end; cc = r_succ(cc))
98 14 : R.col(cc).applyOnTheLeft(r, r + 1, G.adjoint());
99 : // Apply the inverse of the Givens rotation to Q.
100 29 : Q.block(0, 0, Q.rows(), q_idx).applyOnTheRight(r, r + 1, G);
101 : // Advance indices to next diagonal element of R.
102 29 : ++r;
103 29 : c = r_succ(c);
104 : }
105 : // Remove rightmost column of Q, since it corresponds to the bottom row
106 : // of R, which was set to zero by the Givens rotations
107 15 : --q_idx;
108 : // Remove the first column of R.
109 15 : r_idx_start = r_succ(r_idx_start);
110 15 : }
111 :
112 : /// Solve the least squares problem Ax = b.
113 : template <class VecB, class VecX>
114 8 : void solve_col(const VecB &b, VecX &x) const {
115 : // Iterate over the diagonal of R, starting at the bottom right,
116 : // this is standard back substitution
117 : // (recall that R is stored in a circular buffer, so R.col(i) is
118 : // not the mathematical i-th column)
119 8 : auto rev_bgn = ring_reverse_iter().begin();
120 8 : auto rev_end = ring_reverse_iter().end();
121 8 : auto fwd_end = ring_iter().end();
122 32 : for (auto it_d = rev_bgn; it_d != rev_end; ++it_d) {
123 : // Row/column index of diagonal element of R
124 192 : auto [rR, cR] = *it_d;
125 : // (r is the zero-based mathematical index, c is the index in
126 : // the circular buffer)
127 24 : x(rR) = Q.col(rR).transpose() * b; // Compute rhs Qᵀb
128 : // In the current row of R, iterate over the elements to the
129 : // right of the diagonal
130 : // Iterating from left to right seems to give better results
131 48 : for (auto it_c = it_d.forwardit; it_c != fwd_end; ++it_c) {
132 72 : auto [rX2, cR2] = *it_c;
133 24 : x(rR) -= R(rR, cR2) * x(rX2);
134 24 : }
135 24 : x(rR) /= R(rR, cR); // Divide by diagonal element
136 24 : }
137 8 : }
138 :
139 : /// Solve the least squares problem AX = B.
140 : template <class MatB, class MatX>
141 2 : void solve(const MatB &B, MatX &X) const {
142 2 : assert(B.cols() <= X.cols());
143 2 : assert(B.rows() >= Q.rows());
144 2 : assert(X.rows() >= Eigen::Index(num_columns()));
145 : // Each column of the right hand side is solved as an individual system
146 6 : for (Eigen::Index cB = 0; cB < B.cols(); ++cB) {
147 4 : auto b = B.col(cB);
148 4 : auto x = X.col(cB);
149 4 : solve_col(b, x);
150 4 : }
151 2 : }
152 :
153 : template <class Derived>
154 : using solve_ret_t = std::conditional_t<
155 : Eigen::internal::traits<Derived>::ColsAtCompileTime == 1, vec, mat>;
156 :
157 : /// Solve the least squares problem AX = B.
158 : template <class Derived>
159 1 : solve_ret_t<Derived> solve(const Eigen::DenseBase<Derived> &B) {
160 1 : solve_ret_t<Derived> X(m(), B.cols());
161 1 : solve(B, X);
162 1 : return X;
163 1 : }
164 :
165 : /// Get the full, raw storage for the orthogonal matrix Q.
166 : const mat &get_raw_Q() const { return Q; }
167 : /// Get the full, raw storage for the upper triangular matrix R.
168 : /// The columns of this matrix are permuted because it's stored as a
169 : /// circular buffer for efficiently appending columns to the end and
170 : /// popping columns from the front.
171 : const mat &get_raw_R() const { return R; }
172 :
173 : /// Get the full storage for the upper triangular matrix R but with the
174 : /// columns in the correct order.
175 : /// @note Meant for tests only, creates a permuted copy.
176 38 : mat get_full_R() const {
177 38 : if (r_idx_start == 0)
178 16 : return R;
179 : // Using a permutation matrix here isn't as efficient as rotating the
180 : // matrix manually, but this function is only used in tests, so it
181 : // shouldn't matter.
182 22 : Eigen::PermutationMatrix<Eigen::Dynamic> P(R.cols());
183 22 : P.setIdentity();
184 44 : std::rotate(P.indices().data(), P.indices().data() + r_idx_start,
185 22 : P.indices().data() + P.size());
186 22 : return R * P;
187 38 : }
188 : /// Get the matrix R such that Q times R is the original matrix.
189 : /// @note Meant for tests only, creates a permuted copy.
190 38 : mat get_R() const {
191 38 : return get_full_R()
192 38 : .block(0, 0, q_idx, q_idx)
193 38 : .triangularView<Eigen::Upper>();
194 0 : }
195 : /// Get the matrix Q such that Q times R is the original matrix.
196 : /// @note Meant for tests only, creates a copy.
197 84 : mat get_Q() const { return Q.block(0, 0, n(), q_idx); }
198 :
199 : /// Multiply the matrix R by a scalar.
200 1 : void scale_R(real_t scal) {
201 10 : for (auto [i, r_idx] : ring_iter())
202 3 : R.col(r_idx).topRows(i + 1) *= scal;
203 1 : }
204 :
205 : /// Get the number of MGS reorthogonalizations.
206 : size_t get_reorth_count() const { return reorth_count; }
207 : /// Reset the number of MGS reorthogonalizations.
208 : void clear_reorth_count() { reorth_count = 0; }
209 :
210 : /// Reset all indices, clearing the Q and R matrices.
211 0 : void reset() {
212 0 : q_idx = 0;
213 0 : r_idx_start = 0;
214 0 : r_idx_end = 0;
215 0 : reorth_count = 0;
216 0 : }
217 :
218 : /// Re-allocate storage for a problem with a different size. Causes
219 : /// a @ref reset.
220 : void resize(size_t n, size_t m) {
221 : Q.resize(n, m);
222 : R.resize(m, m);
223 : reset();
224 : }
225 :
226 : /// Get the number of columns that are currently stored.
227 20 : size_t num_columns() const { return q_idx; }
228 : /// Get the head index of the circular buffer (points to the oldest
229 : /// element).
230 : size_t ring_head() const { return r_idx_start; }
231 : /// Get the tail index of the circular buffer (points to one past the most
232 : /// recent element).
233 8 : size_t ring_tail() const { return r_idx_end; }
234 : /// Get the next index in the circular buffer.
235 : size_t ring_next(size_t i) const { return r_succ(i); }
236 : /// Get the previous index in the circular buffer.
237 : size_t ring_prev(size_t i) const { return r_pred(i); }
238 :
239 : /// Get iterators in the circular buffer.
240 33 : CircularRange<size_t> ring_iter() const {
241 33 : return {q_idx, r_idx_start, r_idx_end, m()};
242 : }
243 : /// Get reverse iterators in the circular buffer.
244 16 : ReverseCircularRange<size_t> ring_reverse_iter() const {
245 16 : return ring_iter();
246 : }
247 :
248 : private:
249 : mat Q; ///< Storage for orthogonal factor Q.
250 : mat R; ///< Storage for upper triangular factor R.
251 :
252 6 : size_t q_idx = 0; ///< Number of columns of Q being stored.
253 6 : size_t r_idx_start = 0; ///< Index of the first column of R.
254 6 : size_t r_idx_end = 0; ///< Index of the one-past-last column of R.
255 :
256 6 : size_t reorth_count = 0; ///< Number of MGS reorthogonalizations.
257 :
258 : /// Get the next index in the circular storage for R.
259 266 : size_t r_succ(size_t i) const { return i + 1 < m() ? i + 1 : 0; }
260 : /// Get the previous index in the circular storage for R.
261 : size_t r_pred(size_t i) const { return i == 0 ? m() - 1 : i - 1; }
262 : };
263 :
264 : } // namespace pa
|
