Support real / imag for expansions

Project.toml

@@ -1,6 +1,6 @@
 name = "ContinuumArrays"
 uuid = "7ae1f121-cc2c-504b-ac30-9b923412ae5c"
-version = "0.20.7"
+version = "0.20.8"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

src/ContinuumArrays.jl

@@ -3,7 +3,7 @@ using IntervalSets, DomainSets, LinearAlgebra, LazyArrays, FillArrays, BandedMat
 import Base: @_inline_meta, @_propagate_inbounds_meta, axes, size, getindex, convert, prod, *, /, \, +, -, ==, ^,
                 IndexStyle, IndexLinear, ==, OneTo, tail, similar, copyto!, copy, diff,
                 first, last, show, isempty, findfirst, findlast, findall, Slice, union, minimum, maximum, extrema, sum, _sum, _maximum, _minimum,
-                getproperty, isone, iszero, zero, abs, <, ≤, >, ≥, string, summary, to_indices, view, @propagate_inbounds, collect
+                getproperty, isone, iszero, zero, abs, <, ≤, >, ≥, string, summary, to_indices, view, @propagate_inbounds, collect, _all, isreal
 import Base.Broadcast: materialize, BroadcastStyle, broadcasted, Broadcasted
 import LazyArrays: MemoryLayout, Applied, ApplyStyle, flatten, _flatten, colsupport, combine_mul_styles, AbstractArrayApplyStyle,
                         adjointlayout, arguments, _mul_arguments, call, broadcastlayout, layout_getindex, UnknownLayout,

src/bases/bases.jl

@@ -867,6 +867,48 @@ function vec_layout(::ExpansionLayout, f)
 end
 
 
+######
+# real/zero
+#####
+
+isreal(::Basis{<:Real}) = true
+isreal(P::Basis) = error("Overload isreal(::$(typeof(P)))")
+_all(::typeof(iszero), ::Basis, ::Colon) = false # basis are linearly independent
+_all(::typeof(isreal), P::Basis, ::Colon) = isreal(P)
+
+for op in (:iszero, :isreal)
+    @eval _all(::typeof($op), V::SubQuasiArray{T,N,<:Basis}, ::Colon) where {T,N} = isempty(V) || $op(parent(V))
+end
+
+layout_broadcasted(::Tuple{AbstractBasisLayout}, ::typeof(real), P) = isreal(P) ? convert(AbstractQuasiMatrix{real(eltype(P))}, P) : error("Overload real for $(typeof(P))")
+layout_broadcasted(::Tuple{AbstractBasisLayout}, ::typeof(imag), P) = isreal(P) ? zero(real(P)) : error("Overload imag for $(typeof(P))")
+
+function layout_broadcasted(::Tuple{ExpansionLayout}, ::typeof(real), f)
+    P,c = basis(f),coefficients(f)
+    if isreal(P)
+        real(P) * real(c)
+    else
+        real(P)*real(c) - imag(P)*imag(c)
+    end
+end
+
+function layout_broadcasted(::Tuple{ExpansionLayout}, ::typeof(imag), f)
+    P,c = basis(f),coefficients(f)
+    if isreal(P)
+        real(P) * imag(c)
+    else
+        real(P)*imag(c) + imag(P)*real(c)
+    end
+end
+
+for op in (:real, :imag)
+    @eval begin
+        layout_broadcasted(::Tuple{SubBasisLayout}, ::typeof($op), P) = view($op(parent(P)), parentindices(P)...)
+        layout_broadcasted(::Tuple{MappedBasisLayouts}, ::typeof($op), P) = $op(demap(P))[basismap(P),:]
+    end
+end
+
+
 
 include("basisconcat.jl")
 include("basiskron.jl")

src/bases/basiskron.jl

@@ -26,4 +26,19 @@ sum_layout(::KronExpansionLayout, F, dims...) = sum_layout(ApplyLayout{typeof(*)
 
 diff_layout(::KronExpansionLayout, F, order...; dims...) = diff_layout(ApplyLayout{typeof(*)}(), F, order...; dims...)
 
-copy(L::Ldiv{Lay,<:KronExpansionLayout}) where Lay<:AbstractBasisLayout = copy(Ldiv{Lay,ApplyLayout{typeof(*)}}(L.A, L.B))
+copy(L::Ldiv{Lay,<:KronExpansionLayout}) where Lay<:AbstractBasisLayout = copy(Ldiv{Lay,ApplyLayout{typeof(*)}}(L.A, L.B))
+
+
+function plotgrid_layout(::KronExpansionLayout, P)
+    A,X,Bt = arguments(P)
+    plotgrid(A,size(X,1)), plotgrid(parent(Bt), size(X,2))
+end
+
+
+for op in (:real, :imag)
+    @eval function layout_broadcasted(::Tuple{KronExpansionLayout}, ::typeof($op), f)
+        A,X,Bt = arguments(f)
+        @assert isreal(A) && isreal(Bt)
+        real(A) * $op(X) * real(Bt)
+    end
+end

src/bases/splines.jl

@@ -26,6 +26,7 @@ end
 axes(B::Spline{o}) where o =
     (Inclusion(first(B.points)..last(B.points)), OneTo(length(B.points)+o-1))
 ==(A::Spline{o}, B::Spline{o}) where o = A.points == B.points
+isreal(::Spline) = true
 
 function getindex(B::LinearSpline{T}, x::Number, k::Int) where T
     @boundscheck (x ∈ axes(B,1) && 1 ≤ k ≤ size(B,2)) || throw(BoundsError(B, (x, k)))

src/plotting.jl

@@ -32,6 +32,7 @@ plotgrid_layout(::ExpansionLayout, P) = plotgrid_layout(ApplyLayout{typeof(*)}()
 
 plotvalues_size(::Tuple{InfiniteCardinal{1}}, g, x=plotgrid(g)) = g[x]
 plotvalues_size(::Tuple{InfiniteCardinal{1},Int}, g, x=plotgrid(g)) = g[x,:]
+plotvalues_size(::NTuple{2, InfiniteCardinal{1}}, g, x=plotgrid(g)) = g[x...]
 plotvalues_layout(lay, g, x...) = plotvalues_size(size(g), g, x...)
 # plotvalues_layout(::WeightedBasisLayouts, wP, n...) = plotvalues(unweighted(wP), n...)
 plotvalues_layout(::ExpansionLayout{MappedBasisLayout}, g, x...) = plotvalues(demap(g))
@@ -53,3 +54,4 @@ end
 
 _split_svec(x) = (x,)
 _split_svec(x::AbstractArray{<:StaticVector{2}}) = (map(first,x), map(last,x))
+_split_svec(x::Tuple) = x

test/runtests.jl

@@ -83,6 +83,7 @@ end
 
 include("test_splines.jl")
 include("test_chebyshev.jl")
+include("test_fourier.jl")
 include("test_basisconcat.jl")
 include("test_basiskron.jl")
 

test/test_basiskron.jl

@@ -1,4 +1,5 @@
 using ContinuumArrays, QuasiArrays, StaticArrays, Test
+using ContinuumArrays: KronExpansionLayout
 
 @testset "Basis Kron" begin
     L = LinearSpline(range(0,1; length=4))
@@ -27,4 +28,18 @@ end
     @test diff(F;dims=2)[0.1,0.2]  ≈ L[0.1,:]'*C*diff(L)[0.2,:]
 
     @test L\F/L' == (L\F)/L' == L\(F/L') == C
+
+    @testset "real/imag" begin
+        F = L * (C .+ im) * L'
+        @test real(F)[0.1,0.2] == (L * C * L')[0.1,0.2]
+        @test imag(F)[0.1,0.2] == (L * one.(C) * L')[0.1,0.2]
+        @test MemoryLayout(real(F)) isa KronExpansionLayout 
+        @test MemoryLayout(imag(F)) isa KronExpansionLayout
+    end
+
+    @testset "plot" begin
+        F = L * C * L'
+        ((x,y), Z) = ContinuumArrays.plotgridvalues(F)
+        @test F[x,y] == Z
+    end
 end

test/test_chebyshev.jl

@@ -307,4 +307,10 @@ Base.:(==)(::FooBasis, ::FooBasis) = true
         @test_throws ErrorException expand(T,exp) + expand(P,cos)
         @test_throws ErrorException expand(P,cos) + expand(T,exp)
     end
+
+    @testset "isreal/iszero" begin
+        T = Chebyshev(5)
+        @test isreal(T)
+        @test !iszero(T)
+    end
 end

test/test_fourier.jl

@@ -0,0 +1,32 @@
+using ContinuumArrays, LazyArrays, Test
+import ContinuumArrays: Basis
+
+struct Fourier <: Basis{ComplexF64}
+    n::Int
+end
+
+struct CosBasis <: Basis{Float64}
+    n::Int
+end
+
+struct SinBasis <: Basis{Float64}
+    n::Int
+end
+
+Base.isreal(::Fourier) = false
+Base.real(F::Fourier) = CosBasis(F.n)
+Base.imag(F::Fourier) = [zeros(axes(F,1)) SinBasis(F.n-1)]
+
+Base.getindex(::Fourier, θ::Real, j::Int) = exp(im*(j-1)*θ)
+Base.getindex(::CosBasis, θ::Real, j::Int) = cos((j-1)*θ)
+Base.getindex(::SinBasis, θ::Real, j::Int) = sin(j*θ)
+
+Base.axes(F::Union{Fourier,CosBasis,SinBasis}) = (Inclusion(0..2π), Base.OneTo(F.n))
+@testset "Fourier" begin
+    @testset "reim" begin
+        F = Fourier(5)
+        f = F * (collect(1:5) .+ im)
+        @test real(f)[0.1] ≈ real(f[0.1])
+        @test imag(f)[0.1] ≈ imag(f[0.1])
+    end
+end

test/test_recipesbaseext.jl

@@ -41,4 +41,10 @@ import RecipesBase
         y = affine(0..1, x)
         @test plotgrid(T[y,:]) == (plotgrid(T) .+ 1)/2
     end
+
+    @testset "basiskron" begin
+        F = L*randn(6,6)*L'     
+        rep = RecipesBase.apply_recipe(Dict{Symbol, Any}(), F)
+        @test rep[1].args == (0:5, 0:5, F.args[2])
+    end
 end

test/test_splines.jl

@@ -714,4 +714,23 @@ Random.seed!(24543)
         @test L' / L' ≡ Eye(10)
         @test simplifiable(Rdiv(L', L')) ≡ Val(true)
     end
+
+    @testset "real" begin
+        L = LinearSpline(range(0,1,10))
+        L̃ = LinearSpline{ComplexF64}(range(0,1,10))
+        M = L[affine(0..2, 0..1),:]
+        M̃ = L̃[affine(0..2, 0..1),:]
+        V = L[:,1:10]
+        Ṽ = L̃[:,1:10]
+
+        for P in (L, M, L̃, M̃, V, Ṽ)
+            @test isreal(P) && all(isreal, P)
+            @test real(P) == convert(AbstractQuasiArray{real(eltype(P))}, P)
+            @test imag(P) ≡ zeros(axes(P)...)
+
+            f = P * (collect(1:10) .+ im)
+            @test real(f) == P * collect(1:10)
+            @test imag(f) == P * ones(10)
+        end
+    end
 end