Skip to content

Commit 1c938bc

Browse files
dlfivefiftydlfivefifty
committed
Update ichebyshevu
1 parent 3eb4672 commit 1c938bc

File tree

1 file changed

+64
-45
lines changed

1 file changed

+64
-45
lines changed

‎src/chebyshevtransform.jl

Lines changed: 64 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -161,8 +161,6 @@ end
161161
IChebyshevTransformPlan{T,kind,K}(F::FFTW.r2rFFTWPlan{T,K,inplace,N,R}) where {T,kind,K,inplace,N,R} =
162162
IChebyshevTransformPlan{T,kind,K,inplace,N,R}(F)
163163

164-
size(P::IChebyshevTransformPlan) = isdefined(P, :plan) ? size(P.plan) : (0,)
165-
length(P::IChebyshevTransformPlan) = isdefined(P, :plan) ? length(P.plan) : 0
166164

167165

168166
# second kind Chebyshev transforms share a plan with their inverse
@@ -242,9 +240,8 @@ function *(P::IChebyshevTransformPlan{T,1,K,true,N}, x::AbstractArray{T,N}) wher
242240
end
243241

244242
function mul!(y::AbstractArray{T,N}, P::IChebyshevTransformPlan{T,1,K,false,N}, x::AbstractArray{T,N}) where {T<:fftwNumber,K,N}
245-
n = length(x)
246-
length(y) == n || throw(DimensionMismatch("output must match dimension"))
247-
n == 0 && return y
243+
size(y) == size(x) || throw(DimensionMismatch("output must match dimension"))
244+
isempty(x) && return y
248245

249246
_icheb1_prescale!(P.plan.region, x) # Todo: don't mutate x
250247
_plan_mul!(y, P.plan, x)
@@ -408,7 +405,7 @@ end
408405
@inline function _chebu2_postscale!(_, x::AbstractVector{T}) where T
409406
n = length(x)
410407
c = one(T)/ (n+1)
411-
for k=1:n # sqrt(1-x_j^2) weight
408+
@inbounds for k=1:n # sqrt(1-x_j^2) weight
412409
x[k] /= sinpi(k*c)
413410
end
414411
x
@@ -447,87 +444,109 @@ chebyshevutransform(x, dims...; kws...) = plan_chebyshevutransform(x, dims...; k
447444

448445

449446
## Inverse transforms take ChebyshevU coefficients and produce values at ChebyshevU points of the first and second kinds
447+
const IUFIRSTKIND = FFTW.RODFT01
450448

451-
452-
struct IChebyshevUTransformPlan{T,kind,inplace,P} <: ChebyshevPlan{T}
453-
plan::FFTW.r2rFFTWPlan{T,P,true,1,UnitRange{Int}}
454-
IChebyshevUTransformPlan{T,kind,inplace,P}(plan) where {T,kind,inplace,P} = new{T,kind,inplace,P}(plan)
455-
IChebyshevUTransformPlan{T,kind,inplace,P}() where {T,kind,inplace,P} = new{T,kind,inplace,P}()
449+
struct IChebyshevUTransformPlan{T,kind,K,inplace,N,R} <: ChebyshevPlan{T}
450+
plan::FFTW.r2rFFTWPlan{T,K,inplace,N,R}
451+
IChebyshevUTransformPlan{T,kind,K,inplace,N,R}(plan) where {T,kind,K,inplace,N,R} = new{T,kind,K,inplace,N,R}(plan)
452+
IChebyshevUTransformPlan{T,kind,K,inplace,N,R}() where {T,kind,K,inplace,N,R} = new{T,kind,K,inplace,N,R}()
456453
end
457454

458-
IChebyshevUTransformPlan{T,kind,inplace}(F::FFTW.r2rFFTWPlan{T,P}) where {T,kind,inplace,P} =
459-
IChebyshevUTransformPlan{T,kind,inplace,P}(F)
460-
461-
IChebyshevUTransformPlan{T,kind,true}(F::IChebyshevUTransformPlan{T,kind,false,P}) where {T,kind,P} =
462-
IChebyshevUTransformPlan{T,kind,true,P}(F.plan)
455+
IChebyshevUTransformPlan{T,kind,K}(F::FFTW.r2rFFTWPlan{T,K,inplace,N,R}) where {T,kind,K,inplace,N,R} =
456+
IChebyshevUTransformPlan{T,kind,K,inplace,N,R}(F)
463457

464-
function plan_ichebyshevutransform!(x::AbstractVector{T}, ::Val{1}) where T<:fftwNumber
458+
function plan_ichebyshevutransform!(x::AbstractArray{T,N}, ::Val{1}, dims...; kws...) where {T<:fftwNumber,N}
465459
if isempty(x)
466-
IChebyshevUTransformPlan{T,1,true,(8,)}()
460+
IChebyshevUTransformPlan{T,1,kindtuple(IUFIRSTKIND,N,dims...),true,N,isempty(dims) ? UnitRange{Int} : typeof(dims)}()
467461
else
468-
IChebyshevUTransformPlan{T,1,true,(8,)}(FFTW.plan_r2r!(x, FFTW.RODFT01))
462+
IChebyshevUTransformPlan{T,1,kindtuple(IUFIRSTKIND,N,dims...)}(FFTW.plan_r2r!(x, IUFIRSTKIND, dims...; kws...))
469463
end
470464
end
471-
function plan_ichebyshevutransform!(x::AbstractVector{T}, ::Val{2}) where T<:fftwNumber
472-
length(x) ≤ 1 && throw(ArgumentError("Vector must contain at least 2 entries"))
473-
IChebyshevUTransformPlan{T,2,true,(7,)}(FFTW.plan_r2r!(x, FFTW.RODFT00))
465+
function plan_ichebyshevutransform!(x::AbstractArray{T,N}, ::Val{2}, dims...; kws...) where {T<:fftwNumber,N}
466+
any(≤(1),size(x)) && throw(ArgumentError("Array must contain at least 2 entries"))
467+
IChebyshevUTransformPlan{T,2,kindtuple(USECONDKIND,N,dims...)}(FFTW.plan_r2r!(x, USECONDKIND))
474468
end
475469

476-
function plan_ichebyshevutransform(x::AbstractVector{T}, ::Val{1}) where T<:fftwNumber
470+
function plan_ichebyshevutransform(x::AbstractArray{T,N}, ::Val{1}, dims...; kws...) where {T<:fftwNumber,N}
477471
if isempty(x)
478-
IChebyshevUTransformPlan{T,1,false,(8,)}()
472+
IChebyshevUTransformPlan{T,1,kindtuple(IUFIRSTKIND,N,dims...),false,N,isempty(dims) ? UnitRange{Int} : typeof(dims)}()
479473
else
480-
IChebyshevUTransformPlan{T,1,false,(8,)}(FFTW.plan_r2r!(x, FFTW.RODFT01))
474+
IChebyshevUTransformPlan{T,1,kindtuple(IUFIRSTKIND,N,dims...)}(FFTW.plan_r2r(x, IUFIRSTKIND, dims...; kws...))
481475
end
482476
end
483-
function plan_ichebyshevutransform(x::AbstractVector{T}, ::Val{2}) where T<:fftwNumber
484-
length(x) ≤ 1 && throw(ArgumentError("Vector must contain at least 2 entries"))
485-
IChebyshevUTransformPlan{T,2,false,(7,)}(FFTW.plan_r2r!(x, FFTW.RODFT00))
477+
function plan_ichebyshevutransform(x::AbstractArray{T,N}, ::Val{2}, dims...; kws...) where {T<:fftwNumber,N}
478+
any(≤(1),size(x)) && throw(ArgumentError("Array must contain at least 2 entries"))
479+
IChebyshevUTransformPlan{T,2,kindtuple(USECONDKIND,N,dims...)}(FFTW.plan_r2r(x, USECONDKIND))
486480
end
487481

488-
plan_ichebyshevutransform!(x::AbstractVector) = plan_ichebyshevutransform!(x, Val(1))
489-
plan_ichebyshevutransform(x::AbstractVector) = plan_ichebyshevutransform(x, Val(1))
490482

483+
plan_ichebyshevutransform!(x::AbstractArray, dims...; kws...) = plan_ichebyshevutransform!(x, Val(1), dims...; kws...)
484+
plan_ichebyshevutransform(x::AbstractArray, dims...; kws...) = plan_ichebyshevutransform(x, Val(1), dims...; kws...)
491485

492-
function *(P::IChebyshevUTransformPlan{T,1,true}, x::AbstractVector{T}) where T<:fftwNumber
486+
function _ichebyu1_postscale!(_, x::AbstractVector{T}) where T
493487
n = length(x)
494-
n ≤ 1 && return x
495-
496-
x = P.plan * x
497-
for k=1:n # sqrt(1-x_j^2) weight
488+
@inbounds for k=1:n # sqrt(1-x_j^2) weight
498489
x[k] /= 2sinpi(one(T)/(2n) + (k-one(T))/n)
499490
end
500491
x
501492
end
493+
function *(P::IChebyshevUTransformPlan{T,1,K,true}, x::AbstractVector{T}) where {T<:fftwNumber,K}
494+
n = length(x)
495+
n ≤ 1 && return x
502496

497+
x = P.plan * x
498+
_ichebyu1_postscale!(P.plan.region, x)
499+
end
503500

504-
505-
function *(P::IChebyshevUTransformPlan{T,2,true}, x::AbstractVector{T}) where T<:fftwNumber
501+
function mul!(y::AbstractVector{T}, P::IChebyshevUTransformPlan{T,1,K,false}, x::AbstractVector{T}) where {T<:fftwNumber,K}
506502
n = length(x)
503+
length(y) == n || throw(DimensionMismatch("output must match dimension"))
507504
n ≤ 1 && return x
508505

506+
_plan_mul!(y, P.plan, x)
507+
_ichebyu1_postscale!(P.plan.region, y)
508+
end
509+
510+
function _ichebu2_rescale!(_, x::AbstractVector{T}) where T
511+
n = length(x)
509512
c = one(T)/ (n+1)
510-
lmul!((n+1)/(2n+2*one(T)), x)
511-
x = P.plan * x
512513
for k=1:n # sqrt(1-x_j^2) weight
513514
x[k] /= sinpi(k*c)
514515
end
516+
ldiv!(2, x)
515517
x
516518
end
517519

518-
ichebyshevutransform!(x::AbstractVector{T}, kind=Val(1)) where {T<:fftwNumber} =
519-
plan_ichebyshevutransform!(x, kind)*x
520+
function *(P::IChebyshevUTransformPlan{T,2,K,true}, x::AbstractVector{T}) where {T<:fftwNumber,K}
521+
n = length(x)
522+
n ≤ 1 && return x
520523

521-
ichebyshevutransform(x, kind=Val(1)) = ichebyshevutransform!(Array(x), kind)
524+
x = P.plan * x
525+
_ichebu2_rescale!(P.plan.region, x)
526+
end
522527

523-
*(P::IChebyshevUTransformPlan{T,k,false},x::AbstractVector{T}) where {T,k} =
524-
IChebyshevUTransformPlan{T,k,true}(P)*Array(x)
528+
function mul!(y::AbstractVector{T}, P::IChebyshevUTransformPlan{T,2,K,false}, x::AbstractVector{T}) where {T<:fftwNumber,K}
529+
n = length(x)
530+
length(y) == n || throw(DimensionMismatch("output must match dimension"))
531+
n ≤ 1 && return x
532+
533+
_plan_mul!(y, P.plan, x)
534+
_ichebu2_rescale!(P.plan.region, y)
535+
end
536+
537+
ichebyshevutransform!(x::AbstractVector{T}, dims...; kwds...) where {T<:fftwNumber} =
538+
plan_ichebyshevutransform!(x, dims...; kwds...)*x
539+
540+
ichebyshevutransform(x, dims...; kwds...) = plan_ichebyshevutransform(x, dims...; kwds...)*x
541+
542+
*(P::IChebyshevUTransformPlan{T,k,K,false,N}, x::AbstractArray{T,N}) where {T,k,K,N} =
543+
mul!(similar(x), P, x)
525544

526545

527546
## Code generation for integer inputs
528547

529548
for func in (:chebyshevtransform,:ichebyshevtransform,:chebyshevutransform,:ichebyshevutransform)
530-
@eval $func(x::AbstractVector{T}, kind=Val(1)) where {T<:Integer} = $func(convert(AbstractVector{Float64},x), kind)
549+
@eval $func(x::AbstractVector{T}, dims...; kwds...) where {T<:Integer} = $func(convert(AbstractVector{Float64},x), dims...; kwds...)
531550
end
532551

533552

0 commit comments

Comments
 (0)