Add permutedims/factorize overrides (#64)

* Add permutedims/factorize overrides

* import permutedims

Author: dlfivefifty

src/LazyArrays.jl

@@ -17,7 +17,7 @@ import Base: AbstractArray, AbstractMatrix, AbstractVector,
          getindex, setindex!, intersect, @_inline_meta, inv,
          sort, sort!, issorted, sortperm, diff, cumsum, sum, in, broadcast,
          eltype, parent, real, imag,
-         conj, transpose, adjoint, vec,
+         conj, transpose, adjoint, permutedims, vec,
          exp, log, sqrt, cos, sin, tan, csc, sec, cot,
                    cosh, sinh, tanh, csch, sech, coth,
                    acos, asin, atan, acsc, asec, acot,
@@ -36,7 +36,7 @@ import Base.Broadcast: BroadcastStyle, AbstractArrayStyle, Broadcasted, broadcas
                         combine_eltypes, DefaultArrayStyle, instantiate, materialize,
                         materialize!, eltypes
 
-import LinearAlgebra: AbstractTriangular, AbstractQ, checksquare, pinv, fill!, tilebufsize, Abuf, Bbuf, Cbuf, dot
+import LinearAlgebra: AbstractTriangular, AbstractQ, checksquare, pinv, fill!, tilebufsize, Abuf, Bbuf, Cbuf, dot, factorize, qr, lu, cholesky
 
 import LinearAlgebra.BLAS: BlasFloat, BlasReal, BlasComplex
 

src/lazyapplying.jl

@@ -264,10 +264,10 @@ end
 @inline lazy_getindex(A, I...) = sub_materialize(view(A, I...))
 
 
-@inline getindex(A::ApplyMatrix, kr::Colon, jr::Colon) = lazy_getindex(A, kr, jr)
-@inline getindex(A::ApplyMatrix, kr::Colon, jr::AbstractUnitRange) = lazy_getindex(A, kr, jr)
-@inline getindex(A::ApplyMatrix, kr::AbstractUnitRange, jr::Colon) = lazy_getindex(A, kr, jr)
-@inline getindex(A::ApplyMatrix, kr::AbstractUnitRange, jr::AbstractUnitRange) = lazy_getindex(A, kr, jr)
+@inline getindex(A::LazyMatrix, kr::Colon, jr::Colon) = lazy_getindex(A, kr, jr)
+@inline getindex(A::LazyMatrix, kr::Colon, jr::AbstractUnitRange) = lazy_getindex(A, kr, jr)
+@inline getindex(A::LazyMatrix, kr::AbstractUnitRange, jr::Colon) = lazy_getindex(A, kr, jr)
+@inline getindex(A::LazyMatrix, kr::AbstractUnitRange, jr::AbstractUnitRange) = lazy_getindex(A, kr, jr)
 
 
 diagonallayout(::LazyLayout) = DiagonalLayout{LazyLayout}()

src/lazybroadcasting.jl

@@ -164,3 +164,10 @@ _broadcast_rowsupport(sz, A::AbstractMatrix, j) = size(A,2) == 1 ? OneTo(sz[2])
 
 colsupport(::BroadcastLayout{typeof(*)}, A, j) = intersect(_broadcast_colsupport.(Ref(size(A)), A.args, j)...)
 rowsupport(::BroadcastLayout{typeof(*)}, A, j) = intersect(_broadcast_rowsupport.(Ref(size(A)), A.args, j)...)
+
+for op in (:+, :-)
+    @eval begin
+        colsupport(::BroadcastLayout{typeof($op)}, A, j) = convexunion(_broadcast_colsupport.(Ref(size(A)), A.args, j)...)
+        rowsupport(::BroadcastLayout{typeof($op)}, A, j) = convexunion(_broadcast_rowsupport.(Ref(size(A)), A.args, j)...)
+    end
+end

src/lazyconcat.jl

@@ -118,6 +118,27 @@ end
 # based on Base/array.jl, Base/abstractarray.jl
 
 function copyto!(dest::AbstractMatrix, V::Vcat{<:Any,2})
+    arrays = V.args
+    nargs = length(arrays)
+    nrows = size(dest,1)
+    nrows == sum(a->size(a, 1), arrays) || throw(DimensionMismatch("sum of rows each matrix must equal $nrows"))
+    ncols = size(dest, 2)
+    for a in arrays
+        if size(a, 2) != ncols
+            throw(DimensionMismatch("number of columns of each array must match (got $(map(x->size(x,2), A)))"))
+        end
+    end
+    pos = 1
+    for a in arrays
+        p1 = pos+size(a,1)-1
+        dest[pos:p1, :] .= a
+        pos = p1+1
+    end
+    return dest
+end
+
+# this is repeated to avoid allocation in .=
+function copyto!(dest::AbstractMatrix, V::Vcat{<:Any,2,<:Tuple{Vararg{<:AbstractMatrix}}})
     arrays = V.args
     nargs = length(arrays)
     nrows = size(dest,1)
@@ -259,6 +280,12 @@ _vec(a::AbstractArray) = vec(a)
 _vec(a::Adjoint{<:Number,<:AbstractVector}) = _vec(parent(a))
 vec(A::Hcat) = Vcat(_vec.(A.args)...)
 
+_permutedims(a) = a
+_permutedims(a::AbstractArray) = permutedims(a)
+
+permutedims(A::Hcat{T}) where T = Vcat{T}(map(_permutedims,A.args)...)
+permutedims(A::Vcat{T}) where T = Hcat{T}(map(_permutedims,A.args)...)
+
 
 #####
 # broadcasting

src/linalg/factorizations.jl

@@ -33,4 +33,17 @@ function copyto!(dest::AbstractArray, M::Ldiv{QLayout})
     materialize!(Ldiv(A,dest))
 end
 
-materialize!(M::Ldiv{QLayout}) = materialize!(Lmul(M.A',M.B))
+materialize!(M::Ldiv{QLayout}) = materialize!(Lmul(M.A',M.B))
+
+factorizestyle(_) = DefaultArrayApplyStyle()
+
+for op in (:factorize, :qr, :lu, :cholesky)
+    @eval begin
+        $op(B::LazyMatrix) = apply($op, B)
+        ApplyStyle(::typeof($op), B::Type{<:AbstractMatrix}) = factorizestyle(MemoryLayout(B))
+        materialize(A::Applied{DefaultArrayApplyStyle,typeof($op),<:Tuple{<:AbstractMatrix{T}}}) where T = 
+            Base.invoke($op, Tuple{AbstractMatrix{T}}, A.args...)
+
+        eltype(A::Applied{<:Any,typeof($op)}) = float(eltype(first(A.args)))
+    end
+end

test/concattests.jl

@@ -19,6 +19,7 @@ import LazyArrays: MemoryLayout, DenseColumnMajor, PaddedLayout, materialize!, M
         @test copy(A) !== A
         @test vec(A) === A
         @test A' == transpose(A) == Vector(A)'
+        @test permutedims(A) == permutedims(Vector(A))
 
         A = @inferred(Vcat(1:10, 1:20))
         @test @inferred(length(A)) == 30
@@ -35,6 +36,7 @@ import LazyArrays: MemoryLayout, DenseColumnMajor, PaddedLayout, materialize!, M
         @test A' == transpose(A) == Vector(A)'
         @test A' === Hcat((1:10)', (1:20)')
         @test transpose(A) === Hcat(transpose(1:10), transpose(1:20))
+        @test permutedims(A) == permutedims(Vector(A))
 
         A = Vcat(randn(2,10), randn(4,10))
         @test @inferred(length(A)) == 60
@@ -51,6 +53,7 @@ import LazyArrays: MemoryLayout, DenseColumnMajor, PaddedLayout, materialize!, M
         @test copy(A) !== A
         @test vec(A) == vec(Matrix(A))
         @test A' == transpose(A) == Matrix(A)'
+        @test permutedims(A) == permutedims(Matrix(A))
 
         A = Vcat(randn(2,10).+im.*randn(2,10), randn(4,10).+im.*randn(4,10))
         @test eltype(A) == ComplexF64
@@ -69,6 +72,7 @@ import LazyArrays: MemoryLayout, DenseColumnMajor, PaddedLayout, materialize!, M
         @test vec(A) == vec(Matrix(A))
         @test A' == Matrix(A)'
         @test transpose(A) == transpose(Matrix(A))
+        @test permutedims(A) == permutedims(Matrix(A))
 
         @test Vcat() isa Vcat{Any,1,Tuple{}}
 
@@ -77,6 +81,7 @@ import LazyArrays: MemoryLayout, DenseColumnMajor, PaddedLayout, materialize!, M
         @test A[1,1] == 1.0
         @test A[2,1] == 0.0
         @test axes(A) == (Base.OneTo(4),Base.OneTo(1))
+        @test permutedims(A) == permutedims(Matrix(A))
     end
     @testset "Hcat" begin
         A = @inferred(Hcat(1:10, 2:11))