mat: add BandCholesky type

mat/cholesky.go

@@ -20,6 +20,11 @@ const (
 var (
 	_ Matrix    = (*Cholesky)(nil)
 	_ Symmetric = (*Cholesky)(nil)
+
+	_ Matrix    = (*BandCholesky)(nil)
+	_ Symmetric = (*BandCholesky)(nil)
+	_ Banded    = (*BandCholesky)(nil)
+	_ SymBanded = (*BandCholesky)(nil)
 )
 
 // Cholesky is a symmetric positive definite matrix represented by its
@@ -100,7 +105,7 @@ func (c *Cholesky) At(i, j int) float64 {
 	return val
 }
 
-// T returns the the receiver, the transpose of a symmetric matrix.
+// T returns the receiver, the transpose of a symmetric matrix.
 func (c *Cholesky) T() Matrix {
 	return c
 }
@@ -264,7 +269,7 @@ func (a *Cholesky) SolveCholTo(dst *Dense, b *Cholesky) error {
 	return nil
 }
 
-// SolveVecTo finds the vector X that solves A * x = b where A is represented
+// SolveVecTo finds the vector x that solves A * x = b where A is represented
 // by the Cholesky decomposition. The result is stored in-place into
 // dst.
 func (c *Cholesky) SolveVecTo(dst *VecDense, b Vector) error {
@@ -697,3 +702,214 @@ func (c *Cholesky) SymRankOne(orig *Cholesky, alpha float64, x Vector) (ok bool)
 func (c *Cholesky) valid() bool {
 	return c.chol != nil && !c.chol.IsEmpty()
 }
+
+// BandCholesky is a symmetric positive-definite band matrix represented by its
+// Cholesky decomposition.
+//
+// Note that this matrix representation is useful for certain operations, in
+// particular finding solutions to linear equations. It is very inefficient at
+// other operations, in particular At is slow.
+//
+// BandCholesky methods may only be called on a value that has been successfully
+// initialized by a call to Factorize that has returned true. Calls to methods
+// of an unsuccessful Cholesky factorization will panic.
+type BandCholesky struct {
+	// The chol pointer must never be retained as a pointer outside the Cholesky
+	// struct, either by returning chol outside the struct or by setting it to
+	// a pointer coming from outside. The same prohibition applies to the data
+	// slice within chol.
+	chol *TriBandDense
+	cond float64
+}
+
+// Factorize calculates the Cholesky decomposition of the matrix A and returns
+// whether the matrix is positive definite. If Factorize returns false, the
+// factorization must not be used.
+func (ch *BandCholesky) Factorize(a SymBanded) (ok bool) {
+	n, k := a.SymBand()
+	if ch.chol == nil {
+		ch.chol = NewTriBandDense(n, k, Upper, nil)
+	} else {
+		ch.chol.Reset()
+		ch.chol.ReuseAsTriBand(n, k, Upper)
+	}
+	copySymBandIntoTriBand(ch.chol, a)
+	cSym := blas64.SymmetricBand{
+		Uplo:   blas.Upper,
+		N:      n,
+		K:      k,
+		Data:   ch.chol.RawTriBand().Data,
+		Stride: ch.chol.RawTriBand().Stride,
+	}
+	_, ok = lapack64.Pbtrf(cSym)
+	if !ok {
+		ch.Reset()
+		return false
+	}
+	work := getFloats(3*n, false)
+	iwork := getInts(n, false)
+	aNorm := lapack64.Lansb(CondNorm, cSym, work)
+	ch.cond = 1 / lapack64.Pbcon(cSym, aNorm, work, iwork)
+	putInts(iwork)
+	putFloats(work)
+	return true
+}
+
+// SolveTo finds the matrix X that solves A * X = B where A is represented by
+// the Cholesky decomposition. The result is stored in-place into dst.
+func (ch *BandCholesky) SolveTo(dst *Dense, b Matrix) error {
+	if !ch.valid() {
+		panic(badCholesky)
+	}
+	br, bc := b.Dims()
+	if br != ch.chol.mat.N {
+		panic(ErrShape)
+	}
+	dst.reuseAsNonZeroed(br, bc)
+	if b != dst {
+		dst.Copy(b)
+	}
+	lapack64.Pbtrs(ch.chol.mat, dst.mat)
+	if ch.cond > ConditionTolerance {
+		return Condition(ch.cond)
+	}
+	return nil
+}
+
+// SolveVecTo finds the vector x that solves A * x = b where A is represented by
+// the Cholesky decomposition. The result is stored in-place into dst.
+func (ch *BandCholesky) SolveVecTo(dst *VecDense, b Vector) error {
+	if !ch.valid() {
+		panic(badCholesky)
+	}
+	n := ch.chol.mat.N
+	if br, bc := b.Dims(); br != n || bc != 1 {
+		panic(ErrShape)
+	}
+	if b, ok := b.(RawVectorer); ok && dst != b {
+		dst.checkOverlap(b.RawVector())
+	}
+	dst.reuseAsNonZeroed(n)
+	if dst != b {
+		dst.CopyVec(b)
+	}
+	lapack64.Pbtrs(ch.chol.mat, dst.asGeneral())
+	if ch.cond > ConditionTolerance {
+		return Condition(ch.cond)
+	}
+	return nil
+}
+
+// Cond returns the condition number of the factorized matrix.
+func (ch *BandCholesky) Cond() float64 {
+	if !ch.valid() {
+		panic(badCholesky)
+	}
+	return ch.cond
+}
+
+// Reset resets the factorization so that it can be reused as the receiver of
+// a dimensionally restricted operation.
+func (ch *BandCholesky) Reset() {
+	if ch.chol != nil {
+		ch.chol.Reset()
+	}
+	ch.cond = math.Inf(1)
+}
+
+// Dims returns the dimensions of the matrix.
+func (ch *BandCholesky) Dims() (r, c int) {
+	if !ch.valid() {
+		panic(badCholesky)
+	}
+	r, c = ch.chol.Dims()
+	return r, c
+}
+
+// At returns the element at row i, column j.
+func (ch *BandCholesky) At(i, j int) float64 {
+	if !ch.valid() {
+		panic(badCholesky)
+	}
+	n, k, _ := ch.chol.TriBand()
+	if uint(i) >= uint(n) {
+		panic(ErrRowAccess)
+	}
+	if uint(j) >= uint(n) {
+		panic(ErrColAccess)
+	}
+
+	if i > j {
+		i, j = j, i
+	}
+	if j-i > k {
+		return 0
+	}
+	var aij float64
+	for k := max(0, j-k); k <= i; k++ {
+		aij += ch.chol.at(k, i) * ch.chol.at(k, j)
+	}
+	return aij
+}
+
+// T returns the receiver, the transpose of a symmetric matrix.
+func (ch *BandCholesky) T() Matrix {
+	return ch
+}
+
+// TBand returns the receiver, the transpose of a symmetric band matrix.
+func (ch *BandCholesky) TBand() Banded {
+	return ch
+}
+
+// Symmetric implements the Symmetric interface and returns the number of rows
+// in the matrix (this is also the number of columns).
+func (ch *BandCholesky) Symmetric() int {
+	n, _ := ch.chol.Triangle()
+	return n
+}
+
+// Bandwidth returns the lower and upper bandwidth values for the matrix.
+// The total bandwidth of the matrix is kl+ku+1.
+func (ch *BandCholesky) Bandwidth() (kl, ku int) {
+	_, k, _ := ch.chol.TriBand()
+	return k, k
+}
+
+// SymBand returns the number of rows/columns in the matrix, and the size of the
+// bandwidth. The total bandwidth of the matrix is 2*k+1.
+func (ch *BandCholesky) SymBand() (n, k int) {
+	n, k, _ = ch.chol.TriBand()
+	return n, k
+}
+
+// IsEmpty returns whether the receiver is empty. Empty matrices can be the
+// receiver for dimensionally restricted operations. The receiver can be emptied
+// using Reset.
+func (ch *BandCholesky) IsEmpty() bool {
+	return ch == nil || ch.chol.IsEmpty()
+}
+
+// Det returns the determinant of the matrix that has been factorized.
+func (ch *BandCholesky) Det() float64 {
+	if !ch.valid() {
+		panic(badCholesky)
+	}
+	return math.Exp(ch.LogDet())
+}
+
+// LogDet returns the log of the determinant of the matrix that has been factorized.
+func (ch *BandCholesky) LogDet() float64 {
+	if !ch.valid() {
+		panic(badCholesky)
+	}
+	var det float64
+	for i := 0; i < ch.chol.mat.N; i++ {
+		det += 2 * math.Log(ch.chol.mat.Data[i*ch.chol.mat.Stride])
+	}
+	return det
+}
+
+func (ch *BandCholesky) valid() bool {
+	return ch.chol != nil && !ch.chol.IsEmpty()
+}