Merge pull request #124 from devmotion/indexing

Fix linear indexing of `PDiagMat` and `ScalMat`

Author: andreasnoack

src/pdiagmat.jl

@@ -30,8 +30,10 @@ LinearAlgebra.cholesky(a::PDiagMat) = cholesky(Diagonal(a.diag))
 
 ### Inheriting from AbstractMatrix
 
-Base.size(a::PDiagMat) = (a.dim, a.dim)
-Base.getindex(a::PDiagMat{T},i::Integer) where {T} = i % a.dim == (i ÷ a.dim + 1) ? a.diag[i % a.dim] : zero(T)
+function Base.getindex(a::PDiagMat, i::Integer)
+    ncol, nrow = fldmod1(i, a.dim)
+    ncol == nrow ? a.diag[nrow] : zero(eltype(a))
+end
 Base.getindex(a::PDiagMat{T},i::Integer,j::Integer) where {T} = i == j ? a.diag[i] : zero(T)
 
 ### Arithmetics

src/pdmat.jl

@@ -34,7 +34,6 @@ LinearAlgebra.cholesky(a::PDMat) = a.chol
 
 ### Inheriting from AbstractMatrix
 
-Base.size(a::PDMat) = size(a.mat)
 Base.getindex(a::PDMat, i::Int) = getindex(a.mat, i)
 Base.getindex(a::PDMat, I::Vararg{Int, N}) where {N} = getindex(a.mat, I...)
 

src/pdsparsemat.jl

@@ -32,7 +32,6 @@ LinearAlgebra.cholesky(a::PDSparseMat) = a.chol
 
 ### Inheriting from AbstractMatrix
 
-Base.size(a::PDSparseMat) = (a.dim, a.dim)
 Base.getindex(a::PDSparseMat,i::Integer) = getindex(a.mat, i)
 Base.getindex(a::PDSparseMat,I::Vararg{Int, N}) where {N} = getindex(a.mat, I...)
 

src/scalmat.jl

@@ -22,8 +22,10 @@ LinearAlgebra.cholesky(a::ScalMat) = cholesky(Diagonal(fill(a.value, a.dim)))
 
 ### Inheriting from AbstractMatrix
 
-Base.size(a::ScalMat) = (a.dim, a.dim)
-Base.getindex(a::ScalMat{T}, i::Integer) where {T} = i%a.dim == (i ÷ a.dim + 1) ? a.value : zero(T)
+function Base.getindex(a::ScalMat, i::Integer)
+    ncol, nrow = fldmod1(i, a.dim)
+    ncol == nrow ? a.value : zero(eltype(a))
+end
 Base.getindex(a::ScalMat{T}, i::Integer, j::Integer) where {T} = i == j ? a.value : zero(T)
 
 ### Arithmetics

test/addition.jl

@@ -7,7 +7,7 @@ for T in [Float64,Float32]
   printstyled("Testing addition with eltype = $T\n", color=:blue)
   M = convert(Array{T,2},[4. -2. -1.; -2. 5. -1.; -1. -1. 6.])
   V = convert(Array{T,1},[1.5, 2.5, 2.0])
-  X = convert(T,2.0)
+  local X = convert(T,2.0)
 
   pm1 = PDMat(M)
   pm2 = PDiagMat(V)

test/testutils.jl

@@ -78,7 +78,11 @@ function pdtest_basics(C::AbstractPDMat, Cmat::Matrix, d::Int, verbose::Int)
 
     _pdt(verbose, "eltype")
     @test eltype(C) == eltype(Cmat)
-#    @test eltype(typeof(C)) == eltype(typeof(Cmat))
+    #    @test eltype(typeof(C)) == eltype(typeof(Cmat))
+
+    _pdt(verbose, "index")
+    @test all(C[i] == Cmat[i] for i in 1:(d^2))
+    @test all(C[i, j] == Cmat[i, j] for j in 1:d, i in 1:d)
 end
 
 
@@ -180,9 +184,15 @@ function pdtest_mul(C::AbstractPDMat, Cmat::Matrix, X::Matrix, verbose::Int)
     _pdt(verbose, "multiply")
     @test C * X ≈ Cmat * X
 
+    y = similar(C * X, size(C, 1))
+    ymat = similar(Cmat * X, size(Cmat, 1))
     for i = 1:size(X,2)
         xi = vec(copy(X[:,i]))
         @test C * xi ≈ Cmat * xi
+
+        mul!(y, C, xi)
+        mul!(ymat, Cmat, xi)
+        @test y ≈ ymat
     end
 end