mat: add non-zero doers for TriBandDense

gonum · Jun 18, 2022 · 520a77a · 520a77a
1 parent a2c6f81
commit 520a77a
Show file tree

Hide file tree

Showing 2 changed files with 80 additions and 0 deletions.
diff --git 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
@@ -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() {