Implement diag (#36)

Author: atteson

src/WoodburyMatrices.jl

@@ -1,7 +1,7 @@
 module WoodburyMatrices
 
 using LinearAlgebra
-import LinearAlgebra: det, logdet, logabsdet, ldiv!, mul!, adjoint, transpose
+import LinearAlgebra: det, logdet, logabsdet, ldiv!, mul!, adjoint, transpose, diag
 import Base: +, *, \, inv, convert, copy, show, similar, axes, size
 using SparseArrays
 

src/woodbury.jl

@@ -76,6 +76,10 @@ function inv(W::AbstractWoodbury)
     return Woodbury(A′, U′, C′, V′)
 end
 
+# multiplying everything out would be O(m^2 k + m k^2) operations
+# best I could think of was O(m^2 k) operations
+diag(W::AbstractWoodbury) = diag(W.A) + vec(sum(W.U .* (W.C * W.V)', dims=2))
+
 +(W::Woodbury, X::Woodbury)       = Woodbury(W.A + X.A, [W.U X.U],
                                              cat(W.C, X.C; dims=(1,2)), [W.V; X.V])
 *(α::Real, W::Woodbury)           = Woodbury(α*W.A, W.U, α*W.C, W.V)

test/woodbury.jl

@@ -55,6 +55,7 @@ for elty in (Float32, Float64, ComplexF32, ComplexF64, Int)
     iFv = F\v
     @test norm(W\v - iFv)/norm(iFv) <= n*cond(F)*ε # Revisit. Condition number is wrong
     @test norm(inv(W)*v - iFv)/norm(iFv) <= n*cond(F)*ε
+    @test norm(diag(W) - diag(F)) <= n*cond(F)*ε # haven't thought about this error bound
     @test abs((det(W) - det(F))/det(F)) <= n*cond(F)*ε # Revisit. Condition number is wrong
     iWv = similar(iFv)
     if elty != Int