1- include (" FourierTransforms.jl/FourierTransforms.jl" )
1+ module DFTKGenericLinearAlgebraExt
2+ using DFTK
3+ using DFTK: DummyInplace
4+ using LinearAlgebra
5+ using AbstractFFTs
6+ import AbstractFFTs: Plan, ScaledPlan,
7+ fft, ifft, bfft, fft!, ifft!, bfft!,
8+ plan_fft, plan_ifft, plan_bfft, plan_fft!, plan_ifft!, plan_bfft!,
9+ rfft, irfft, brfft, plan_rfft, plan_irfft, plan_brfft,
10+ fftshift, ifftshift,
11+ rfft_output_size, brfft_output_size,
12+ plan_inv, normalization
13+ import Base: show, summary, size, ndims, length, eltype, * , inv, \
14+ import LinearAlgebra: mul!
215
3- # This is needed to flag that the fft_generic.jl file has already been loaded
4- const GENERIC_FFT_LOADED = true
5-
6- if ! isdefined (Main, :GenericLinearAlgebra )
7- @warn (" Code paths for generic floating-point types activated in DFTK. Remember to " *
8- " add 'using GenericLinearAlgebra' to your user script. " *
9- " See https://docs.dftk.org/stable/examples/arbitrary_floattype/ for details." )
10- end
16+ include (" ctfft.jl" ) # Main file of FourierTransforms.jl
1117
1218# Utility functions to setup FFTs for DFTK. Most functions in here
1319# are needed to correct for the fact that FourierTransforms is not
1420# yet fully compliant with the AbstractFFTs interface and has still
1521# various bugs we work around.
1622
17- function next_working_fft_size (:: Any , size:: Integer )
23+ function DFTK . next_working_fft_size (:: Any , size:: Integer )
1824 # TODO FourierTransforms has a bug, which is triggered
1925 # only in some factorizations, see
2026 # https://github.com/JuliaComputing/FourierTransforms.jl/issues/10
2127 nextpow (2 , size) # We fall back to powers of two to be safe
2228end
23- default_primes (:: Any ) = (2 , )
29+ DFTK . default_primes (:: Any ) = (2 , )
2430
2531# Generic fallback function, Float32 and Float64 specialization in fft.jl
26- function build_fft_plans! (tmp:: AbstractArray{<:Complex} )
32+ function DFTK . build_fft_plans! (tmp:: AbstractArray{<:Complex} )
2733 # Note: FourierTransforms has no support for in-place FFTs at the moment
2834 # ... also it's extension to multi-dimensional arrays is broken and
2935 # the algo only works for some cases
@@ -40,14 +46,12 @@ function build_fft_plans!(tmp::AbstractArray{<:Complex})
4046 ipFFT, opFFT, ipBFFT, opBFFT
4147end
4248
43-
44-
4549struct GenericPlan{T}
4650 subplans
4751 factor:: T
4852end
4953
50- function generic_apply (p:: GenericPlan , X:: AbstractArray )
54+ function Base.: * (p:: GenericPlan , X:: AbstractArray )
5155 pl1, pl2, pl3 = p. subplans
5256 ret = similar (X)
5357 for i = 1 : size (X, 1 ), j = 1 : size (X, 2 )
6569LinearAlgebra. mul! (Y, p:: GenericPlan , X) = Y .= p * X
6670LinearAlgebra. ldiv! (Y, p:: GenericPlan , X) = Y .= p \ X
6771
68- import Base: * , \ , inv, length
6972length (p:: GenericPlan ) = prod (length, p. subplans)
70- * (p:: GenericPlan , X:: AbstractArray ) = generic_apply (p, X)
7173* (p:: GenericPlan{T} , fac:: Number ) where {T} = GenericPlan {T} (p. subplans, p. factor * T (fac))
7274* (fac:: Number , p:: GenericPlan{T} ) where {T} = p * fac
7375\ (p:: GenericPlan , X) = inv (p) * X
7476inv (p:: GenericPlan{T} ) where {T} = GenericPlan {T} (inv .(p. subplans), 1 / p. factor)
7577
7678function generic_plan_fft (data:: AbstractArray{T, 3} ) where {T}
77- GenericPlan {T} ([FourierTransforms . plan_fft (data[:, 1 , 1 ]),
78- FourierTransforms . plan_fft (data[1 , :, 1 ]),
79- FourierTransforms . plan_fft (data[1 , 1 , :])], T (1 ))
79+ GenericPlan {T} ([plan_fft (data[:, 1 , 1 ]),
80+ plan_fft (data[1 , :, 1 ]),
81+ plan_fft (data[1 , 1 , :])], T (1 ))
8082end
8183function generic_plan_bfft (data:: AbstractArray{T, 3} ) where {T}
82- GenericPlan {T} ([FourierTransforms. plan_bfft (data[:, 1 , 1 ]),
83- FourierTransforms. plan_bfft (data[1 , :, 1 ]),
84- FourierTransforms. plan_bfft (data[1 , 1 , :])], T (1 ))
84+ GenericPlan {T} ([plan_bfft (data[:, 1 , 1 ]),
85+ plan_bfft (data[1 , :, 1 ]),
86+ plan_bfft (data[1 , 1 , :])], T (1 ))
87+ end
88+
8589end
0 commit comments