Better support for columns of inverses (#225)

* Better support for columns of inverses

* Add tests

* add tests, complete BroadcastArray(/, ...) and BroadcastArray(\, ...) support

Project.toml

@@ -1,6 +1,6 @@
 name = "LazyArrays"
 uuid = "5078a376-72f3-5289-bfd5-ec5146d43c02"
-version = "0.22.12"
+version = "0.22.13"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"

src/lazybroadcasting.jl

@@ -195,6 +195,26 @@ _broadcast_rowsupport(ax, ::Tuple{<:Any,<:Any}, A, j) = rowsupport(A,j)
 colsupport(lay::BroadcastLayout{typeof(*)}, A, j) = intersect(_broadcast_colsupport.(Ref(axes(A)), axes.(arguments(lay,A)), arguments(lay,A), Ref(j))...)
 rowsupport(lay::BroadcastLayout{typeof(*)}, A, j) = intersect(_broadcast_rowsupport.(Ref(axes(A)), axes.(arguments(lay,A)), arguments(lay,A), Ref(j))...)
 
+function colsupport(lay::BroadcastLayout{typeof(\)}, A, j)
+    _,b = arguments(lay,A)
+    _broadcast_colsupport(axes(A), axes(b), b, j)
+end
+
+function rowsupport(lay::BroadcastLayout{typeof(\)}, A, j)
+    _,b = arguments(lay,A)
+    _broadcast_rowsupport(axes(A), axes(b), b, j)
+end
+
+function colsupport(lay::BroadcastLayout{typeof(/)}, A, j)
+    b,_ = arguments(lay,A)
+    _broadcast_colsupport(axes(A), axes(b), b, j)
+end
+
+function rowsupport(lay::BroadcastLayout{typeof(/)}, A, j)
+    b,_ = arguments(lay,A)
+    _broadcast_rowsupport(axes(A), axes(b), b, j)
+end
+
 for op in (:+, :-)
     @eval begin
         rowsupport(lay::BroadcastLayout{typeof($op)}, A, j) = convexunion(_broadcast_rowsupport.(Ref(axes(A)), axes.(arguments(lay,A)), arguments(lay,A), Ref(j))...)

src/linalg/inv.jl

@@ -155,9 +155,17 @@ inv(D::Diagonal{T,<:LazyVector}) where T = Diagonal(inv.(D.diag))
 ###
 
 # \ is likely to be specialised
-@propagate_inbounds getindex(Ai::InvMatrix{T}, ::Colon, j::Integer) where T = parent(Ai) \ [Zeros{T}(j-1); one(T); Zeros{T}(size(parent(Ai),1)-j)]
-@propagate_inbounds getindex(Ai::PInvMatrix{T}, ::Colon, j::Integer) where T = parent(Ai) \ [Zeros{T}(j-1); one(T); Zeros{T}(size(parent(Ai),1)-j)]
-getindex(Ai::SubArray{<:Any,2,<:InvMatrix}, ::Colon, j::Integer) = parent(Ai)[:, parentindices(Ai)[2][j]]
+struct InvColumnLayout <: AbstractLazyLayout end
+
+
+sublayout(::AbstractInvLayout, I::Type{<:Tuple{Slice, Int}}) = InvColumnLayout()
+
+function sub_materialize(::InvColumnLayout, v::AbstractVector, _)
+    _,j = parentindices(v)
+    Ai = parent(v)
+    T = eltype(v)
+    parent(Ai) \ [Zeros{T}(j-1); one(T); Zeros{T}(size(parent(Ai),1)-j)]
+end
 
 @propagate_inbounds getindex(A::PInvMatrix{T}, k::Integer, j::Integer) where T = A[:,j][k]
 @propagate_inbounds getindex(A::InvMatrix{T}, k::Integer, j::Integer) where T = A[:,j][k]

test/broadcasttests.jl

@@ -327,4 +327,16 @@ import Base: broadcasted
         @test A + I isa BroadcastArray
         @test I + A isa BroadcastArray
     end
+
+    @testset "/ and \\" begin
+        for A in (BroadcastArray(\, ones(5,5), UpperTriangular(randn(5,5))),
+                    BroadcastArray(\, ones(5), UpperTriangular(randn(5,5))),
+                    BroadcastArray(\, 6, UpperTriangular(randn(5,5))),
+                    BroadcastArray(/, UpperTriangular(randn(5,5)), ones(5,5)),
+                    BroadcastArray(/, UpperTriangular(randn(5,5)), ones(5)),
+                    BroadcastArray(/, UpperTriangular(randn(5,5)), 5))
+            @test colsupport(A, 3) == 1:3
+            @test rowsupport(A, 3) == 3:5
+        end
+    end
 end

test/ldivtests.jl

@@ -75,6 +75,9 @@ import Base.Broadcast: materialize
         @test Ai ≈ pinv(A)
         @test_throws DimensionMismatch inv(PInv(A))
         @test_throws DimensionMismatch InvMatrix(A)
+
+        A = UpperTriangular(Ones(20,20))
+        @test ApplyArray(inv,A)[1:10,1:10] ≈ diagm(0 => ones(10), 1 => -ones(9))
     end
 
     @testset "Int" begin