mat: calculate Q elements lazily when calling QR.At

tvkn · kortschak · commit 1f29d7b1d172 · 2024-08-17T06:04:00.000+09:30
When a matrix is very tall, calculating Q will currently allocate a
large Q at the end of the factorisation, even if it is not going to be
used, and a large Q matrix can lead to out of memory issues.

For this reason, Q is never eagerly computed unless explicitly required
to by the user, with QR.ToQ.
To keep fulfilling the Matrix interface, the QR.At method will compute
the requested element only, which only require computing a single row of
Q.
diff --git a/mat/qr.go b/mat/qr.go
@@ -31,8 +31,13 @@ func (qr *QR) Dims() (r, c int) {
 	return qr.qr.Dims()
 }
 
-// At returns the element at row i, column j.
+// At returns the element at row i, column j. At will panic if the receiver
+// does not contain a successful factorization.
 func (qr *QR) At(i, j int) float64 {
+	if !qr.isValid() {
+		panic(badQR)
+	}
+
 	m, n := qr.Dims()
 	if uint(i) >= uint(m) {
 		panic(ErrRowAccess)
@@ -41,13 +46,46 @@ func (qr *QR) At(i, j int) float64 {
 		panic(ErrColAccess)
 	}
 
+	if qr.q == nil || qr.q.IsEmpty() {
+		// Calculate Qi, Q i-th row
+		qi := getFloat64s(m, true)
+		qr.qRowTo(i, qi)
+
+		// Compute QR(i,j)
+		var val float64
+		for k := 0; k <= j; k++ {
+			val += qi[k] * qr.qr.at(k, j)
+		}
+		putFloat64s(qi)
+		return val
+	}
+
 	var val float64
 	for k := 0; k <= j; k++ {
 		val += qr.q.at(i, k) * qr.qr.at(k, j)
 	}
 	return val
 }
 
+// qRowTo extracts the i-th row of the orthonormal matrix Q from a QR
+// decomposition.
+func (qr *QR) qRowTo(i int, dst []float64) {
+	c := blas64.General{
+		Rows:   1,
+		Cols:   len(dst),
+		Stride: len(dst),
+		Data:   dst,
+	}
+	c.Data[i] = 1 // C is the i-th unit vector
+
+	// Construct Qi from the elementary reflectors: Qi = C * (H(1) H(2) ... H(nTau))
+	work := []float64{0}
+	lapack64.Ormqr(blas.Right, blas.NoTrans, qr.qr.mat, qr.tau, c, work, -1)
+	work = getFloat64s(int(work[0]), false)
+	lapack64.Ormqr(blas.Right, blas.NoTrans, qr.qr.mat, qr.tau, c, work, len(work))
+	putFloat64s(work)
+}
+
 // T performs an implicit transpose by returning the receiver inside a
 // Transpose.
 func (qr *QR) T() Matrix {
@@ -98,7 +136,9 @@ func (qr *QR) factorize(a Matrix, norm lapack.MatrixNorm) {
 	lapack64.Geqrf(qr.qr.mat, qr.tau, work, len(work))
 	putFloat64s(work)
 	qr.updateCond(norm)
-	qr.updateQ()
+	if qr.q != nil {
+		qr.q.Reset()
+	}
 }
 
 func (qr *QR) updateQ() {
@@ -192,6 +232,10 @@ func (qr *QR) QTo(dst *Dense) {
 			panic(ErrShape)
 		}
 	}
+
+	if qr.q == nil || qr.q.IsEmpty() {
+		qr.updateQ()
+	}
 	dst.Copy(qr.q)
 }
 
diff --git a/mat/qr_test.go b/mat/qr_test.go
@@ -18,9 +18,11 @@ func TestQR(t *testing.T) {
 	rnd := rand.New(rand.NewSource(1))
 	for _, test := range []struct {
 		m, n int
+		big  bool
 	}{
-		{5, 5},
-		{10, 5},
+		{m: 5, n: 5},
+		{m: 10, n: 5},
+		{m: 1e5, n: 3, big: true}, // Test that very tall matrices do not OoM.
 	} {
 		m := test.m
 		n := test.n
@@ -35,6 +37,13 @@ func TestQR(t *testing.T) {
 
 		var qr QR
 		qr.Factorize(a)
+		if test.big {
+			_ = qr.At(0, 0)     // should not panic, even for big matrices
+			_ = qr.At(m-1, n-1) // should not panic, even for big matrices
+			// We cannot proceed past here for big matrices.
+			continue
+		}
+
 		var q, r Dense
 		qr.QTo(&q)
 
@@ -56,6 +65,23 @@ func TestQR(t *testing.T) {
 		if !EqualApprox(&got, &want, 1e-12) {
 			t.Errorf("QR does not equal original matrix. \nWant: %v\nGot: %v", want, got)
 		}
+
+		// Verify indirect QR.At()
+		got.Reset()
+		got.ReuseAs(m, n)
+		qr.q.Reset() // reset q matrix to force lazy computation
+		for i := 0; i < m; i++ {
+			for j := 0; j < n; j++ {
+				got.set(i, j, qr.At(i, j))
+			}
+		}
+
+		if !EqualApprox(a, &got, 1e-14) {
+			t.Errorf("m=%d,n=%d: A and QR (computed with QR.At()) are not equal", m, n)
+		}
+		if !EqualApprox(a.T(), got.T(), 1e-14) {
+			t.Errorf("m=%d,n=%d: Aᵀ and (QR)ᵀ (computed with QR.At()) are not equal", m, n)
+		}
 	}
 }
 
