Respond to PR comments

gonum · Nov 25, 2018 · 67021b8 · 67021b8
1 parent 26da546
commit 67021b8
Show file tree

Hide file tree

Showing 4 changed files with 56 additions and 10 deletions.
diff --git a/mat/diagonal.go b/mat/diagonal.go
@@ -209,7 +209,16 @@ func (d *DiagDense) DiagFrom(m Matrix) {
 			Inc:  mat.Stride + 1,
 			Data: mat.Data[:(n-1)*mat.Stride+n],
 		}
-	// TODO(kortschak): Add banded triangular handling when the type exists.
+	case RawTriBander:
+		mat := r.RawTriBand()
+		data := mat.Data
+		if mat.Uplo == blas.Lower {
+			data = data[mat.K:]
+		}
+		vec = blas64.Vector{
+			Inc:  mat.Stride,
+			Data: data[:(n-1)*mat.Stride+1],
+		}
 	case RawTriangular:
 		mat := r.RawTriangular()
 		if mat.Diag == blas.Unit {

diff --git a/mat/diagonal_test.go b/mat/diagonal_test.go
@@ -244,6 +244,42 @@ func TestDiagFrom(t *testing.T) {
 				0, 0, 0, 0, 0, 6,
 			}),
 		},
+		{
+			mat: NewTriBandDense(6, 2, Upper, []float64{
+				1, math.NaN(), math.NaN(),
+				2, math.NaN(), math.NaN(),
+				3, math.NaN(), math.NaN(),
+				4, math.NaN(), math.NaN(),
+				5, math.NaN(), math.NaN(),
+				6, math.NaN(), math.NaN(),
+			}),
+			want: NewDense(6, 6, []float64{
+				1, 0, 0, 0, 0, 0,
+				0, 2, 0, 0, 0, 0,
+				0, 0, 3, 0, 0, 0,
+				0, 0, 0, 4, 0, 0,
+				0, 0, 0, 0, 5, 0,
+				0, 0, 0, 0, 0, 6,
+			}),
+		},
+		{
+			mat: NewTriBandDense(6, 2, Lower, []float64{
+				math.NaN(), math.NaN(), 1,
+				math.NaN(), math.NaN(), 2,
+				math.NaN(), math.NaN(), 3,
+				math.NaN(), math.NaN(), 4,
+				math.NaN(), math.NaN(), 5,
+				math.NaN(), math.NaN(), 6,
+			}),
+			want: NewDense(6, 6, []float64{
+				1, 0, 0, 0, 0, 0,
+				0, 2, 0, 0, 0, 0,
+				0, 0, 3, 0, 0, 0,
+				0, 0, 0, 4, 0, 0,
+				0, 0, 0, 0, 5, 0,
+				0, 0, 0, 0, 0, 6,
+			}),
+		},
 		{
 			mat: NewTriDense(6, Upper, []float64{
 				1, math.NaN(), math.NaN(), math.NaN(), math.NaN(), math.NaN(),

diff --git a/mat/triangular.go b/mat/triangular.go
@@ -298,6 +298,9 @@ func (t *TriDense) isolatedWorkspace(a Triangular) (w *TriDense, restore func())
 
 // DiagView returns the diagonal as a matrix backed by the original data.
 func (t *TriDense) DiagView() Diagonal {
+	if t.mat.Diag == blas.Unit {
+		panic("mat: cannot take view of Unit diagonal")
+	}
 	n := t.mat.N
 	return &DiagDense{
 		mat: blas64.Vector{

diff --git a/mat/triband.go b/mat/triband.go
@@ -308,20 +308,18 @@ func (t *TriBandDense) RawTriBand() blas64.TriangularBand {
 
 // DiagView returns the diagonal as a matrix backed by the original data.
 func (t *TriBandDense) DiagView() Diagonal {
+	if t.mat.Diag == blas.Unit {
+		panic("mat: cannot take view of Unit diagonal")
+	}
 	n := t.mat.N
-	if t.isUpper() {
-		return &DiagDense{
-			mat: blas64.Vector{
-				Inc:  t.mat.Stride,
-				Data: t.mat.Data[:(n-1)*t.mat.Stride+1],
-			},
-			n: n,
-		}
+	data := t.mat.Data
+	if !t.isUpper() {
+		data = data[t.mat.K:]
 	}
 	return &DiagDense{
 		mat: blas64.Vector{
 			Inc:  t.mat.Stride,
-			Data: t.mat.Data[t.mat.K : t.mat.K+(n-1)*t.mat.Stride+1],
+			Data: data[:(n-1)*t.mat.Stride+1],
 		},
 		n: n,
 	}