418 lines
10 KiB
Go
418 lines
10 KiB
Go
// Copyright ©2020 The Gonum Authors. All rights reserved.
|
||
// Use of this source code is governed by a BSD-style
|
||
// license that can be found in the LICENSE file.
|
||
|
||
package mat
|
||
|
||
import (
|
||
"math"
|
||
|
||
"gonum.org/v1/gonum/blas"
|
||
"gonum.org/v1/gonum/blas/blas64"
|
||
"gonum.org/v1/gonum/internal/asm/f64"
|
||
"gonum.org/v1/gonum/lapack/lapack64"
|
||
)
|
||
|
||
var (
|
||
tridiagDense *Tridiag
|
||
_ Matrix = tridiagDense
|
||
_ allMatrix = tridiagDense
|
||
_ denseMatrix = tridiagDense
|
||
_ Banded = tridiagDense
|
||
_ MutableBanded = tridiagDense
|
||
_ RawTridiagonaler = tridiagDense
|
||
)
|
||
|
||
// A RawTridiagonaler can return a lapack64.Tridiagonal representation of the
|
||
// receiver. Changes to the elements of DL, D, DU in lapack64.Tridiagonal will
|
||
// be reflected in the original matrix, changes to the N field will not.
|
||
type RawTridiagonaler interface {
|
||
RawTridiagonal() lapack64.Tridiagonal
|
||
}
|
||
|
||
// Tridiag represents a tridiagonal matrix by its three diagonals.
|
||
type Tridiag struct {
|
||
mat lapack64.Tridiagonal
|
||
}
|
||
|
||
// NewTridiag creates a new n×n tridiagonal matrix with the first sub-diagonal
|
||
// in dl, the main diagonal in d and the first super-diagonal in du. If all of
|
||
// dl, d, and du are nil, new backing slices will be allocated for them. If dl
|
||
// and du have length n-1 and d has length n, they will be used as backing
|
||
// slices, and changes to the elements of the returned Tridiag will be reflected
|
||
// in dl, d, du. If neither of these is true, NewTridiag will panic.
|
||
func NewTridiag(n int, dl, d, du []float64) *Tridiag {
|
||
if n <= 0 {
|
||
if n == 0 {
|
||
panic(ErrZeroLength)
|
||
}
|
||
panic(ErrNegativeDimension)
|
||
}
|
||
if dl != nil || d != nil || du != nil {
|
||
if len(dl) != n-1 || len(d) != n || len(du) != n-1 {
|
||
panic(ErrShape)
|
||
}
|
||
} else {
|
||
d = make([]float64, n)
|
||
if n > 1 {
|
||
dl = make([]float64, n-1)
|
||
du = make([]float64, n-1)
|
||
}
|
||
}
|
||
return &Tridiag{
|
||
mat: lapack64.Tridiagonal{
|
||
N: n,
|
||
DL: dl,
|
||
D: d,
|
||
DU: du,
|
||
},
|
||
}
|
||
}
|
||
|
||
// Dims returns the number of rows and columns in the matrix.
|
||
func (a *Tridiag) Dims() (r, c int) {
|
||
return a.mat.N, a.mat.N
|
||
}
|
||
|
||
// Bandwidth returns 1, 1 - the upper and lower bandwidths of the matrix.
|
||
func (a *Tridiag) Bandwidth() (kl, ku int) {
|
||
return 1, 1
|
||
}
|
||
|
||
// T performs an implicit transpose by returning the receiver inside a Transpose.
|
||
func (a *Tridiag) T() Matrix {
|
||
// An alternative would be to return the receiver with DL,DU swapped; the
|
||
// untranspose function would then always return false. With Transpose the
|
||
// diagonal swapping will be done in tridiagonal routines in lapack like
|
||
// lapack64.Gtsv or gonum.Dlagtm based on the trans parameter.
|
||
return Transpose{a}
|
||
}
|
||
|
||
// TBand performs an implicit transpose by returning the receiver inside a
|
||
// TransposeBand.
|
||
func (a *Tridiag) TBand() Banded {
|
||
// An alternative would be to return the receiver with DL,DU swapped; see
|
||
// explanation in T above.
|
||
return TransposeBand{a}
|
||
}
|
||
|
||
// RawTridiagonal returns the underlying lapack64.Tridiagonal used by the
|
||
// receiver. Changes to elements in the receiver following the call will be
|
||
// reflected in the returned matrix.
|
||
func (a *Tridiag) RawTridiagonal() lapack64.Tridiagonal {
|
||
return a.mat
|
||
}
|
||
|
||
// SetRawTridiagonal sets the underlying lapack64.Tridiagonal used by the
|
||
// receiver. Changes to elements in the receiver following the call will be
|
||
// reflected in the input.
|
||
func (a *Tridiag) SetRawTridiagonal(mat lapack64.Tridiagonal) {
|
||
a.mat = mat
|
||
}
|
||
|
||
// IsEmpty returns whether the receiver is empty. Empty matrices can be the
|
||
// receiver for size-restricted operations. The receiver can be zeroed using
|
||
// Reset.
|
||
func (a *Tridiag) IsEmpty() bool {
|
||
return a.mat.N == 0
|
||
}
|
||
|
||
// Reset empties the matrix so that it can be reused as the receiver of a
|
||
// dimensionally restricted operation.
|
||
//
|
||
// Reset should not be used when the matrix shares backing data. See the Reseter
|
||
// interface for more information.
|
||
func (a *Tridiag) Reset() {
|
||
a.mat.N = 0
|
||
a.mat.DL = a.mat.DL[:0]
|
||
a.mat.D = a.mat.D[:0]
|
||
a.mat.DU = a.mat.DU[:0]
|
||
}
|
||
|
||
// CloneFromTridiag makes a copy of the input Tridiag into the receiver,
|
||
// overwriting the previous value of the receiver. CloneFromTridiag does not
|
||
// place any restrictions on receiver shape.
|
||
func (a *Tridiag) CloneFromTridiag(from *Tridiag) {
|
||
n := from.mat.N
|
||
switch n {
|
||
case 0:
|
||
panic(ErrZeroLength)
|
||
case 1:
|
||
a.mat = lapack64.Tridiagonal{
|
||
N: 1,
|
||
DL: use(a.mat.DL, 0),
|
||
D: use(a.mat.D, 1),
|
||
DU: use(a.mat.DU, 0),
|
||
}
|
||
a.mat.D[0] = from.mat.D[0]
|
||
default:
|
||
a.mat = lapack64.Tridiagonal{
|
||
N: n,
|
||
DL: use(a.mat.DL, n-1),
|
||
D: use(a.mat.D, n),
|
||
DU: use(a.mat.DU, n-1),
|
||
}
|
||
copy(a.mat.DL, from.mat.DL)
|
||
copy(a.mat.D, from.mat.D)
|
||
copy(a.mat.DU, from.mat.DU)
|
||
}
|
||
}
|
||
|
||
// DiagView returns the diagonal as a matrix backed by the original data.
|
||
func (a *Tridiag) DiagView() Diagonal {
|
||
return &DiagDense{
|
||
mat: blas64.Vector{
|
||
N: a.mat.N,
|
||
Data: a.mat.D[:a.mat.N],
|
||
Inc: 1,
|
||
},
|
||
}
|
||
}
|
||
|
||
// Zero sets all of the matrix elements to zero.
|
||
func (a *Tridiag) Zero() {
|
||
zero(a.mat.DL)
|
||
zero(a.mat.D)
|
||
zero(a.mat.DU)
|
||
}
|
||
|
||
// Trace returns the trace of the matrix.
|
||
//
|
||
// Trace will panic with ErrZeroLength if the matrix has zero size.
|
||
func (a *Tridiag) Trace() float64 {
|
||
if a.IsEmpty() {
|
||
panic(ErrZeroLength)
|
||
}
|
||
return f64.Sum(a.mat.D)
|
||
}
|
||
|
||
// Norm returns the specified norm of the receiver. Valid norms are:
|
||
//
|
||
// 1 - The maximum absolute column sum
|
||
// 2 - The Frobenius norm, the square root of the sum of the squares of the elements
|
||
// Inf - The maximum absolute row sum
|
||
//
|
||
// Norm will panic with ErrNormOrder if an illegal norm is specified and with
|
||
// ErrZeroLength if the matrix has zero size.
|
||
func (a *Tridiag) Norm(norm float64) float64 {
|
||
if a.IsEmpty() {
|
||
panic(ErrZeroLength)
|
||
}
|
||
return lapack64.Langt(normLapack(norm, false), a.mat)
|
||
}
|
||
|
||
// MulVecTo computes A⋅x or Aᵀ⋅x storing the result into dst.
|
||
func (a *Tridiag) MulVecTo(dst *VecDense, trans bool, x Vector) {
|
||
n := a.mat.N
|
||
if x.Len() != n {
|
||
panic(ErrShape)
|
||
}
|
||
dst.reuseAsNonZeroed(n)
|
||
t := blas.NoTrans
|
||
if trans {
|
||
t = blas.Trans
|
||
}
|
||
xMat, _ := untransposeExtract(x)
|
||
if xVec, ok := xMat.(*VecDense); ok && dst != xVec {
|
||
dst.checkOverlap(xVec.mat)
|
||
lapack64.Lagtm(t, 1, a.mat, xVec.asGeneral(), 0, dst.asGeneral())
|
||
} else {
|
||
xCopy := getVecDenseWorkspace(n, false)
|
||
xCopy.CloneFromVec(x)
|
||
lapack64.Lagtm(t, 1, a.mat, xCopy.asGeneral(), 0, dst.asGeneral())
|
||
putVecDenseWorkspace(xCopy)
|
||
}
|
||
}
|
||
|
||
// SolveTo solves a tridiagonal system A⋅X = B or Aᵀ⋅X = B where A is an
|
||
// n×n tridiagonal matrix represented by the receiver and B is a given n×nrhs
|
||
// matrix. If A is non-singular, the result will be stored into dst and nil will
|
||
// be returned. If A is singular, the contents of dst will be undefined and a
|
||
// Condition error will be returned.
|
||
func (a *Tridiag) SolveTo(dst *Dense, trans bool, b Matrix) error {
|
||
n, nrhs := b.Dims()
|
||
if n != a.mat.N {
|
||
panic(ErrShape)
|
||
}
|
||
|
||
dst.reuseAsNonZeroed(n, nrhs)
|
||
bU, bTrans := untranspose(b)
|
||
if dst == bU {
|
||
if bTrans {
|
||
work := getDenseWorkspace(n, nrhs, false)
|
||
defer putDenseWorkspace(work)
|
||
work.Copy(b)
|
||
dst.Copy(work)
|
||
}
|
||
} else {
|
||
if rm, ok := bU.(RawMatrixer); ok {
|
||
dst.checkOverlap(rm.RawMatrix())
|
||
}
|
||
dst.Copy(b)
|
||
}
|
||
|
||
var aCopy Tridiag
|
||
aCopy.CloneFromTridiag(a)
|
||
var ok bool
|
||
if trans {
|
||
ok = lapack64.Gtsv(blas.Trans, aCopy.mat, dst.mat)
|
||
} else {
|
||
ok = lapack64.Gtsv(blas.NoTrans, aCopy.mat, dst.mat)
|
||
}
|
||
if !ok {
|
||
return Condition(math.Inf(1))
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// SolveVecTo solves a tridiagonal system A⋅X = B or Aᵀ⋅X = B where A is an
|
||
// n×n tridiagonal matrix represented by the receiver and b is a given n-vector.
|
||
// If A is non-singular, the result will be stored into dst and nil will be
|
||
// returned. If A is singular, the contents of dst will be undefined and a
|
||
// Condition error will be returned.
|
||
func (a *Tridiag) SolveVecTo(dst *VecDense, trans bool, b Vector) error {
|
||
n, nrhs := b.Dims()
|
||
if n != a.mat.N || nrhs != 1 {
|
||
panic(ErrShape)
|
||
}
|
||
if b, ok := b.(RawVectorer); ok && dst != b {
|
||
dst.checkOverlap(b.RawVector())
|
||
}
|
||
dst.reuseAsNonZeroed(n)
|
||
if dst != b {
|
||
dst.CopyVec(b)
|
||
}
|
||
var aCopy Tridiag
|
||
aCopy.CloneFromTridiag(a)
|
||
var ok bool
|
||
if trans {
|
||
ok = lapack64.Gtsv(blas.Trans, aCopy.mat, dst.asGeneral())
|
||
} else {
|
||
ok = lapack64.Gtsv(blas.NoTrans, aCopy.mat, dst.asGeneral())
|
||
}
|
||
if !ok {
|
||
return Condition(math.Inf(1))
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// DoNonZero calls the function fn for each of the non-zero elements of A. The
|
||
// function fn takes a row/column index and the element value of A at (i,j).
|
||
func (a *Tridiag) DoNonZero(fn func(i, j int, v float64)) {
|
||
for i, aij := range a.mat.DU {
|
||
if aij != 0 {
|
||
fn(i, i+1, aij)
|
||
}
|
||
}
|
||
for i, aii := range a.mat.D {
|
||
if aii != 0 {
|
||
fn(i, i, aii)
|
||
}
|
||
}
|
||
for i, aij := range a.mat.DL {
|
||
if aij != 0 {
|
||
fn(i+1, i, aij)
|
||
}
|
||
}
|
||
}
|
||
|
||
// DoRowNonZero calls the function fn for each of the non-zero elements of row i
|
||
// of A. The function fn takes a row/column index and the element value of A at
|
||
// (i,j).
|
||
func (a *Tridiag) DoRowNonZero(i int, fn func(i, j int, v float64)) {
|
||
n := a.mat.N
|
||
if uint(i) >= uint(n) {
|
||
panic(ErrRowAccess)
|
||
}
|
||
if n == 1 {
|
||
v := a.mat.D[0]
|
||
if v != 0 {
|
||
fn(0, 0, v)
|
||
}
|
||
return
|
||
}
|
||
switch i {
|
||
case 0:
|
||
v := a.mat.D[0]
|
||
if v != 0 {
|
||
fn(i, 0, v)
|
||
}
|
||
v = a.mat.DU[0]
|
||
if v != 0 {
|
||
fn(i, 1, v)
|
||
}
|
||
case n - 1:
|
||
v := a.mat.DL[n-2]
|
||
if v != 0 {
|
||
fn(n-1, n-2, v)
|
||
}
|
||
v = a.mat.D[n-1]
|
||
if v != 0 {
|
||
fn(n-1, n-1, v)
|
||
}
|
||
default:
|
||
v := a.mat.DL[i-1]
|
||
if v != 0 {
|
||
fn(i, i-1, v)
|
||
}
|
||
v = a.mat.D[i]
|
||
if v != 0 {
|
||
fn(i, i, v)
|
||
}
|
||
v = a.mat.DU[i]
|
||
if v != 0 {
|
||
fn(i, i+1, v)
|
||
}
|
||
}
|
||
}
|
||
|
||
// DoColNonZero calls the function fn for each of the non-zero elements of
|
||
// column j of A. The function fn takes a row/column index and the element value
|
||
// of A at (i, j).
|
||
func (a *Tridiag) DoColNonZero(j int, fn func(i, j int, v float64)) {
|
||
n := a.mat.N
|
||
if uint(j) >= uint(n) {
|
||
panic(ErrColAccess)
|
||
}
|
||
if n == 1 {
|
||
v := a.mat.D[0]
|
||
if v != 0 {
|
||
fn(0, 0, v)
|
||
}
|
||
return
|
||
}
|
||
switch j {
|
||
case 0:
|
||
v := a.mat.D[0]
|
||
if v != 0 {
|
||
fn(0, 0, v)
|
||
}
|
||
v = a.mat.DL[0]
|
||
if v != 0 {
|
||
fn(1, 0, v)
|
||
}
|
||
case n - 1:
|
||
v := a.mat.DU[n-2]
|
||
if v != 0 {
|
||
fn(n-2, n-1, v)
|
||
}
|
||
v = a.mat.D[n-1]
|
||
if v != 0 {
|
||
fn(n-1, n-1, v)
|
||
}
|
||
default:
|
||
v := a.mat.DU[j-1]
|
||
if v != 0 {
|
||
fn(j-1, j, v)
|
||
}
|
||
v = a.mat.D[j]
|
||
if v != 0 {
|
||
fn(j, j, v)
|
||
}
|
||
v = a.mat.DL[j]
|
||
if v != 0 {
|
||
fn(j+1, j, v)
|
||
}
|
||
}
|
||
}
|