Skip to content

Commit 3518349

Browse files
dlfivefiftydlfivefifty
committed
Update chebyshevutransform
1 parent 709fe03 commit 3518349

File tree

2 files changed

+105
-60
lines changed

2 files changed

+105
-60
lines changed

‎src/chebyshevtransform.jl

Lines changed: 97 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ function plan_chebyshevtransform!(x::AbstractArray{T,N}, ::Val{1}, dims...; kws.
2929
end
3030
end
3131
function plan_chebyshevtransform!(x::AbstractArray{T,N}, ::Val{2}, dims...; kws...) where {T<:fftwNumber,N}
32-
length(x) ≤ 1 && throw(ArgumentError("Array must contain at least 2 entries"))
32+
any(≤(1),size(x)) && throw(ArgumentError("Array must contain at least 2 entries"))
3333
ChebyshevTransformPlan{T,2,kindtuple(SECONDKIND,N,dims...)}(FFTW.plan_r2r!(x, FFTW.REDFT00, dims...; kws...))
3434
end
3535

@@ -42,7 +42,7 @@ function plan_chebyshevtransform(x::AbstractArray{T,N}, ::Val{1}, dims...; kws..
4242
end
4343
end
4444
function plan_chebyshevtransform(x::AbstractArray{T,N}, ::Val{2}, dims...; kws...) where {T<:fftwNumber,N}
45-
length(x) ≤ 1 && throw(ArgumentError("Vector must contain at least 2 entries"))
45+
any(≤(1),size(x)) && throw(ArgumentError("Array must contain at least 2 entries"))
4646
ChebyshevTransformPlan{T,2,kindtuple(SECONDKIND,N,dims...)}(FFTW.plan_r2r(x, FFTW.REDFT00, dims...; kws...))
4747
end
4848

@@ -202,22 +202,22 @@ end
202202
plan_ichebyshevtransform!(x::AbstractArray, dims...; kws...) = plan_ichebyshevtransform!(x, Val(1), dims...; kws...)
203203
plan_ichebyshevtransform(x::AbstractArray, dims...; kws...) = plan_ichebyshevtransform(x, Val(1), dims...; kws...)
204204

205-
@inline _icheb1_prerescale!(_, x::AbstractVector) = (x[1] *= 2)
206-
@inline function _icheb1_prerescale!(d::Number, x::AbstractMatrix)
205+
@inline _icheb1_prescale!(_, x::AbstractVector) = (x[1] *= 2)
206+
@inline function _icheb1_prescale!(d::Number, x::AbstractMatrix)
207207
if isone(d)
208208
lmul!(2, view(x,1,:))
209209
else
210210
lmul!(2, view(x,:,1))
211211
end
212212
x
213213
end
214-
@inline function _icheb1_prerescale!(d::UnitRange, x::AbstractMatrix)
214+
@inline function _icheb1_prescale!(d::UnitRange, x::AbstractMatrix)
215215
lmul!(2, view(x,:,1))
216216
lmul!(2, view(x,1,:))
217217
x
218218
end
219-
@inline _icheb1_postrescale!(_, x::AbstractVector) = (x[1] /= 2)
220-
@inline function _icheb1_postrescale!(d::Number, x::AbstractMatrix)
219+
@inline _icheb1_postscale!(_, x::AbstractVector) = (x[1] /= 2)
220+
@inline function _icheb1_postscale!(d::Number, x::AbstractMatrix)
221221
if isone(d)
222222
ldiv!(2, view(x,1,:))
223223
else
@@ -226,7 +226,7 @@ end
226226
x
227227
end
228228

229-
@inline function _icheb1_postrescale!(d::UnitRange, x::AbstractMatrix)
229+
@inline function _icheb1_postscale!(d::UnitRange, x::AbstractMatrix)
230230
ldiv!(2, view(x,1,:))
231231
ldiv!(2, view(x,:,1))
232232
x
@@ -236,7 +236,7 @@ function *(P::IChebyshevTransformPlan{T,1,K,true,N}, x::AbstractArray{T,N}) wher
236236
n = length(x)
237237
n == 0 && return x
238238

239-
_icheb1_prerescale!(P.plan.region, x)
239+
_icheb1_prescale!(P.plan.region, x)
240240
x = ldiv!(2^length(P.plan.region), P.plan*x)
241241
x
242242
end
@@ -246,14 +246,14 @@ function mul!(y::AbstractArray{T,N}, P::IChebyshevTransformPlan{T,1,K,false,N},
246246
length(y) == n || throw(DimensionMismatch("output must match dimension"))
247247
n == 0 && return y
248248

249-
_icheb1_prerescale!(P.plan.region, x) # Todo: don't mutate x
249+
_icheb1_prescale!(P.plan.region, x) # Todo: don't mutate x
250250
_plan_mul!(y, P.plan, x)
251-
_icheb1_postrescale!(P.plan.region, x)
251+
_icheb1_postscale!(P.plan.region, x)
252252
ldiv!(2^length(P.plan.region), y)
253253
end
254254

255-
@inline _icheb2_prerescale!(_, x::AbstractVector) = (x[1] *= 2; x[end] *= 2)
256-
@inline function _icheb2_prerescale!(d::Number, x::AbstractMatrix)
255+
@inline _icheb2_prescale!(_, x::AbstractVector) = (x[1] *= 2; x[end] *= 2)
256+
@inline function _icheb2_prescale!(d::Number, x::AbstractMatrix)
257257
if isone(d)
258258
lmul!(2, @view(x[1,:]))
259259
lmul!(2, @view(x[end,:]))
@@ -263,7 +263,7 @@ end
263263
end
264264
x
265265
end
266-
@inline function _icheb2_prerescale!(d::UnitRange, x::AbstractMatrix)
266+
@inline function _icheb2_prescale!(d::UnitRange, x::AbstractMatrix)
267267
lmul!(2, @view(x[1,:]))
268268
lmul!(2, @view(x[end,:]))
269269
lmul!(2, @view(x[:,1]))
@@ -289,20 +289,20 @@ end
289289
x
290290
end
291291
@inline function _icheb2_rescale!(d::Number, y::AbstractArray{T}) where T
292-
_icheb2_prerescale!(d, y)
292+
_icheb2_prescale!(d, y)
293293
lmul!(convert(T, size(y,d) - 1)/2, y)
294294
y
295295
end
296296
@inline function _icheb2_rescale!(d::UnitRange, y::AbstractArray{T}) where T
297-
_icheb2_prerescale!(d, y)
297+
_icheb2_prescale!(d, y)
298298
lmul!(prod(convert.(T, size(y) .- 1)./2), y)
299299
y
300300
end
301301

302302
function *(P::IChebyshevTransformPlan{T,2,K,true,N}, x::AbstractArray{T,N}) where {T<:fftwNumber,K,N}
303303
n = length(x)
304304

305-
_icheb2_prerescale!(P.plan.region, x)
305+
_icheb2_prescale!(P.plan.region, x)
306306
x = inv(P)*x
307307
_icheb2_rescale!(P.plan.region, x)
308308
end
@@ -311,7 +311,7 @@ function mul!(y::AbstractArray{T,N}, P::IChebyshevTransformPlan{T,2,K,false,N},
311311
n = length(x)
312312
length(y) == n || throw(DimensionMismatch("output must match dimension"))
313313

314-
_icheb2_prerescale!(P.plan.region, x)
314+
_icheb2_prescale!(P.plan.region, x)
315315
_plan_mul!(y, inv(P), x)
316316
_icheb2_postrescale!(P.plan.region, x)
317317
_icheb2_rescale!(P.plan.region, y)
@@ -324,71 +324,117 @@ ichebyshevtransform(x, dims...; kwds...) = plan_ichebyshevtransform(x, dims...;
324324

325325
## Chebyshev U
326326

327-
struct ChebyshevUTransformPlan{T,kind,inplace,P} <: ChebyshevPlan{T}
328-
plan::FFTW.r2rFFTWPlan{T,P,true,1,UnitRange{Int}}
329-
ChebyshevUTransformPlan{T,kind,inplace,P}(plan) where {T,kind,inplace,P} = new{T,kind,inplace,P}(plan)
330-
ChebyshevUTransformPlan{T,kind,inplace,P}() where {T,kind,inplace,P} = new{T,kind,inplace,P}()
331-
end
332-
333-
ChebyshevUTransformPlan{T,kind,inplace}(plan::FFTW.r2rFFTWPlan{T,P}) where {T,kind,inplace,P} =
334-
ChebyshevUTransformPlan{T,kind,inplace,P}(plan)
327+
const UFIRSTKIND = 9
328+
const USECONDKIND = 7
335329

336-
ChebyshevUTransformPlan{T,kind,inplace}(plan::ChebyshevUTransformPlan{T,kind,inp,P}) where {T,kind,inplace,inp,P} =
337-
ChebyshevUTransformPlan{T,kind,inplace,P}(plan.plan)
330+
struct ChebyshevUTransformPlan{T,kind,K,inplace,N,R} <: ChebyshevPlan{T}
331+
plan::FFTW.r2rFFTWPlan{T,K,inplace,N,R}
332+
ChebyshevUTransformPlan{T,kind,K,inplace,N,R}(plan) where {T,kind,K,inplace,N,R} = new{T,kind,K,inplace,N,R}(plan)
333+
ChebyshevUTransformPlan{T,kind,K,inplace,N,R}() where {T,kind,K,inplace,N,R} = new{T,kind,K,inplace,N,R}()
334+
end
338335

336+
ChebyshevUTransformPlan{T,kind,K}(plan::FFTW.r2rFFTWPlan{T,K,inplace,N,R}) where {T,kind,K,inplace,N,R} =
337+
ChebyshevUTransformPlan{T,kind,K,inplace,N,R}(plan)
339338

340339

341-
function plan_chebyshevutransform!(x::AbstractVector{T}, ::Val{1}) where T<:fftwNumber
340+
function plan_chebyshevutransform!(x::AbstractArray{T,N}, ::Val{1}, dims...; kws...) where {T<:fftwNumber,N}
342341
if isempty(x)
343-
ChebyshevUTransformPlan{T,1,true,(9,)}()
342+
ChebyshevUTransformPlan{T,1,kindtuple(UFIRSTKIND,N,dims...),true,N,isempty(dims) ? UnitRange{Int} : typeof(dims)}()
344343
else
345-
ChebyshevUTransformPlan{T,1,true,(9,)}(FFTW.plan_r2r!(x, FFTW.RODFT10))
344+
ChebyshevUTransformPlan{T,1,kindtuple(UFIRSTKIND,N,dims...)}(FFTW.plan_r2r!(x, FFTW.RODFT10, dims...; kws...))
346345
end
347346
end
348-
function plan_chebyshevutransform!(x::AbstractVector{T}, ::Val{2}) where T<:fftwNumber
349-
length(x) ≤ 1 && throw(ArgumentError("Vector must contain at least 2 entries"))
350-
ChebyshevUTransformPlan{T,2,true,(7,)}(FFTW.plan_r2r!(x, FFTW.RODFT00))
347+
function plan_chebyshevutransform!(x::AbstractArray{T,N}, ::Val{2}, dims...; kws...) where {T<:fftwNumber,N}
348+
any(≤(1),size(x)) && throw(ArgumentError("Array must contain at least 2 entries"))
349+
ChebyshevUTransformPlan{T,2,kindtuple(USECONDKIND,N,dims...)}(FFTW.plan_r2r!(x, FFTW.RODFT00, dims...; kws...))
351350
end
352351

353-
function plan_chebyshevutransform(x::AbstractVector{T}, ::Val{1}) where T<:fftwNumber
352+
function plan_chebyshevutransform(x::AbstractArray{T,N}, ::Val{1}, dims...; kws...) where {T<:fftwNumber,N}
354353
if isempty(x)
355-
ChebyshevUTransformPlan{T,1,false,(9,)}()
354+
ChebyshevUTransformPlan{T,1,kindtuple(UFIRSTKIND,N,dims...),false,N,isempty(dims) ? UnitRange{Int} : typeof(dims)}()
356355
else
357-
ChebyshevUTransformPlan{T,1,false,(9,)}(FFTW.plan_r2r!(x, FFTW.RODFT10))
356+
ChebyshevUTransformPlan{T,1,kindtuple(UFIRSTKIND,N,dims...)}(FFTW.plan_r2r(x, FFTW.RODFT10, dims...; kws...))
358357
end
359358
end
360-
function plan_chebyshevutransform(x::AbstractVector{T}, ::Val{2}) where T<:fftwNumber
361-
length(x) ≤ 1 && throw(ArgumentError("Vector must contain at least 2 entries"))
362-
ChebyshevUTransformPlan{T,2,false,(7,)}(FFTW.plan_r2r!(x, FFTW.RODFT00))
359+
function plan_chebyshevutransform(x::AbstractArray{T,N}, ::Val{2}, dims...; kws...) where {T<:fftwNumber,N}
360+
any(≤(1),size(x)) && throw(ArgumentError("Array must contain at least 2 entries"))
361+
ChebyshevUTransformPlan{T,2,kindtuple(USECONDKIND,N,dims...)}(FFTW.plan_r2r(x, FFTW.RODFT00, dims...; kws...))
363362
end
364363

365-
plan_chebyshevutransform!(x::AbstractVector) = plan_chebyshevutransform!(x, Val(1))
366-
plan_chebyshevutransform(x::AbstractVector) = plan_chebyshevutransform(x, Val(1))
364+
plan_chebyshevutransform!(x::AbstractArray, dims...; kws...) = plan_chebyshevutransform!(x, Val(1), dims...; kws...)
365+
plan_chebyshevutransform(x::AbstractArray, dims...; kws...) = plan_chebyshevutransform(x, Val(1), dims...; kws...)
367366

368367

369-
function *(P::ChebyshevUTransformPlan{T,1,true},x::AbstractVector{T}) where T
368+
@inline function _chebu1_prescale!(_, x::AbstractVector{T}) where T
370369
n = length(x)
371-
n ≤ 1 && return x
372-
373370
for k=1:n # sqrt(1-x_j^2) weight
374371
x[k] *= sinpi(one(T)/(2n) + (k-one(T))/n)/n
375372
end
373+
x
374+
end
375+
376+
@inline function _chebu1_postscale!(_, x::AbstractVector{T}) where T
377+
n = length(x)
378+
for k=1:n # sqrt(1-x_j^2) weight
379+
x[k] /= sinpi(one(T)/(2n) + (k-one(T))/n)/n
380+
end
381+
x
382+
end
383+
384+
function *(P::ChebyshevUTransformPlan{T,1,K,true}, x::AbstractVector{T}) where {T,K}
385+
length(x) ≤ 1 && return x
386+
_chebu1_prescale!(P.plan.region, x)
376387
P.plan * x
377388
end
378389

379-
function *(P::ChebyshevUTransformPlan{T,2,true}, x::AbstractVector{T}) where T
390+
function mul!(y::AbstractVector{T}, P::ChebyshevUTransformPlan{T,1,K,false}, x::AbstractVector{T}) where {T,K}
380391
n = length(x)
381-
n ≤ 1 && return x
392+
length(x) ≤ 1 && return copyto!(y, x)
393+
_chebu1_prescale!(P.plan.region, x)
394+
_plan_mul!(y, P.plan, x)
395+
_chebu1_postscale!(P.plan.region, x)
396+
y
397+
end
382398

399+
@inline function _chebu2_prescale!(_, x::AbstractVector{T}) where T
400+
n = length(x)
383401
c = one(T)/ (n+1)
384402
for k=1:n # sqrt(1-x_j^2) weight
385403
x[k] *= sinpi(k*c)
386404
end
387-
lmul!(c, P.plan * x)
405+
x
406+
end
407+
408+
@inline function _chebu2_postscale!(_, x::AbstractVector{T}) where T
409+
n = length(x)
410+
c = one(T)/ (n+1)
411+
for k=1:n # sqrt(1-x_j^2) weight
412+
x[k] /= sinpi(k*c)
413+
end
414+
x
388415
end
389416

390-
chebyshevutransform!(x::AbstractVector{T}, kind=Val(1)) where {T<:fftwNumber} =
391-
plan_chebyshevutransform!(x, kind)*x
417+
function *(P::ChebyshevUTransformPlan{T,2,K,true}, x::AbstractVector{T}) where {T,K}
418+
n = length(x)
419+
n ≤ 1 && return x
420+
_chebu2_prescale!(P.plan.region, x)
421+
lmul!(one(T)/ (n+1), P.plan * x)
422+
end
423+
424+
function mul!(y::AbstractVector{T}, P::ChebyshevUTransformPlan{T,2,K,false}, x::AbstractVector{T}) where {T,K}
425+
n = length(x)
426+
n ≤ 1 && return copyto!(y, x)
427+
_chebu2_prescale!(P.plan.region, x)
428+
_plan_mul!(y, P.plan, x)
429+
_chebu2_postscale!(P.plan.region, x)
430+
lmul!(one(T)/ (n+1), y)
431+
end
432+
433+
*(P::ChebyshevUTransformPlan{T,kind,K,false,N}, x::AbstractArray{T,N}) where {T,kind,K,N} =
434+
mul!(similar(x), P, x)
435+
436+
chebyshevutransform!(x::AbstractVector{T}, dims...; kws...) where {T<:fftwNumber} =
437+
plan_chebyshevutransform!(x, dims...; kws...)*x
392438

393439

394440
"""
@@ -397,9 +443,8 @@ chebyshevutransform!(x::AbstractVector{T}, kind=Val(1)) where {T<:fftwNumber} =
397443
transforms from values on a Chebyshev grid of the first or second kind to Chebyshev
398444
coefficients of the 2nd kind (Chebyshev U expansion).
399445
"""
400-
chebyshevutransform(x, kind=Val(1)) = chebyshevutransform!(Array(x), kind)
446+
chebyshevutransform(x, dims...; kws...) = plan_chebyshevutransform(x, dims...; kws...)*x
401447

402-
*(P::ChebyshevUTransformPlan{T,k,false}, x::AbstractVector{T}) where {T,k} = ChebyshevUTransformPlan{T,k,true}(P)*Array(x)
403448

404449
## Inverse transforms take ChebyshevU coefficients and produce values at ChebyshevU points of the first and second kinds
405450

‎test/chebyshevtests.jl

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -113,12 +113,12 @@ using FastTransforms, Test
113113
f̃ = copy(f)
114114
f̄ = copy(f̌)
115115
P = @inferred(plan_chebyshevutransform(f))
116-
@test P*f == f̌
116+
@test P*f ≈ f̌
117117
@test f == f̃
118118
@test_throws ArgumentError P * T[1,2]
119119
P = @inferred(plan_chebyshevutransform!(f))
120-
@test P*f == f̌
121-
@test f == f̌
120+
@test P*f ≈ f̌
121+
@test f ≈ f̌
122122
@test_throws ArgumentError P * T[1,2]
123123
Pi = @inferred(plan_ichebyshevutransform(f̌))
124124
@test Pi*f̌ ≈ f̃
@@ -149,16 +149,16 @@ using FastTransforms, Test
149149
f̃ = copy(f)
150150
f̄ = copy(f̌)
151151
P = @inferred(plan_chebyshevutransform(f, Val(2)))
152-
@test @inferred(P*f) == f̌
153-
@test f == f̃
152+
@test @inferred(P*f) ≈ f̌
153+
@test f ≈ f̃
154154
@test_throws ArgumentError P * T[1,2]
155155
P = @inferred(plan_chebyshevutransform!(f, Val(2)))
156-
@test @inferred(P*f) == f̌
157-
@test f == f̌
156+
@test @inferred(P*f) ≈ f̌
157+
@test f ≈ f̌
158158
@test_throws ArgumentError P * T[1,2]
159159
Pi = @inferred(plan_ichebyshevutransform(f̌, Val(2)))
160160
@test @inferred(Pi*f̌) ≈ f̃
161-
@test f̌ == f̄
161+
@test f̌ ≈ f̄
162162
@test_throws ArgumentError Pi * T[1,2]
163163
Pi = @inferred(plan_ichebyshevutransform!(f̌, Val(2)))
164164
@test @inferred(Pi*f̌) ≈ f̃

0 commit comments

Comments
 (0)