diff --git a/mat/matrix_test.go b/mat/matrix_test.go index 5c9666230..8cc2d8c7e 100644 --- a/mat/matrix_test.go +++ b/mat/matrix_test.go @@ -612,6 +612,14 @@ func TestDoer(t *testing.T) { NewSymBandDense(3, 1, ones(3*(1+1))), NewSymBandDense(6, 1, ones(6*(1+1))), NewSymBandDense(6, 2, ones(6*(2+1))), + NewTriBandDense(3, 0, Upper, ones(3)), + NewTriBandDense(3, 1, Upper, ones(3*(1+1))), + NewTriBandDense(6, 1, Upper, ones(6*(1+1))), + NewTriBandDense(6, 2, Upper, ones(6*(2+1))), + NewTriBandDense(3, 0, Lower, ones(3)), + NewTriBandDense(3, 1, Lower, ones(3*(1+1))), + NewTriBandDense(6, 1, Lower, ones(6*(1+1))), + NewTriBandDense(6, 2, Lower, ones(6*(2+1))), NewTridiag(1, nil, ones(1), nil), NewTridiag(2, ones(1), ones(2), ones(1)), NewTridiag(3, ones(2), ones(3), ones(2)), @@ -639,6 +647,8 @@ func TestDoer(t *testing.T) { m.SetBand(i, j, v) case MutableSymBanded: m.SetSymBand(i, j, v) + case MutableTriBanded: + m.SetTriBand(i, j, v) default: panic("bad test: need mutable type") } diff --git a/mat/triband.go b/mat/triband.go index bffe93fbb..73495ac44 100644 --- a/mat/triband.go +++ b/mat/triband.go @@ -353,6 +353,76 @@ func (t *TriBandDense) reuseAsNonZeroed(n, k int, kind TriKind) { } } +// DoNonZero calls the function fn for each of the non-zero elements of t. The function fn +// takes a row/column index and the element value of t at (i, j). +func (t *TriBandDense) DoNonZero(fn func(i, j int, v float64)) { + if t.isUpper() { + for i := 0; i < t.mat.N; i++ { + for j := i; j < min(i+t.mat.K+1, t.mat.N); j++ { + v := t.at(i, j) + if v != 0 { + fn(i, j, v) + } + } + } + } else { + for i := 0; i < t.mat.N; i++ { + for j := max(0, i-t.mat.K); j <= i; j++ { + v := t.at(i, j) + if v != 0 { + fn(i, j, v) + } + } + } + } +} + +// DoRowNonZero calls the function fn for each of the non-zero elements of row i of t. The function fn +// takes a row/column index and the element value of t at (i, j). +func (t *TriBandDense) DoRowNonZero(i int, fn func(i, j int, v float64)) { + if i < 0 || t.mat.N <= i { + panic(ErrRowAccess) + } + if t.isUpper() { + for j := i; j < min(i+t.mat.K+1, t.mat.N); j++ { + v := t.at(i, j) + if v != 0 { + fn(i, j, v) + } + } + } else { + for j := max(0, i-t.mat.K); j <= i; j++ { + v := t.at(i, j) + if v != 0 { + fn(i, j, v) + } + } + } +} + +// DoColNonZero calls the function fn for each of the non-zero elements of column j of t. The function fn +// takes a row/column index and the element value of t at (i, j). +func (t *TriBandDense) DoColNonZero(j int, fn func(i, j int, v float64)) { + if j < 0 || t.mat.N <= j { + panic(ErrColAccess) + } + if t.isUpper() { + for i := 0; i < t.mat.N; i++ { + v := t.at(i, j) + if v != 0 { + fn(i, j, v) + } + } + } else { + for i := 0; i < t.mat.N; i++ { + v := t.at(i, j) + if v != 0 { + fn(i, j, v) + } + } + } +} + // Zero sets all of the matrix elements to zero. func (t *TriBandDense) Zero() { if t.isUpper() {