Add MulQuasiArray

Author: dlfivefifty

src/ContinuumArrays.jl

@@ -10,7 +10,7 @@ import BandedMatrices: AbstractBandedLayout, _BandedMatrix
 include("QuasiArrays/QuasiArrays.jl")
 using .QuasiArrays
 import .QuasiArrays: cardinality, checkindex, QuasiAdjoint, QuasiTranspose, slice, QSlice, SubQuasiArray,
-                    QuasiDiagonal
+                    QuasiDiagonal, MulQuasiArray, MulQuasiMatrix, MulQuasiVector, QuasiMatMulMat
 
 export Spline, LinearSpline, HeavisideSpline, DiracDelta, Derivative, JacobiWeight, Jacobi, Legendre
 

src/QuasiArrays/QuasiArrays.jl

@@ -9,7 +9,7 @@ import Base: @_inline_meta, DimOrInd, OneTo, @_propagate_inbounds_meta, @_noinli
                 index_shape, to_shape, unsafe_length, @nloops, @ncall, Slice, unalias
 import Base: ViewIndex, Slice, ScalarIndex, RangeIndex, view, viewindexing, ensure_indexable, index_dimsum,
                 check_parent_index_match, reindex, _isdisjoint, unsafe_indices,
-                parentindices, reverse
+                parentindices, reverse, ndims
 import Base: *, /, \, +, -, inv
 import Base: exp, log, sqrt,
           cos, sin, tan, csc, sec, cot,
@@ -23,7 +23,8 @@ import Base.Broadcast: materialize
 import LinearAlgebra: transpose, adjoint, checkeltype_adjoint, checkeltype_transpose, Diagonal,
                         AbstractTriangular
 
-import LazyArrays: MemoryLayout, UnknownLayout, Mul2
+import LazyArrays: MemoryLayout, UnknownLayout, Mul2, _materialize, MulLayout, ⋆, rmaterialize,
+                    _rmaterialize, _lmaterialize, flatten, _flatten
 
 export AbstractQuasiArray, AbstractQuasiMatrix, AbstractQuasiVector, materialize
 

src/QuasiArrays/matmul.jl

@@ -2,6 +2,9 @@
 const QuasiArrayMulArray{styleA, styleB, p, q, T, V} =
     Mul2{styleA, styleB, <:AbstractQuasiArray{T,p}, <:AbstractArray{V,q}}
 
+const ArrayMulQuasiArray{styleA, styleB, p, q, T, V} =
+    Mul2{styleA, styleB, <:AbstractArray{T,p}, <:AbstractQuasiArray{V,q}}
+
 const QuasiArrayMulQuasiArray{styleA, styleB, p, q, T, V} =
     Mul2{styleA, styleB, <:AbstractQuasiArray{T,p}, <:AbstractQuasiArray{V,q}}
 ####
@@ -13,7 +16,7 @@ const QuasiMatMulVec{styleA, styleB, T, V} = QuasiArrayMulArray{styleA, styleB,
 function getindex(M::QuasiMatMulVec, k::Real)
     A,B = M.factors
     ret = zero(eltype(M))
-    @inbounds for j = 1:size(A,2)
+    @inbounds for j in axes(A,2)
         ret += A[k,j] * B[j]
     end
     ret
@@ -22,7 +25,7 @@ end
 function getindex(M::QuasiMatMulVec, k::AbstractArray)
     A,B = M.factors
     ret = zeros(eltype(M),length(k))
-    @inbounds for j = 1:size(A,2)
+    @inbounds for j in axes(A,2)
         ret .+= view(A,k,j) .* B[j]
     end
     ret
@@ -41,3 +44,114 @@ inv(A::AbstractQuasiArray) = materialize(Inv(A))
 
 *(A::AbstractQuasiArray, B::Mul) = materialize(Mul(A, B.factors...))
 *(A::Mul, B::AbstractQuasiArray) = materialize(Mul(A.factors..., B))
+
+
+####
+# MulQuasiArray
+#####
+
+struct MulQuasiArray{T, N, MUL<:Mul} <: AbstractQuasiArray{T,N}
+    mul::MUL
+end
+
+const MulQuasiVector{T, MUL<:Mul} = MulQuasiArray{T, 1, MUL}
+const MulQuasiMatrix{T, MUL<:Mul} = MulQuasiArray{T, 2, MUL}
+
+const Vec = MulQuasiVector
+
+
+MulQuasiArray{T,N}(M::MUL) where {T,N,MUL<:Mul} = MulQuasiArray{T,N,MUL}(M)
+MulQuasiArray{T}(M::Mul) where {T} = MulQuasiArray{T,ndims(M)}(M)
+MulQuasiArray(M::Mul) = MulQuasiArray{eltype(M)}(M)
+MulQuasiVector(M::Mul) = MulQuasiVector{eltype(M)}(M)
+MulQuasiMatrix(M::Mul) = MulQuasiMatrix{eltype(M)}(M)
+
+MulQuasiArray(factors...) = MulQuasiArray(Mul(factors...))
+MulQuasiArray{T}(factors...) where T = MulQuasiArray{T}(Mul(factors...))
+MulQuasiArray{T,N}(factors...) where {T,N} = MulQuasiArray{T,N}(Mul(factors...))
+MulQuasiVector(factors...) = MulQuasiVector(Mul(factors...))
+MulQuasiMatrix(factors...) = MulQuasiMatrix(Mul(factors...))
+
+_MulArray(factors...) = MulQuasiArray(factors...)
+_MulArray(factors::AbstractArray...) = MulArray(factors...)
+
+axes(A::MulQuasiArray) = axes(A.mul)
+size(A::MulQuasiArray) = map(length, axes(A))
+
+IndexStyle(::MulQuasiArray{<:Any,1}) = IndexLinear()
+
+==(A::MulQuasiArray, B::MulQuasiArray) = A.mul == B.mul
+
+@propagate_inbounds getindex(A::MulQuasiArray, kj::Real...) = A.mul[kj...]
+
+*(A::MulQuasiArray, B::MulQuasiArray) = A.mul * B.mul
+*(A::MulQuasiArray, B::Mul) = A.mul * B
+*(A::Mul, B::MulQuasiArray) = A * B.mul
+*(A::MulQuasiArray, B::AbstractQuasiArray) = A.mul * B
+*(A::AbstractQuasiArray, B::MulQuasiArray) = A * B.mul
+*(A::MulQuasiArray, B::AbstractArray) = A.mul * B
+*(A::AbstractArray, B::MulQuasiArray) = A * B.mul
+
+adjoint(A::MulQuasiArray) = MulQuasiArray(reverse(adjoint.(A.mul.factors))...)
+transpose(A::MulQuasiArray) = MulQuasiArray(reverse(transpose.(A.mul.factors))...)
+
+
+MemoryLayout(M::MulQuasiArray) = MulLayout(MemoryLayout.(M.mul.factors))
+
+
+
+####
+# Matrix * Array
+####
+
+_flatten(A::MulQuasiArray, B...) = _flatten(A.mul.factors..., B...)
+flatten(A::MulQuasiArray) = MulQuasiArray(Mul(_flatten(A.mul.factors...)))
+
+
+# the default is always Array
+
+_materialize(M::QuasiArrayMulArray, _) = MulQuasiArray(M)
+_materialize(M::ArrayMulQuasiArray, _) = MulQuasiArray(M)
+_materialize(M::QuasiArrayMulQuasiArray, _) = MulQuasiArray(M)
+
+
+
+# if multiplying two MulQuasiArrays simplifies the arguments, we materialize,
+# otherwise we leave it as a lazy object
+_mulquasi_join(As, M::MulQuasiArray, Cs) = MulQuasiArray(As..., M.mul.factors..., Cs...)
+_mulquasi_join(As, B, Cs) = *(As..., B, Cs...)
+
+
+function _materialize(M::Mul2{<:Any,<:Any,<:MulQuasiArray,<:MulQuasiArray}, _)
+    As, Bs = M.factors
+    _mul_join(reverse(tail(reverse(As))), last(As) * first(Bs), tail(Bs))
+end
+
+
+function _materialize(M::Mul2{<:Any,<:Any,<:MulQuasiArray,<:AbstractQuasiArray}, _)
+    As, B = M.factors
+    ⋆(As.mul.factors..., B)
+end
+
+function _materialize(M::Mul2{<:Any,<:Any,<:AbstractQuasiArray,<:MulQuasiArray}, _)
+    A, Bs = M.factors
+    *(A, Bs.mul.factors...)
+end
+
+# A MulQuasiArray can't be materialized further left-to-right, so we do right-to-left
+function _materialize(M::Mul2{<:Any,<:Any,<:MulQuasiArray,<:AbstractArray}, _)
+    As, B = M.factors
+    ⋆(As.mul.factors..., B)
+end
+
+function _lmaterialize(A::MulQuasiArray, B, C...)
+    As = A.mul.factors
+    flatten(_MulArray(reverse(tail(reverse(As)))..., _lmaterialize(last(As), B, C...)))
+end
+
+
+
+function _rmaterialize(Z::MulQuasiArray, Y, W...)
+    Zs = Z.mul.factors
+    flatten(_MulArray(_rmaterialize(first(Zs), Y, W...), tail(Zs)...))
+end

src/bases/jacobi.jl

@@ -42,7 +42,7 @@ end
 function materialize(M::Mul2{<:Any,<:Any,<:Derivative{<:Any,<:ChebyshevInterval},<:Jacobi})
     D, S = M.factors
     A = PInv(Jacobi(S.b+1,S.a+1))*D*S
-    Mul(Jacobi(S.b+1,S.a+1), A)
+    MulQuasiMatrix(Jacobi(S.b+1,S.a+1), A)
 end
 
 # pinv(Legendre())D*W*Jacobi(true,true)
@@ -62,7 +62,7 @@ function materialize(M::Mul{<:Tuple,<:Tuple{<:Derivative{<:Any,<:ChebyshevInterv
     w = parent(W)
     (w.a && S.a && w.b && S.b) || throw(ArgumentError())
     A = pinv(Legendre{eltype(M)}())*D*W*S
-    Mul(Legendre(), A)
+    MulQuasiMatrix(Legendre(), A)
 end
 
 function materialize(M::Mul{<:Tuple,<:Tuple{<:PInv{<:Any,<:Jacobi{Bool}},

src/bases/splines.jl

@@ -78,10 +78,10 @@ function materialize(M::Mul2{<:Any,<:Any,<:QuasiAdjoint{<:Any,<:HeavisideSpline}
 end
 
 ## Derivative
-function copyto!(dest::Mul2{<:Any,<:Any,<:HeavisideSpline},
+function copyto!(dest::MulQuasiMatrix{<:Any,<:Mul2{<:Any,<:Any,<:HeavisideSpline}},
                  M::Mul2{<:Any,<:Any,<:Derivative,<:LinearSpline})
     D, L = M.factors
-    H, A = dest.factors
+    H, A = dest.mul.factors
     x = H.points
 
     axes(dest) == axes(M) || throw(DimensionMismatch("axes must be same"))
@@ -98,7 +98,7 @@ end
 function similar(M::Mul2{<:Any,<:Any,<:Derivative,<:LinearSpline}, ::Type{T}) where T
     D, B = M.factors
     n = size(B,2)
-    Mul(HeavisideSpline{T}(B.points),
+    MulQuasiMatrix(HeavisideSpline{T}(B.points),
         BandedMatrix{T}(undef, (n-1,n), (0,1)))
 end
 

test/runtests.jl

@@ -1,7 +1,8 @@
 using ContinuumArrays, LazyArrays, IntervalSets, FillArrays, LinearAlgebra, BandedMatrices, Test,
     InfiniteArrays
     import ContinuumArrays: ℵ₁, materialize
-    import ContinuumArrays.QuasiArrays: SubQuasiArray
+    import ContinuumArrays.QuasiArrays: SubQuasiArray, MulQuasiMatrix, Vec
+    import LazyArrays: rmaterialize, ⋆
 
 @testset "DiracDelta" begin
     δ = DiracDelta(-1..3)
@@ -75,7 +76,16 @@ end
     L = LinearSpline([1,2,3])
     f = L*[1,2,4]
     D = Derivative(axes(L,1))
+    @test D*L isa MulQuasiMatrix
+
+    fp = (D*L)*[1,2,4]
+    @test fp isa Vec
+    @test length(fp.mul.factors) == 2
+    @test fp[1.1] ≈ 1
+    @test fp[2.2] ≈ 2
+
     fp = D*f
+    @test length(fp.mul.factors) == 2
 
     @test fp[1.1] ≈ 1
     @test fp[2.2] ≈ 2
@@ -86,17 +96,19 @@ end
     L = LinearSpline(0:2)
 
     D = Derivative(axes(L,1))
-    M = materialize(Mul(D',D,L))
-    DL = D*L
-    @test M.factors == tuple(D', (D*L).factors...)
+    M = rmaterialize(Mul(D',D*L))
+    @test length(M.mul.factors) == 3
+    @test last(M.mul.factors) isa BandedMatrix
 
-    @test materialize(Mul(L', D', D, L)) == (L'D'*D*L) ==
-        [1.0 -1 0; -1.0 2.0 -1.0; 0.0 -1.0 1.0]
+    @test M.mul.factors == rmaterialize(Mul(D',D,L)).mul.factors ==
+        ⋆(D',D,L).mul.factors == *(D',D,L).mul.factors
 
-    @test materialize(Mul(L', D', D, L)) isa BandedMatrix
-    @test (L'D'*D*L) isa BandedMatrix
+    @test (L'D') isa MulQuasiMatrix
+    A = (L'D') * (D*L)
+    @test A == (D*L)'*(D*L) == [1.0 -1 0; -1.0 2.0 -1.0; 0.0 -1.0 1.0]
 
-    @test bandwidths(materialize(L'D'*D*L)) == (1,1)
+    @test A isa MulArray
+    @test bandwidths(A) == (1,1)
 end
 
 @testset "Views" begin
@@ -121,7 +133,7 @@ end
 
 @testset "Subindex of splines" begin
     L = LinearSpline(range(0,stop=1,length=10))
-    @test L[:,2:end-1] isa Mul
+    @test L[:,2:end-1] isa MulQuasiMatrix
     @test_broken L[:,2:end-1][0.1,1] == L[0.1,2]
     v = randn(8)
     f = L[:,2:end-1] * v
@@ -132,7 +144,7 @@ end
     L = LinearSpline(range(0,stop=1,length=10))
     B = L[:,2:end-1] # Zero dirichlet by dropping first and last spline
     D = Derivative(axes(L,1))
-    Δ = -(B'D'D*B) # Weak Laplacian
+    Δ = -((B'D')*(D*B)) # Weak Laplacian
 
     f = L*exp.(L.points) # project exp(x)
     u = B * (Δ \ (B'f))
@@ -145,7 +157,8 @@ end
     L = LinearSpline(range(0,stop=1,length=10))
     B = L[:,2:end-1] # Zero dirichlet by dropping first and last spline
     D = Derivative(axes(L,1))
-    A = -(B'D'D*B) + 100^2*B'B # Weak Laplacian
+
+    A = -((B'D')*(D*B)) + 100^2*B'B # Weak Laplacian
 
     f = L*exp.(L.points) # project exp(x)
     u = B * (A \ (B'f))
@@ -159,7 +172,7 @@ end
     D = Derivative(axes(W,1))
     P = Legendre()
 
-    A = @inferred(PInv(Jacobi(2,2))*D*S)
+    A = @inferred(PInv(Jacobi(2,2))*(D*S))
     @test A isa BandedMatrix
     @test size(A) == (∞,∞)
     @test A[1:10,1:10] == diagm(1 => 1:0.5:5)