Support BroadcastLayout multiplication (#68)

* support lazy broadcasting with mul

* Update quasibroadcast.jl

* avoid broadcasted

* Remove QuasiIndices, unused code

* simplifiable for D^2, use in it broadcast

* Add reducedim (sum, etc.)

* Update QuasiArrays.jl

* move memorylayout _sum overloading

* Update quasireducedim.jl

* increase coverage

* Require Julia v1.6

* Increase coverage

* Test (x .* D) * y

* (x .* D^2) * y

Author: dlfivefifty

.github/workflows/ci.yml

@@ -10,8 +10,7 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.5'
-          - '^1.6.0-0'
+          - '1.6'
         os:
           - ubuntu-latest
           - macOS-latest

Project.toml

@@ -1,7 +1,7 @@
 name = "QuasiArrays"
 uuid = "c4ea9172-b204-11e9-377d-29865faadc5c"
 authors = ["Sheehan Olver <solver@mac.com>"]
-version = "0.5.2"
+version = "0.6"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
@@ -15,9 +15,9 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 ArrayLayouts = "0.7"
 DomainSets = "0.4"
 FillArrays = "0.11"
-LazyArrays = "0.21"
+LazyArrays = "0.21.5"
 StaticArrays = "1"
-julia = "1.5"
+julia = "1.6"
 
 [extras]
 Base64 = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

src/QuasiArrays.jl

@@ -14,7 +14,7 @@ import Base: Slice, IdentityUnitRange, ScalarIndex, RangeIndex, view, viewindexi
                 check_parent_index_match, reindex, _isdisjoint, unsafe_indices, _unsafe_ind2sub,
                 _ind2sub, _sub2ind, _ind2sub_recurse, _lookup, SubArray,
                 parentindices, reverse, ndims, checkbounds, uncolon,
-                promote_shape, maybeview, unsafe_view, checkindex, checkbounds_indices,
+                maybeview, unsafe_view, checkindex, checkbounds_indices,
                 throw_boundserror, rdims, replace_in_print_matrix, show, summary,
                 hcat, vcat, hvcat
 import Base: *, /, \, +, -, ^, inv
@@ -26,7 +26,10 @@ import Base: exp, log, sqrt,
 import Base: Array, Matrix, Vector
 import Base: union, intersect, sort, sort!
 import Base: conj, real, imag
-import Base: sum, cumsum, diff
+# reducedim.jl imports
+import Base: prod, sum, cumsum, diff, add_sum, mul_prod, mapreduce, max, min, count, _count, any, _any, all, _all, _sum, _prod, _mapreduce, reduced_index, check_reducedims
+import Base: BitInteger, IEEEFloat, uniontypes, _InitialValue, safe_tail, reducedim1, _simple_count
+
 import Base: ones, zeros, one, zero, fill
 
 import Base.Broadcast: materialize, materialize!, BroadcastStyle, AbstractArrayStyle, Style, broadcasted, Broadcasted, Unknown,
@@ -45,7 +48,8 @@ import LazyArrays: MemoryLayout, UnknownLayout, Mul, ApplyLayout, BroadcastLayou
                     rowsupport, colsupport, tuple_type_memorylayouts, applylayout, broadcastlayout,
                     LdivStyle, most, InvLayout, PInvLayout, sub_materialize, lazymaterialize,
                     _mul, rowsupport, DiagonalLayout, adjointlayout, transposelayout, conjlayout,
-                    sublayout, call, LazyArrayStyle, layout_getindex, _broadcast2broadcastarray, _applyarray_summary, _broadcastarray_summary
+                    sublayout, call, LazyArrayStyle, layout_getindex, _broadcast2broadcastarray, _applyarray_summary, _broadcastarray_summary,
+                    _broadcasted_mul, simplifiable, simplify
 
 import Base.IteratorsMD
 
@@ -86,6 +90,7 @@ include("quasireshapedarray.jl")
 include("quasipermutedims.jl")
 include("quasibroadcast.jl")
 include("abstractquasiarraymath.jl")
+include("quasireducedim.jl")
 
 include("quasiarray.jl")
 include("quasiarraymath.jl")

src/abstractquasiarray.jl

@@ -489,14 +489,6 @@ their component parts.  A typical definition for an array that wraps a parent is
 dataids(A::AbstractQuasiArray) = (UInt(objectid(A)),)
 
 
-## structured matrix methods ##
-replace_in_print_matrix(A::AbstractQuasiMatrix,i,j,s::AbstractString) = s
-replace_in_print_matrix(A::AbstractQuasiVector,i,j,s::AbstractString) = s
-
-## Concatenation ##
-eltypeof(x::AbstractQuasiArray) = eltype(x)
-
-
 ## Reductions and accumulates ##
 
 function isequal(A::AbstractQuasiArray, B::AbstractQuasiArray)
@@ -512,17 +504,6 @@ function isequal(A::AbstractQuasiArray, B::AbstractQuasiArray)
     return true
 end
 
-function cmp(A::AbstractQuasiVector, B::AbstractQuasiVector)
-    for (a, b) in zip(A, B)
-        if !isequal(a, b)
-            return isless(a, b) ? -1 : 1
-        end
-    end
-    return cmp(length(A), length(B))
-end
-
-isless(A::AbstractQuasiVector, B::AbstractQuasiVector) = cmp(A, B) < 0
-
 function (==)(A::AbstractQuasiArray, B::AbstractQuasiArray)
     if axes(A) != axes(B)
         return false
@@ -539,192 +520,6 @@ function (==)(A::AbstractQuasiArray, B::AbstractQuasiArray)
     return anymissing ? missing : true
 end
 
-_lookup(ind, r::Inclusion) = ind
-
-_ind2sub(dims::NTuple{N,Any}, ind) where N = (@_inline_meta; _ind2sub_recurse(dims, ind-1))
-_ind2sub(inds::QuasiIndices, ind)     = (@_inline_meta; _ind2sub_recurse(inds, ind-1))
-_ind2sub(inds::Tuple{Inclusion{<:Any},AbstractUnitRange{<:Integer}}, ind)     = (@_inline_meta; _ind2sub_recurse(inds, ind-1))
-_ind2sub(inds::Tuple{AbstractUnitRange{<:Integer},Inclusion{<:Any}}, ind)     = (@_inline_meta; _ind2sub_recurse(inds, ind-1))
-
-function _ind2sub(inds::Union{NTuple{N,Any},QuasiIndices{N}}, ind::AbstractQuasiVector) where N
-    M = length(ind)
-    t = ntuple(n->similar(ind),Val(N))
-    for (i,idx) in pairs(IndexLinear(), ind)
-        sub = _ind2sub(inds, idx)
-        for j = 1:N
-            t[j][i] = sub[j]
-        end
-    end
-    t
-end
-
-
-## map over arrays ##
-
-## transform any set of dimensions
-## dims specifies which dimensions will be transformed. for example
-## dims==1:2 will call f on all slices A[:,:,...]
-
-function mapslices(f, A::AbstractQuasiArray; dims)
-    if isempty(dims)
-        return map(f,A)
-    end
-    if !isa(dims, AbstractQuasiVector)
-        dims = [dims...]
-    end
-
-    dimsA = [axes(A)...]
-    ndimsA = ndims(A)
-    alldims = [1:ndimsA;]
-
-    otherdims = setdiff(alldims, dims)
-
-    idx = Any[first(ind) for ind in axes(A)]
-    itershape   = tuple(dimsA[otherdims]...)
-    for d in dims
-        idx[d] = Slice(axes(A, d))
-    end
-
-    # Apply the function to the first slice in order to determine the next steps
-    Aslice = A[idx...]
-    r1 = f(Aslice)
-    # In some cases, we can re-use the first slice for a dramatic performance
-    # increase. The slice itself must be mutable and the result cannot contain
-    # any mutable containers. The following errs on the side of being overly
-    # strict (#18570 & #21123).
-    safe_for_reuse = isa(Aslice, StridedArray) &&
-                     (isa(r1, Number) || (isa(r1, AbstractQuasiArray) && eltype(r1) <: Number))
-
-    # determine result size and allocate
-    Rsize = copy(dimsA)
-    # TODO: maybe support removing dimensions
-    if !isa(r1, AbstractQuasiArray) || ndims(r1) == 0
-        # If the result of f on a single slice is a scalar then we add singleton
-        # dimensions. When adding the dimensions, we have to respect the
-        # index type of the input array (e.g. in the case of OffsetArrays)
-        tmp = similar(Aslice, typeof(r1), reduced_indices(Aslice, 1:ndims(Aslice)))
-        tmp[firstindex(tmp)] = r1
-        r1 = tmp
-    end
-    nextra = max(0, length(dims)-ndims(r1))
-    if eltype(Rsize) == Int
-        Rsize[dims] = [size(r1)..., ntuple(d->1, nextra)...]
-    else
-        Rsize[dims] = [axes(r1)..., ntuple(d->OneTo(1), nextra)...]
-    end
-    R = similar(r1, tuple(Rsize...,))
-
-    ridx = Any[map(first, axes(R))...]
-    for d in dims
-        ridx[d] = axes(R,d)
-    end
-
-    concatenate_setindex!(R, r1, ridx...)
-
-    nidx = length(otherdims)
-    indices = Iterators.drop(CartesianIndices(itershape), 1) # skip the first element, we already handled it
-    inner_mapslices!(safe_for_reuse, indices, nidx, idx, otherdims, ridx, Aslice, A, f, R)
-end
-
-concatenate_setindex!(R, X::AbstractQuasiArray, I...) = (R[I...] = X)
-
-## 1 argument
-
-function map!(f::F, dest::AbstractQuasiArray, A::AbstractQuasiArray) where F
-    for (i,j) in zip(eachindex(dest),eachindex(A))
-        dest[i] = f(A[j])
-    end
-    return dest
-end
-
-# map on collections
-map(f, A::AbstractQuasiArray) = collect_similar(A, Generator(f,A))
-
-## 2 argument
-function map!(f::F, dest::AbstractQuasiArray, A::AbstractQuasiArray, B::AbstractQuasiArray) where F
-    for (i, j, k) in zip(eachindex(dest), eachindex(A), eachindex(B))
-        dest[i] = f(A[j], B[k])
-    end
-    return dest
-end
-
-
-function map_n!(f::F, dest::AbstractQuasiArray, As) where F
-    for i = LinearIndices(As[1])
-        dest[i] = f(ith_all(i, As)...)
-    end
-    return dest
-end
-
-map!(f::F, dest::AbstractQuasiArray, As::AbstractQuasiArray...) where {F} = map_n!(f, dest, As)
-
-
-## hashing AbstractQuasiArray ##
-
-function hash(A::AbstractQuasiArray, h::UInt)
-    h = hash(AbstractQuasiArray, h)
-    # Axes are themselves AbstractQuasiArrays, so hashing them directly would stack overflow
-    # Instead hash the tuple of firsts and lasts along each dimension
-    h = hash(map(first, axes(A)), h)
-    h = hash(map(last, axes(A)), h)
-    isempty(A) && return h
-
-    # Goal: Hash approximately log(N) entries with a higher density of hashed elements
-    # weighted towards the end and special consideration for repeated values. Colliding
-    # hashes will often subsequently be compared by equality -- and equality between arrays
-    # works elementwise forwards and is short-circuiting. This means that a collision
-    # between arrays that differ by elements at the beginning is cheaper than one where the
-    # difference is towards the end. Furthermore, blindly choosing log(N) entries from a
-    # sparse array will likely only choose the same element repeatedly (zero in this case).
-
-    # To achieve this, we work backwards, starting by hashing the last element of the
-    # array. After hashing each element, we skip `fibskip` elements, where `fibskip`
-    # is pulled from the Fibonacci sequence -- Fibonacci was chosen as a simple
-    # ~O(log(N)) algorithm that ensures we don't hit a common divisor of a dimension
-    # and only end up hashing one slice of the array (as might happen with powers of
-    # two). Finally, we find the next distinct value from the one we just hashed.
-
-    # This is a little tricky since skipping an integer number of values inherently works
-    # with linear indices, but `findprev` uses `keys`. Hoist out the conversion "maps":
-    ks = keys(A)
-    key_to_linear = LinearIndices(ks) # Index into this map to compute the linear index
-    linear_to_key = vec(ks)           # And vice-versa
-
-    # Start at the last index
-    keyidx = last(ks)
-    linidx = key_to_linear[keyidx]
-    fibskip = prevfibskip = oneunit(linidx)
-    n = 0
-    while true
-        n += 1
-        # Hash the current key-index and its element
-        elt = A[keyidx]
-        h = hash(keyidx=>elt, h)
-
-        # Skip backwards a Fibonacci number of indices -- this is a linear index operation
-        linidx = key_to_linear[keyidx]
-        linidx <= fibskip && break
-        linidx -= fibskip
-        keyidx = linear_to_key[linidx]
-
-        # Only increase the Fibonacci skip once every N iterations. This was chosen
-        # to be big enough that all elements of small arrays get hashed while
-        # obscenely large arrays are still tractable. With a choice of N=4096, an
-        # entirely-distinct 8000-element array will have ~75% of its elements hashed,
-        # with every other element hashed in the first half of the array. At the same
-        # time, hashing a `typemax(Int64)`-length Float64 range takes about a second.
-        if rem(n, 4096) == 0
-            fibskip, prevfibskip = fibskip + prevfibskip, fibskip
-        end
-
-        # Find a key index with a value distinct from `elt` -- might be `keyidx` itself
-        keyidx = findprev(!isequal(elt), A, keyidx)
-        keyidx === nothing && break
-    end
-
-    return h
-end
-
 
 ##
 # show

src/indices.jl

@@ -11,12 +11,8 @@ IndexStyle(A::AbstractQuasiArray, B::AbstractQuasiArray...) = IndexStyle(IndexSt
 IndexStyle(A::AbstractQuasiArray, B::AbstractArray...) = IndexStyle(IndexStyle(A), IndexStyle(B...))
 
 
-function promote_shape(a::AbstractQuasiArray, b::AbstractQuasiArray)
-    promote_shape(axes(a), axes(b))
-end
-
-const QuasiIndices{N} = NTuple{N,Union{AbstractQuasiVector{<:Number},AbstractVector{<:Number}}}
-function promote_shape(a::QuasiIndices, b::QuasiIndices)
+promote_shape(a::AbstractQuasiArray, b::AbstractQuasiArray) = promote_shape(axes(a), axes(b))
+function promote_shape(a::Tuple, b::Tuple)
     if length(a) < length(b)
         return promote_shape(b, a)
     end

src/lazyquasiarrays.jl

@@ -124,10 +124,10 @@ _BroadcastQuasiArray(bc::Broadcasted) = BroadcastQuasiArray{combine_eltypes(bc.f
 BroadcastQuasiArray(bc::Broadcasted{S}) where S =
     _BroadcastQuasiArray(instantiate(Broadcasted{S}(bc.f, _broadcast2broadcastarray(bc.args...))))
 BroadcastQuasiArray(b::BroadcastQuasiArray) = b
-BroadcastQuasiArray(f, A, As...) = BroadcastQuasiArray(instantiate(broadcasted(f, A, As...)))
+BroadcastQuasiArray(f, A, As...) = BroadcastQuasiArray{combine_eltypes(f, (A, As...))}(f, A, As...)
 BroadcastQuasiVector(f, A, As...) = BroadcastQuasiVector{combine_eltypes(f, (A, As...))}(f, A, As...)
 BroadcastQuasiMatrix(f, A, As...) = BroadcastQuasiMatrix{combine_eltypes(f, (A, As...))}(f, A, As...)
-BroadcastQuasiArray{T}(f, A, As...) where {T} = BroadcastQuasiArray{T}(instantiate(broadcasted(f, A, As...)))
+BroadcastQuasiArray{T}(f, A, As...) where {T} = BroadcastQuasiArray{T,length(axes(broadcasted(f, A, As...)))}(f, A, As...)
 BroadcastQuasiArray{T,N}(f, A, As...) where {T,N} = BroadcastQuasiArray{T,N,typeof(f),typeof((A, As...))}(f, (A, As...))
 
 @inline BroadcastQuasiArray(A::AbstractQuasiArray) = BroadcastQuasiArray(call(A), arguments(A)...)
@@ -201,6 +201,9 @@ _broadcast_mul_arguments(a, B) = __broadcast_mul_arguments(a, _mul_arguments(B).
 _mul_arguments(A::BroadcastQuasiMatrix{<:Any,typeof(*),<:Tuple{AbstractQuasiVector,AbstractQuasiMatrix}}) =
     _broadcast_mul_arguments(A.args...)
 
+broadcasted(::LazyQuasiArrayStyle{2}, ::typeof(*), a::AbstractQuasiVector, B::ApplyQuasiMatrix{<:Any,typeof(*)}) = 
+    *(_broadcast_mul_arguments(a, B)...)
+
 ndims(M::Applied{LazyQuasiArrayApplyStyle,typeof(*)}) = ndims(last(M.args))
 
 call(a::AbstractQuasiArray) = call(MemoryLayout(typeof(a)), a)
@@ -230,3 +233,10 @@ function *(App::ApplyQuasiMatrix{<:Any,typeof(^),<:Tuple{<:AbstractQuasiMatrix{T
     p == 0 && return copy(b)
     return A*(ApplyQuasiMatrix(^,A,p-1)*b)
 end
+
+function simplifiable(::typeof(*), App::ApplyQuasiMatrix{<:Any,typeof(^),<:Tuple{<:AbstractQuasiMatrix{T},<:Integer}}, b::AbstractQuasiArray) where T
+    A,p = arguments(App)
+    simplifiable(*, A, b)
+end
+
+

src/matmul.jl

@@ -99,6 +99,13 @@ copy(M::Mul{<:Any,QuasiArrayLayout}) = _quasi_mul(M, axes(M))
 copy(M::Mul{<:AbstractLazyLayout,QuasiArrayLayout}) = ApplyQuasiArray(M)
 copy(M::Mul{ApplyLayout{typeof(\)},QuasiArrayLayout}) = ApplyQuasiArray(M)
 copy(M::Mul{QuasiArrayLayout,<:AbstractLazyLayout}) = ApplyQuasiArray(M)
+for op in (:+, :-)
+    @eval begin
+        copy(M::Mul{BroadcastLayout{typeof($op)},QuasiArrayLayout}) = simplify(M)
+        copy(M::Mul{QuasiArrayLayout,BroadcastLayout{typeof($op)}}) = simplify(M)
+    end
+end
+@inline copy(M::Mul{BroadcastLayout{typeof(*)},QuasiArrayLayout}) = copy(Mul{BroadcastLayout{typeof(*)},UnknownLayout}(M.A,M.B))
 
 
 LazyArrays._vec_mul_view(a::AbstractQuasiVector, kr, ::Colon) = view(a, kr)

src/quasiarray.jl

@@ -33,6 +33,8 @@ QuasiVector(par::AbstractVector{T}, axes::Tuple{AbstractVector}) where T =
 
 QuasiArray{T}(par::AbstractArray{<:Any,N}, axes::NTuple{N,AbstractVector}) where {T,N} = 
     QuasiArray{T,N,typeof(axes)}(par, axes)
+QuasiArray{T}(par::AbstractArray{T,N}, axes::NTuple{N,AbstractVector}) where {T,N} = 
+    QuasiArray{T,N,typeof(axes)}(par, axes)
 QuasiMatrix{T}(par::AbstractMatrix, axes::NTuple{2,AbstractVector}) where T = 
     QuasiMatrix{T,typeof(axes)}(par, axes)
 QuasiVector{T}(par::AbstractVector, axes::Tuple{AbstractVector}) where T = 
@@ -81,7 +83,7 @@ QuasiVector(par::AbstractVector{T}, axes::AbstractArray) where {T} =
     QuasiVector(par, (axes,))
 
 
-QuasiArray(a::AbstractQuasiArray) = QuasiArray(a[map(collect,axes(a))...], axes(a))
+QuasiArray(a::AbstractQuasiArray{T}) where T = QuasiArray{T}(a[map(collect,axes(a))...], axes(a))
 QuasiArray{T}(a::AbstractQuasiArray) where T = QuasiArray(convert(AbstractArray{T},a[map(collect,axes(a))...]), axes(a))
 QuasiArray{T,N}(a::AbstractQuasiArray{<:Any,N}) where {T,N} = QuasiArray(convert(AbstractArray{T,N},a[map(collect,axes(a))...]), axes(a))
 QuasiArray{T,N,AXES}(a::AbstractQuasiArray{<:Any,N}) where {T,N,AXES} = QuasiArray{T,N,AXES}(Array{T}(a), axes(a))
@@ -134,3 +136,9 @@ end
     axes(A) == axes(B) && A.parent == B
 ==(B::AbstractArray{V,N}, A::QuasiArray{T,N,NTuple{N,OneTo{Int}}}) where {T,V,N} =
     A == B
+
+
+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