Add mul/colsupport for MulQuasiArray, support overloading dot and sum (#94)

dlfivefifty · web-flow · commit 631e150491b2 · 2023-07-04T22:03:48.000+01:00
* Add mul/colsupport for MulQuasiArray

* multiplication with QuasiArrays and Zeros

* Improve show

* increase coverage

* Fix test

* fix tests

* sum -&gt; sum_layout
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -12,7 +12,6 @@ jobs:
         version:
           - '1.6'
           - '1'
-          - '^1.9.0-0'
         os:
           - ubuntu-latest
           - macOS-latest
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "QuasiArrays"
 uuid = "c4ea9172-b204-11e9-377d-29865faadc5c"
 authors = ["Sheehan Olver <solver@mac.com>"]
-version = "0.9.8"
+version = "0.10"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
@@ -12,10 +12,10 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [compat]
-ArrayLayouts = "0.7.6, 0.8, 1"
-DomainSets = "0.5, 0.6"
-FillArrays = "0.12, 0.13, 1"
-LazyArrays = "0.22.2, 1"
+ArrayLayouts = "1"
+DomainSets = "0.6"
+FillArrays = "1"
+LazyArrays = "1.2"
 StaticArrays = "1"
 julia = "1.6"
 
diff --git a/src/QuasiArrays.jl b/src/QuasiArrays.jl
@@ -49,7 +49,7 @@ import LazyArrays: MemoryLayout, UnknownLayout, Mul, ApplyLayout, BroadcastLayou
                     LdivStyle, InvLayout, PInvLayout, sub_materialize, lazymaterialize,
                     _mul, rowsupport, DiagonalLayout, adjointlayout, transposelayout, conjlayout,
                     sublayout, call, LazyArrayStyle, layout_getindex, _broadcast2broadcastarray, _applyarray_summary, _broadcastarray_summary,
-                    _broadcasted_mul, simplifiable, simplify
+                    _broadcasted_mul, simplifiable, simplify, _mul_colsupport, _mul_rowsupport
 
 import Base.IteratorsMD
 
diff --git a/src/abstractquasiarray.jl b/src/abstractquasiarray.jl
@@ -523,8 +523,5 @@ end
 ##
 # show
 ##
-
-show(io::IO, A::AbstractQuasiArray) = summary(io, A)
-
 struct QuasiArrayLayout <: MemoryLayout end
 MemoryLayout(::Type{<:AbstractQuasiArray}) = QuasiArrayLayout()
diff --git a/src/indices.jl b/src/indices.jl
@@ -194,6 +194,7 @@ getindex(S::Inclusion{T}, i::AbstractArray{T}) where T = (@_inline_meta; @bounds
 getindex(S::Inclusion, i::Inclusion) = (@_inline_meta; @boundscheck checkbounds(S, i); copy(S))
 getindex(S::Inclusion, ::Colon) = copy(S)
 Base.unsafe_getindex(S::Inclusion{T}, x) where T = convert(T, x)::T
+show(io::IO, r::Inclusion) = summary(io, r)
 summary(io::IO, r::Inclusion) = print(io, "Inclusion(", r.domain, ")")
 iterate(S::Inclusion, s...) = iterate(S.domain, s...)
 
diff --git a/src/lazyquasiarrays.jl b/src/lazyquasiarrays.jl
@@ -173,13 +173,17 @@ call(b::BroadcastLayout, a::SubQuasiArray) = call(b, parent(a))
 # show
 ####
 
+show(io::IO, A::BroadcastQuasiArray) = summary(io, A)
+show(io::IO, A::ApplyQuasiArray) = summary(io, A)
 
 summary(io::IO, A::BroadcastQuasiArray) = _broadcastarray_summary(io, A)
 summary(io::IO, A::ApplyQuasiArray) = _applyarray_summary(io, A)
 
 _mul_summary(_, io, A) = _applyarray_summary(io, A)
 summary(io::IO, A::ApplyQuasiArray{<:Any,N,typeof(*)}) where N = _mul_summary(MemoryLayout(A), io, A)
 
+
+
 for op in (:+, :-, :*, :\, :/)
     @eval begin
         function summary(io::IO, A::BroadcastQuasiArray{<:Any,N,typeof($op)}) where N
diff --git a/src/matmul.jl b/src/matmul.jl
@@ -64,7 +64,10 @@ _factors(M) = (M,)
 @inline flatten(A::MulQuasiArray) = ApplyQuasiArray(flatten(Applied(A)))
 @inline flatten(A::SubQuasiArray{<:Any,2,<:MulQuasiArray}) = materialize(flatten(Applied(A)))
 
-
+colsupport(B::MulQuasiArray, j) = _mul_colsupport(j, reverse(B.args)...)
+rowsupport(B::MulQuasiArray, j) = _mul_rowsupport(j, B.args...)
+_mul_colsupport(j, Z::AbstractQuasiArray) = colsupport(Z,j)
+_mul_colsupport(j, Z::AbstractQuasiArray, Y...) = _mul_colsupport(colsupport(Z,j), Y...)
 
 adjoint(A::MulQuasiArray) = ApplyQuasiArray(*, reverse(adjoint.(A.args))...)
 transpose(A::MulQuasiArray) = ApplyQuasiArray(*, reverse(transpose.(A.args))...)
diff --git a/src/quasiadjtrans.jl b/src/quasiadjtrans.jl
@@ -167,15 +167,10 @@ map(f, tvs::QuasiTransposeAbsVec...) = transpose(map((xs...) -> transpose(f(tran
 ## multiplication *
 
 # QuasiAdjoint/QuasiTranspose-vector * vector
+dot(a::AbstractQuasiArray, b::AbstractQuasiArray) = ArrayLayouts.dot(a, b)
+@inline copy(d::Dot{<:Any,<:Any,<:AbstractQuasiArray,<:AbstractQuasiArray}) = _dot(size(d.A,1), d.A, d.B)
+_dot(sz, a, b) = Base.invoke(dot, NTuple{2,Any}, a, b)
 
-# allow re-expansion using Legendre
-_dot(ax, a, b) = Base.invoke(dot, NTuple{2,Any}, a, b)
-
-function dot(a::AbstractQuasiArray, b::AbstractQuasiArray)
-    ax = axes(a,1)
-    ax == axes(b,1) || throw(DimensionMismatch())
-    _dot(ax, a, b)
-end
 
 *(u::QuasiAdjointAbsVec, v::AbstractQuasiVector) = dot(u.parent, v)
 *(u::QuasiTransposeAbsVec{T}, v::AbstractQuasiVector{T}) where {T<:Real} = dot(u.parent, v)
diff --git a/src/quasiarray.jl b/src/quasiarray.jl
@@ -149,3 +149,19 @@ function reshape(A::QuasiVector, ax::Tuple{Any,OneTo{Int}})
     @assert ax == (axes(A,1),Base.OneTo(1))
     QuasiMatrix(reshape(A.parent,size(A.parent,1),1), (A.axes[1], Base.OneTo(1)))
 end
+
+function show(io::IO, A::QuasiVector)
+    print(io, "QuasiVector(")
+    show(io, A.parent)
+    print(io, ", ")
+    show(io, A.axes[1])
+    print(io, ")")
+end
+
+function show(io::IO, A::QuasiMatrix)
+    print(io, "QuasiMatrix(")
+    show(io, A.parent)
+    print(io, ", ")
+    show(io, A.axes)
+    print(io, ")")
+end
diff --git a/src/quasifill.jl b/src/quasifill.jl
@@ -194,6 +194,8 @@ end
 # show
 #######
 
+show(io::IO, F::AbstractQuasiFill) = summary(io, F)
+
 function summary(io::IO, F::QuasiOnes)
     print(io, "ones(")
     summary(io, F.axes[1])
@@ -488,7 +490,7 @@ MemoryLayout(::Type{<:QuasiZeros}) = ZerosLayout()
 MemoryLayout(::Type{<:QuasiOnes}) = OnesLayout()
 
 _quasi_mul(M::Mul{ZerosLayout}, _) = QuasiZeros{eltype(M)}(axes(M))
-_quasi_mul(M::Mul{QuasiArrayLayout,ZerosLayout}, _) = QuasiZeros{eltype(M)}(axes(M))
+_quasi_mul(M::Mul{QuasiArrayLayout,ZerosLayout}, _) = FillArrays.mult_zeros(M.A, M.B)
 _quasi_mul(M::Mul{QuasiArrayLayout,ZerosLayout}, ::NTuple{N,OneTo{Int}}) where N = Zeros{eltype(M)}(axes(M))
 fillzeros(::Type{T}, a::Tuple{AbstractQuasiVector,Vararg{Any}}) where T<:Number = QuasiZeros{T}(a)
 fillzeros(::Type{T}, a::Tuple{Any,AbstractQuasiVector,Vararg{Any}}) where T<:Number = QuasiZeros{T}(a)
@@ -520,4 +522,12 @@ fill(c, x::Inclusion, y::Union{OneTo,IdentityUnitRange,Inclusion}...) = QuasiFil
 fill(c, x::Union{OneTo,IdentityUnitRange}, y::Inclusion, z::Union{OneTo,IdentityUnitRange,Inclusion}...) = QuasiFill(c, (x, y, z...))
 
 iszero(x::AbstractQuasiFill) = iszero(getindex_value(x))
-isone(x::AbstractQuasiFill) = isone(getindex_value(x))
+isone(x::AbstractQuasiFill) = isone(getindex_value(x))
+
+
+function FillArrays.mult_zeros(a::AbstractQuasiArray, b)
+    axes(a, 2) ≠ axes(b, 1) &&
+        throw(DimensionMismatch("A has dimensions $(size(a)) but B has dimensions $(size(b))"))
+    T = promote_type(eltype(a), eltype(b))
+    fillsimilar(QuasiZeros{T}(), axes(a, 1), axes(b)[2:end]...)
+end
diff --git a/src/quasireducedim.jl b/src/quasireducedim.jl
@@ -263,21 +263,21 @@ end
 
 
 # support overloading sum by MemoryLayout
-_sum(V::AbstractQuasiArray, dims) = __sum(MemoryLayout(V), V, dims)
-_sum(V::AbstractQuasiArray, ::Colon) = __sum(MemoryLayout(V), V, :)
+_sum(V::AbstractQuasiArray, dims) = sum_layout(MemoryLayout(V), V, dims)
+_sum(V::AbstractQuasiArray, ::Colon) = sum_layout(MemoryLayout(V), V, :)
 
-_cumsum(A, dims) = __cumsum(MemoryLayout(A), A, dims)
+_cumsum(A, dims) = cumsum_layout(MemoryLayout(A), A, dims)
 cumsum(A::AbstractQuasiArray; dims::Integer) = _cumsum(A, dims)
 cumsum(x::AbstractQuasiVector) = cumsum(x, dims=1)
 
 # sum is equivalent to hitting by ones(n) on the left or right
 
-__cumsum(::QuasiArrayLayout, A, ::Colon) = QuasiArray(cumsum(parent(A)), axes(A))
-__cumsum(::QuasiArrayLayout, A, d::Int) = QuasiArray(cumsum(parent(A),dims=d), axes(A))
+cumsum_layout(::QuasiArrayLayout, A, ::Colon) = QuasiArray(cumsum(parent(A)), axes(A))
+cumsum_layout(::QuasiArrayLayout, A, d::Int) = QuasiArray(cumsum(parent(A),dims=d), axes(A))
 
 for Sum in (:sum, :cumsum)
-    __Sum = Symbol("__", Sum)
-    @eval function $__Sum(LAY::ApplyLayout{typeof(*)}, V::AbstractQuasiMatrix, d::Int)
+    Sum_Lay = Symbol(Sum, "_layout")
+    @eval function $Sum_Lay(LAY::ApplyLayout{typeof(*)}, V::AbstractQuasiMatrix, d::Int)
         a = arguments(LAY, V)
         if d == 1
             *($Sum(first(a); dims=1), tail(a)...)
@@ -288,15 +288,16 @@ for Sum in (:sum, :cumsum)
     end
 end
 
-function __cumsum(LAY::ApplyLayout{typeof(*)}, V::AbstractQuasiVector, dims)
+function cumsum_layout(LAY::ApplyLayout{typeof(*)}, V::AbstractQuasiVector, dims)
     a = arguments(LAY, V)
     apply(*, cumsum(a[1]; dims=dims), tail(a)...)
 end
 
-function __sum(LAY::ApplyLayout{typeof(*)}, V::AbstractQuasiVector, ::Colon)
+function sum_layout(LAY::ApplyLayout{typeof(*)}, V::AbstractQuasiVector, ::Colon)
     a = arguments(LAY, V)
     first(apply(*, sum(a[1]; dims=1), tail(a)...))
 end
 
-__sum(_, A, dims) = _sum(identity, A, dims)
+sum_layout(::MemoryLayout, A, dims) = sum_size(size(A), A, dims)
+sum_size(::NTuple{N,Int}, A, dims) where N = _sum(identity, A, dims)
 
diff --git a/src/subquasiarray.jl b/src/subquasiarray.jl
@@ -333,6 +333,8 @@ end
 # show
 ###
 
+show(io::IO, v::SubQuasiArray) = summary(io, v)
+
 function summary(io::IO, v::SubQuasiArray)
     print(io, "view(")
     summary(io, parent(v))
diff --git a/test/test_dense.jl b/test/test_dense.jl
@@ -1,4 +1,4 @@
-using QuasiArrays, ArrayLayouts, LazyArrays, Test
+using QuasiArrays, ArrayLayouts, LazyArrays, Base64, Test
 import QuasiArrays: QuasiArrayLayout, QuasiArrayApplyStyle
 import LazyArrays: ApplyStyle, MulStyle
 
@@ -24,4 +24,11 @@ import LazyArrays: ApplyStyle, MulStyle
         @test V^1 == A
         @test V^2 == A^2
     end
+
+    @testset "show" begin
+        A = QuasiArray(reshape(1:9,3,3), (0:0.5:1, 0:0.5:1))
+        @test stringmime("text/plain", A) == "QuasiMatrix([1 4 7; 2 5 8; 3 6 9], (0.0:0.5:1.0, 0.0:0.5:1.0))"
+        x = QuasiArray(1:3, 0:0.5:1)
+        @test stringmime("text/plain", x) == "QuasiVector([1, 2, 3], 0.0:0.5:1.0)"
+    end
 end
diff --git a/test/test_quasifill.jl b/test/test_quasifill.jl
@@ -299,7 +299,7 @@ import QuasiArrays: AbstractQuasiFill
         @test_throws DimensionMismatch B * QuasiZeros(ax, ax)
 
         # Check multiplication by Adjoint vectors works as expected.
-        @test QuasiArray(randn(3, 4),ax,bx)' * QuasiZeros(ax) === QuasiZeros(bx)
+        @test QuasiArray(randn(3, 4),ax,bx)' * QuasiZeros(ax) === Zeros((axes(B,2),))
         @test QuasiArray(randn(4),bx)' * QuasiZeros(bx) == zero(Float64)
         @test QuasiArray([1, 2, 3],ax)' * QuasiZeros{Int}(ax) == zero(Int)
         @test_broken QuasiArray([SVector(1,2)', SVector(2,3)', SVector(3,4)'],ax)' * QuasiZeros{Int}(ax) == SVector(0,0)
@@ -662,7 +662,7 @@ import QuasiArrays: AbstractQuasiFill
         @test A*QuasiOnes(bx) ≈ A*QuasiVector(QuasiOnes(bx))
         @test A*QuasiOnes(bx,cx) ≈ A*QuasiMatrix(QuasiOnes(bx,cx))
         @test A*QuasiZeros(bx)  ≡ Zeros((Base.OneTo(1),))
-        @test A*QuasiZeros(bx,cx) ≡ QuasiZeros(Base.OneTo(1),cx)
+        @test A*QuasiZeros(bx,cx) ≡ Zeros((Base.OneTo(1),Base.IdentityUnitRange(cx)))
 
         D = QuasiDiagonal(a)
         @test QuasiZeros(ax,bx)*D ≡ QuasiZeros(ax,bx)
@@ -710,4 +710,21 @@ import QuasiArrays: AbstractQuasiFill
         @test stringmime("text/plain",zeros(Inclusion([1,2,3]))) == "zeros(Inclusion([1, 2, 3]))"
         @test stringmime("text/plain",fill(2,Inclusion([1,2,3]))) == "fill(2, Inclusion([1, 2, 3]))"
     end
+
+    @testset "Mul" begin
+        A = QuasiArray(randn(3,3), 0:0.5:1, Base.OneTo(3))
+        @test A * Zeros(axes(A,2)) ≡ QuasiZeros(axes(A,1))
+        B = QuasiArray(randn(3,3), 0:0.5:1,1:0.5:2)
+        @test B * QuasiZeros(axes(B,2)) ≡ QuasiZeros(axes(B,1))
+
+        @test_throws DimensionMismatch A * Zeros(2)
+        @test_throws DimensionMismatch B * QuasiZeros(0:0.5:1)
+        @test_throws DimensionMismatch FillArrays.mult_zeros(B, QuasiZeros(0:0.5:1))
+    end
+
+    @testset "isone" begin
+        @test isone(QuasiFill(1, 0:0.5:1))
+        @test isone(QuasiOnes(0:0.5:1))
+        @test !isone(QuasiZeros(0:0.5:1))
+    end
 end
diff --git a/test/test_quasilazy.jl b/test/test_quasilazy.jl
@@ -39,6 +39,9 @@ Base.getindex(A::MyQuasiLazyMatrix, x::Float64, y::Float64) = A.A[x,y]
                 @test M/2 ≈ 2\M ≈ A*B/2
                 @test (2\M).args[2] == (M/2).args[2] == B/2
 
+                @test colsupport(M, 1) == axes(M,1)
+                @test rowsupport(M, 1) == axes(M,2)
+
                 M = ApplyQuasiArray(*, B', A')
                 @test M ≈ B'A'
                 @test M[[1,3], [0,0.5]] ≈ (B'A')[[1,3], [0,0.5]]
