Kronecker operator acting on LowRankFun/ProductFun (#242)

* multiply kronop with prodfun/lowrankfun

* kron operator action on LowRankFun/ProductFun

* specialize mul for constant second op

* cast ProductFun to LowRankFun in getindex

* fix space in kron with const 2nd op

* views in ProductFun chop

Author: jishnub

Project.toml

@@ -1,6 +1,6 @@
 name = "ApproxFunBase"
 uuid = "fbd15aa5-315a-5a7d-a8a4-24992e37be05"
-version = "0.7.23"
+version = "0.7.24"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

src/Multivariate/Multivariate.jl

@@ -44,7 +44,7 @@ Fun(f::MultivariateFun,sp::Space) = Fun(Fun(f),sp)
 
 Fun(f,d1::Domain,d2::Domain) = Fun(f,d1*d2)
 
-coefficients(f::BivariateFun,sp::TensorSpace)=coefficients(f,sp[1],sp[2])
+coefficients(f::BivariateFun,sp::TensorSpace)=coefficients(f, factors(sp)...)
 
 
 

src/Multivariate/ProductFun.jl

@@ -62,11 +62,22 @@ true
 ```
 """
 function ProductFun(cfs::AbstractMatrix{T},sp::AbstractProductSpace{Tuple{S,V},DD};
-  tol::Real=100eps(T),chopping::Bool=false) where {S<:UnivariateSpace,V<:UnivariateSpace,T<:Number,DD}
+    tol::Real=100eps(T),chopping::Bool=false) where {S<:UnivariateSpace,V<:UnivariateSpace,T<:Number,DD}
     if chopping
-        ncfs,kend=norm(cfs,Inf),size(cfs,2)
-        if kend > 1 while isempty(chop(cfs[:,kend],ncfs*tol)) kend-=1 end end
-        ret=VFun{S,T}[Fun(columnspace(sp,k),chop(cfs[:,k],ncfs*tol)) for k=1:max(kend,1)]
+        ncfs, kend = norm(cfs,Inf), size(cfs,2)
+        if kend > 1
+            while kend > 0
+                if all(iszero, @view(cfs[:, kend]))
+                    kend -= 1
+                    continue
+                end
+                if !isempty(chop(@view(cfs[:,kend]), ncfs*tol))
+                    break
+                end
+                kend-=1
+            end
+        end
+        ret=VFun{S,T}[Fun(columnspace(sp,k),chop(@view(cfs[:,k]), ncfs*tol)) for k=1:max(kend,1)]
         ProductFun{S,V,typeof(sp),T}(ret,sp)
     else
         ret=VFun{S,T}[Fun(columnspace(sp,k),cfs[:,k]) for k=1:size(cfs,2)]
@@ -135,8 +146,8 @@ ProductFun(f::Function) = ProductFun(dynamic(f),ChebyshevInterval(),ChebyshevInt
 
 ## Conversion from other 2D Funs
 
-ProductFun(f::LowRankFun) = ProductFun(coefficients(f),space(f,1),space(f,2))
-ProductFun(f::Fun{S}) where {S<:AbstractProductSpace} = ProductFun(coefficientmatrix(f),space(f))
+ProductFun(f::LowRankFun; kw...) = ProductFun(coefficients(f),space(f,1),space(f,2); kw...)
+ProductFun(f::Fun{S}; kw...) where {S<:AbstractProductSpace} = ProductFun(coefficientmatrix(f), space(f); kw...)
 
 ## Conversion to other ProductSpaces with the same coefficients
 

src/PDE/KroneckerOperator.jl

@@ -426,21 +426,45 @@ end
 Base.getindex(A::KroneckerOperator, B::MultivariateFun) = A[Fun(B)]
 function Base.getindex(K::KroneckerOperator, f::LowRankFun)
     op1, op2 = K.ops
-    mapreduce(((A,B),)-> op1[A] ⊗ op2[B], +, zip(f.A, f.B))
-end
-function Base.getindex(K::KroneckerOperator, B::ProductFun)
-    op1, op2 = K.ops
-    S2 = factors(B.space)[2]
-    T = cfstype(B)
-    mapreduce(((ind, fi),)-> op1[fi] ⊗ op2[Fun(S2, [zeros(T, ind-1); one(T)])], +,
-        enumerate(B.coefficients))
+    sum(zip(f.A, f.B)) do (A,B)
+        op1[A] ⊗ op2[B]
+    end
 end
-for F in [:LowRankFun, :ProductFun, :MultivariateFun]
+Base.getindex(K::KroneckerOperator, B::ProductFun) = K[LowRankFun(B)]
+
+for F in [:MultivariateFun, :ProductFun, :LowRankFun]
     for O in [:DerivativeWrapper, :DefiniteIntegralWrapper]
         @eval Base.getindex(K::$O{<:KroneckerOperator}, f::$F) = K.op[f]
         @eval (*)(A::$O{<:KroneckerOperator}, B::$F) = A.op * B
-        @eval (*)(A::$F, B::$O{<:KroneckerOperator}) = Fun(A) * B.op
+        @eval (*)(A::$F, B::$O{<:KroneckerOperator}) = A * B.op
     end
-    @eval (*)(A::KroneckerOperator, B::$F) = A * Fun(B)
-    @eval (*)(A::$F, B::KroneckerOperator) = Fun(A) * B
 end
+(*)(A::KroneckerOperator, B::MultivariateFun) = A * Fun(B)
+(*)(A::MultivariateFun, B::KroneckerOperator) = Fun(A) * B
+
+(*)(ko::KroneckerOperator, pf::ProductFun) = ko * LowRankFun(pf)
+(*)(pf::ProductFun, ko::KroneckerOperator) = LowRankFun(pf) * ko
+
+# if the second operator is a constant, we may scale the first operator,
+# and apply it on the coefficients
+function (*)(ko::KroneckerOperator{<:Operator, <:ConstantOperator}, pf::ProductFun)
+    O1, O2 = ko.ops
+    O12 = O2.λ * O1
+    ProductFun(map(x -> O12*x, pf.coefficients), factors(pf.space)[2])
+end
+
+function (*)(ko::KroneckerOperator, lrf::LowRankFun)
+    O1, O2 = ko.ops
+    sum(zip(lrf.A, lrf.B)) do (f1, f2)
+        (O1*f1) ⊗ (O2*f2)
+    end
+end
+function (*)(lrf::LowRankFun, ko::KroneckerOperator)
+    O1, O2 = ko.ops
+    sum(zip(lrf.A, lrf.B)) do (f1, f2)
+        (f1*O1) ⊗ (f2*O2)
+    end
+end
+
+_mulop(P::Operator, ::BivariateSpace, B::ProductFun) = P * LowRankFun(B)
+(*)(P::PlusOperator, lrf::LowRankFun) = sum(op -> op*lrf, P.ops)

src/PDE/PDE.jl

@@ -49,13 +49,17 @@ function timedirichlet(d::Union{ProductDomain,TensorSpace})
 end
 
 
-
-function *(B::Operator,f::ProductFun)
+# Operators on an univariate space may act on the second space of the ProductFun,
+# consistent with treating it as an expansion in the second space
+# We re-route through _mulop to distinguish between operators on UnivariateSpace and
+# those on BivariateSpace
+function _mulop(B::Operator, ::UnivariateSpace, f::ProductFun)
     if isafunctional(B)
         Fun(factor(space(f),2),map(c->Number(B*c),f.coefficients))
     else
         ProductFun(space(f),map(c->B*c,f.coefficients))
     end
 end
+*(B::Operator,f::ProductFun) = _mulop(B, domainspace(B), f)
 
 *(f::ProductFun,B::Operator) = transpose(B*(transpose(f)))

src/Space.jl

@@ -552,9 +552,9 @@ transform!(S::Space,cfs) = plan_transform!(S,cfs)*cfs
 
 _coefficients!!(::Val{true}) = coefficients!
 _coefficients!!(::Val{false}) = coefficients
-_mul(P::CanonicalTransformPlan, ip, vals) = _coefficients!!(ip)(P.plan * vals, P.canonicalspace, P.space)
-_mul(P::ICanonicalTransformPlan, ip, cfs) = P.plan * _coefficients!!(ip)(cfs, P.space, P.canonicalspace)
-*(P::Union{CanonicalTransformPlan, ICanonicalTransformPlan}, vals::AbstractVector) = _mul(P, Val(inplace(P)), vals)
+_multransform(P::CanonicalTransformPlan, ip, vals) = _coefficients!!(ip)(P.plan * vals, P.canonicalspace, P.space)
+_multransform(P::ICanonicalTransformPlan, ip, cfs) = P.plan * _coefficients!!(ip)(cfs, P.space, P.canonicalspace)
+*(P::Union{CanonicalTransformPlan, ICanonicalTransformPlan}, vals::AbstractVector) = _multransform(P, Val(inplace(P)), vals)
 
 
 for OP in (:plan_transform,:plan_itransform,:plan_transform!,:plan_itransform!)