diff --git a/Project.toml b/Project.toml index bdb4d81..1e55d15 100644 --- a/Project.toml +++ b/Project.toml @@ -6,9 +6,10 @@ authors = ["Gabriel Wu and contributors"] [deps] Arblib = "fb37089c-8514-4489-9461-98f9c8763369" DoubleFloats = "497a8b3b-efae-58df-a0af-a86822472b78" -FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838" FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341" +FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +GenericFFT = "a8297547-1b15-4a5a-a998-a2ac5f1cef28" GenericLinearAlgebra = "14197337-ba66-59df-a3e3-ca00e7dcff7a" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881" @@ -22,9 +23,10 @@ Wigxjpf = "af901252-fd8a-4391-8647-10b4fde07a1e" [compat] Arblib = "1" DoubleFloats = "1" -FastGaussQuadrature = "1" FFTW = "1" +FastGaussQuadrature = "1" ForwardDiff = "1" +GenericFFT = "0.1" GenericLinearAlgebra = "0.4" OffsetArrays = "1" Quadmath = "1" diff --git a/src/IITM/arbitrary.jl b/src/IITM/arbitrary.jl index 18fe5d3..13ba014 100644 --- a/src/IITM/arbitrary.jl +++ b/src/IITM/arbitrary.jl @@ -39,7 +39,7 @@ function transition_matrix_iitm(s::AbstractShape{T, CT}, λ, nₘₐₓ, Nr, Nϑ xφ = range(0, 2 * T(π), length = Nφ + 1)[1:(end - 1)] wφ = 2 * T(π) / Nφ - fourier_workspace = CT <: ComplexF64 ? _azimuthal_fourier_workspace(Nφ, Nϑ) : nothing + fourier_workspace = CT <: ComplexF64 ? _azimuthal_fourier_workspace(Nφ, Nϑ, nₘₐₓ) : nothing fourier_modes = CT <: ComplexF64 ? _azimuthal_fourier_mode_bins(nₘₐₓ) : nothing it = OrderDegreeIterator(nₘₐₓ) diff --git a/src/IITM/fourier.jl b/src/IITM/fourier.jl index a7c13c5..a8667f3 100644 --- a/src/IITM/fourier.jl +++ b/src/IITM/fourier.jl @@ -5,18 +5,38 @@ _iitm_ldiv(𝐌::AbstractMatrix{<:Union{Arb, Acb}}, 𝐗) = inv(𝐌) * 𝐗 _iitm_ldiv(𝐌, 𝐗) = 𝐌 \ 𝐗 -struct _AzimuthalFourierWorkspace{P} - contrast::Matrix{ComplexF64} - contrast_inv::Matrix{ComplexF64} - spectrum::Matrix{ComplexF64} - spectrum_inv::Matrix{ComplexF64} +# Capability predicate: true when the FFT path is available for this complex type. +_iitm_fft_capable(::Type{<:Complex{<:AbstractFloat}}) = true # FFTW / GenericFFT +_iitm_fft_capable(::Type{Acb}) = true # Arblib.dft! +_iitm_fft_capable(::Type) = false # direct fallback + +struct _AzimuthalFourierWorkspace{CT, P} + contrast::Matrix{CT} + contrast_inv::Matrix{CT} + spectrum::Matrix{CT} + spectrum_inv::Matrix{CT} plan::P + coeff::OffsetArray{CT, 2, Matrix{CT}} + coeff_inv::OffsetArray{CT, 2, Matrix{CT}} +end + +# Acb workspace: holds Flint-native AcbVector/AcbMatrix scratch so that +# no Arb/Flint allocations happen inside the radial loop. +struct _AcbFourierWorkspace + col_in::Arblib.AcbVector + col_out::Arblib.AcbVector + coeff_mat::AcbMatrix + coeff_inv_mat::AcbMatrix + coeff::OffsetArray{Acb, 2, AcbMatrix} + coeff_inv::OffsetArray{Acb, 2, AcbMatrix} + unit::Acb # cached `1` at the working precision (reused each radial step) end +# ---------- FFTW plan cache (ComplexF64 only) ---------- const _AZIMUTHAL_FFT_PLAN_CACHE = Dict{Tuple{Int, Int}, Any}() const _AZIMUTHAL_FFT_PLAN_CACHE_LOCK = ReentrantLock() -function _azimuthal_fft_plan(Nφ, Nϑ) +function _azimuthal_fft_plan(::Type{ComplexF64}, Nφ, Nϑ) key = (Nφ, Nϑ) lock(_AZIMUTHAL_FFT_PLAN_CACHE_LOCK) try @@ -29,36 +49,84 @@ function _azimuthal_fft_plan(Nφ, Nϑ) end end -function _azimuthal_fourier_workspace(Nφ, Nϑ) - contrast = Matrix{ComplexF64}(undef, Nφ, Nϑ) - contrast_inv = similar(contrast) - spectrum = similar(contrast) - spectrum_inv = similar(contrast) - plan = _azimuthal_fft_plan(Nφ, Nϑ) - return _AzimuthalFourierWorkspace(contrast, contrast_inv, spectrum, spectrum_inv, - plan) +# GenericFFT path: fresh plan each time (bakes in precision-specific twiddles). +# Do NOT cache — a BigFloat plan would go stale if precision changes. +function _azimuthal_fft_plan(::Type{CT}, Nφ, Nϑ) where {CT <: Complex{<:AbstractFloat}} + return plan_fft(zeros(CT, Nφ, Nϑ), 1) # GenericFFT; no flags kwarg +end + +# ---------- Workspace constructors ---------- + +function _azimuthal_fourier_workspace(Nφ, Nϑ, nₘₐₓ; prec = 0) + _azimuthal_fourier_workspace(ComplexF64, Nφ, Nϑ, nₘₐₓ) +end + +function _azimuthal_fourier_workspace(::Type{CT}, Nφ, Nϑ, nₘₐₓ; prec = 0) where {CT <: Complex{<:AbstractFloat}} + contrast = zeros(CT, Nφ, Nϑ) + contrast_inv = zeros(CT, Nφ, Nϑ) + spectrum = zeros(CT, Nφ, Nϑ) + spectrum_inv = zeros(CT, Nφ, Nϑ) + plan = _azimuthal_fft_plan(CT, Nφ, Nϑ) + coeff = OffsetArray(zeros(CT, 4nₘₐₓ + 1, Nϑ), (-2nₘₐₓ):(2nₘₐₓ), 1:Nϑ) + coeff_inv = OffsetArray(zeros(CT, 4nₘₐₓ + 1, Nϑ), (-2nₘₐₓ):(2nₘₐₓ), 1:Nϑ) + return _AzimuthalFourierWorkspace{CT, typeof(plan)}( + contrast, contrast_inv, spectrum, spectrum_inv, plan, coeff, coeff_inv) +end + +# Acb workspace: preallocated AcbVector/AcbMatrix scratch; no Flint allocations +# inside the radial loop. +function _azimuthal_fourier_workspace(::Type{Acb}, Nφ, Nϑ, nₘₐₓ; prec) + col_in = Arblib.AcbVector(Nφ; prec = prec) + col_out = Arblib.AcbVector(Nφ; prec = prec) + coeff_mat = AcbMatrix(4nₘₐₓ + 1, Nϑ; prec = prec) + coeff_inv_mat = AcbMatrix(4nₘₐₓ + 1, Nϑ; prec = prec) + coeff = OffsetArray(coeff_mat, (-2nₘₐₓ):(2nₘₐₓ), 1:Nϑ) + coeff_inv = OffsetArray(coeff_inv_mat, (-2nₘₐₓ):(2nₘₐₓ), 1:Nϑ) + unit = Acb(1; prec = prec) + return _AcbFourierWorkspace(col_in, col_out, coeff_mat, coeff_inv_mat, coeff, coeff_inv, unit) end +# ---------- Mode-bin helpers ---------- + function _azimuthal_fourier_mode_bins(nₘₐₓ, period = 1) qs = [q for q in (-2nₘₐₓ):(2nₘₐₓ) if q % period == 0] bins = [q ÷ period for q in qs] return qs, bins end -function _azimuthal_fourier_coefficients(ε::AbstractMatrix{ComplexF64}, nₘₐₓ, wφ, - workspace::_AzimuthalFourierWorkspace, - mode_bins) +# ---------- Forward column DFT: ComplexF64 (FFTW) ---------- +# FFTW plans support mul!(out, plan, in) in-place via AbstractFFTs. +function _apply_forward_dft!(spectrum::Matrix{ComplexF64}, + contrast::Matrix{ComplexF64}, + plan) + mul!(spectrum, plan, contrast) +end + +# ---------- Forward column DFT: generic Complex{<:AbstractFloat} (GenericFFT) ---------- +# GenericFFT's DummyFFTPlan only supports `plan * A` (allocating), not mul!. +function _apply_forward_dft!(spectrum::Matrix{CT}, + contrast::Matrix{CT}, + plan) where {CT <: Complex{<:AbstractFloat}} + spectrum .= plan * contrast +end + +# ---------- Fourier coefficient extraction: generic Complex{<:AbstractFloat} ---------- + +function _azimuthal_fourier_coefficients(ε::AbstractMatrix{CT}, nₘₐₓ, wφ, + workspace::_AzimuthalFourierWorkspace{CT}, + mode_bins) where {CT <: Complex{<:AbstractFloat}} Nφ, Nϑ = size(ε) qs, bins = mode_bins @. workspace.contrast = ε - 1 @. workspace.contrast_inv = workspace.contrast / ε - mul!(workspace.spectrum, workspace.plan, workspace.contrast) - mul!(workspace.spectrum_inv, workspace.plan, workspace.contrast_inv) + _apply_forward_dft!(workspace.spectrum, workspace.contrast, workspace.plan) + _apply_forward_dft!(workspace.spectrum_inv, workspace.contrast_inv, workspace.plan) - coeff = OffsetArray(zeros(ComplexF64, 4nₘₐₓ + 1, Nϑ), (-2nₘₐₓ):(2nₘₐₓ), - 1:Nϑ) - coeff_inv = similar(coeff) + # Write into preallocated workspace arrays (no allocation). + # Alias safety: callers consume coeff/coeff_inv before the next radial step. + coeff = workspace.coeff + coeff_inv = workspace.coeff_inv for i in 1:Nϑ for iq in eachindex(qs) @@ -71,3 +139,50 @@ function _azimuthal_fourier_coefficients(ε::AbstractMatrix{ComplexF64}, nₘₐ return coeff, coeff_inv end + +# ---------- Fourier coefficient extraction: Acb (Arblib.dft! backend) ---------- +# Uses AcbMatrix (Arblib native) for the coeff arrays — Julia's Matrix{Acb} is +# unsafe due to GC interaction with the Flint memory model. +# Precision is inferred from the ε matrix elements at runtime. + +function _azimuthal_fourier_coefficients(ε::AbstractMatrix{Acb}, nₘₐₓ, wφ, + workspace::_AcbFourierWorkspace, + mode_bins) + Nφ, Nϑ = size(ε) + qs, bins = mode_bins + prec = Arblib.precision(ε[1, 1]) + + # Reuse preallocated AcbVector/AcbMatrix scratch + cached unit (no Flint allocations). + # Alias safety: callers consume coeff/coeff_inv before the next radial step. + col_in = workspace.col_in + col_out = workspace.col_out + coeff = workspace.coeff + coeff_inv = workspace.coeff_inv + one_prec = workspace.unit + + for i in 1:Nϑ + # contrast column: ε[j,i] - 1 + for j in 1:Nφ + col_in[j] = ε[j, i] - one_prec + end + Arblib.dft!(col_out, col_in; len = Nφ, prec = prec) + for iq in eachindex(qs) + q = qs[iq] + idx = mod(-bins[iq], Nφ) + 1 + coeff[q, i] = wφ * col_out[idx] + end + + # contrast_inv column: (ε[j,i] - 1) / ε[j,i] + for j in 1:Nφ + col_in[j] = (ε[j, i] - one_prec) / ε[j, i] + end + Arblib.dft!(col_out, col_in; len = Nφ, prec = prec) + for iq in eachindex(qs) + q = qs[iq] + idx = mod(-bins[iq], Nφ) + 1 + coeff_inv[q, i] = wφ * col_out[idx] + end + end + + return coeff, coeff_inv +end diff --git a/src/IITM/linearization.jl b/src/IITM/linearization.jl index 16eba9e..5f07342 100644 --- a/src/IITM/linearization.jl +++ b/src/IITM/linearization.jl @@ -862,7 +862,7 @@ function _iitm_arbitrary_fixed_geometry_linearization(input, variables) xφ = range(zero(RT), 2 * RT(π), length = Nφ + 1)[1:(end - 1)] wφ = 2 * RT(π) / Nφ - fourier_workspace = CT <: ComplexF64 ? _azimuthal_fourier_workspace(Nφ, Nϑ) : + fourier_workspace = CT <: ComplexF64 ? _azimuthal_fourier_workspace(Nφ, Nϑ, nₘₐₓ) : nothing fourier_modes = CT <: ComplexF64 ? _azimuthal_fourier_mode_bins(nₘₐₓ) : nothing @@ -999,7 +999,7 @@ function _iitm_nfold_fixed_geometry_linearization(input, variables) @. xφ = (xφ + 1) * π / N @. wφ = π / N * wφ end - fourier_workspace = CT <: ComplexF64 ? _azimuthal_fourier_workspace(Nφ, Nϑ) : + fourier_workspace = CT <: ComplexF64 ? _azimuthal_fourier_workspace(Nφ, Nϑ, nₘₐₓ) : nothing fourier_modes = CT <: ComplexF64 ? _azimuthal_fourier_mode_bins(nₘₐₓ, N) : nothing diff --git a/src/IITM/nfold.jl b/src/IITM/nfold.jl index 5bbca90..07d9aff 100644 --- a/src/IITM/nfold.jl +++ b/src/IITM/nfold.jl @@ -37,7 +37,7 @@ function transition_matrix_iitm(s::AbstractNFoldShape{N, T, CT}, λ, nₘₐₓ, ϑ = acos.(x) Nϑ = has_symmetric_plane(s) ? Nϑ ÷ 2 : Nϑ - if CT <: ComplexF64 + if _iitm_fft_capable(CT) xφ = range(0, 2 * T(π) / N, length = Nφ + 1)[1:(end - 1)] wφ = 2 * T(π) / (N * Nφ) else @@ -46,8 +46,10 @@ function transition_matrix_iitm(s::AbstractNFoldShape{N, T, CT}, λ, nₘₐₓ, @. wφ = π / N * wφ end - fourier_workspace = CT <: ComplexF64 ? _azimuthal_fourier_workspace(Nφ, Nϑ) : nothing - fourier_modes = CT <: ComplexF64 ? _azimuthal_fourier_mode_bins(nₘₐₓ, N) : nothing + fourier_workspace = _iitm_fft_capable(CT) ? + _azimuthal_fourier_workspace(CT, Nφ, Nϑ, nₘₐₓ; + prec = CT === Acb ? Arblib.precision(s.m) : 0) : nothing + fourier_modes = _iitm_fft_capable(CT) ? _azimuthal_fourier_mode_bins(nₘₐₓ, N) : nothing it = collect(OrderDegreeIterator(nₘₐₓ)) its = [collect(enumerate([x for x in it if (x[2] % N + N) % N == i])) @@ -132,7 +134,7 @@ function transition_matrix_iitm(s::AbstractNFoldShape{N, T, CT}, λ, nₘₐₓ, ε = [refractive_index(s, (r * sin(ϑ[i]) * cos(φ), r * sin(ϑ[i]) * sin(φ), r * x[i]))^2 for φ in xφ, i in 1:Nϑ] - fourier_coeffs = CT <: ComplexF64 ? + fourier_coeffs = _iitm_fft_capable(CT) ? _azimuthal_fourier_coefficients(ε, nₘₐₓ, wφ, fourier_workspace, fourier_modes) : nothing @@ -213,6 +215,58 @@ function transition_matrix_iitm(s::AbstractNFoldShape{N, T, CT}, λ, nₘₐₓ, return TransitionMatrix{CT, nₘₐₓ, typeof(𝐓′)}(𝐓′) end +@testitem "Generic FFT: Complex{Double64} Prism matches ComplexF64 reference" begin + using TransitionMatrices: Prism, calc_T_iitm, calc_Csca, calc_Cext, Double64 + + # Reference: ComplexF64 (FFTW path) + m_f64 = complex(1.5) + p_f64 = Prism(4, 1.0, 1.5, m_f64) + nₘₐₓ = 4 + Nr = 20 + Nϑ = 20 + Nφ = 16 + 𝐓_ref = calc_T_iitm(p_f64, 2π, nₘₐₓ, Nr, Nϑ, Nφ) + Csca_ref = calc_Csca(𝐓_ref) + Cext_ref = calc_Cext(𝐓_ref) + + # GenericFFT path: Complex{Double64} + m_d64 = complex(Double64(3) / 2) + p_d64 = Prism(4, Double64(1), Double64(3) / 2, m_d64) + 𝐓_d64 = calc_T_iitm(p_d64, 2 * Double64(π), nₘₐₓ, Nr, Nϑ, Nφ) + Csca_d64 = Float64(calc_Csca(𝐓_d64)) + Cext_d64 = Float64(calc_Cext(𝐓_d64)) + + @test abs(Csca_d64 - Csca_ref) < 1e-10 * abs(Csca_ref) + @test abs(Cext_d64 - Cext_ref) < 1e-10 * abs(Cext_ref) +end + +@testitem "Generic FFT: Acb Prism matches ComplexF64 reference" begin + using TransitionMatrices: Prism, calc_T_iitm, calc_Csca, calc_Cext, Acb, Arb + + # Reference: ComplexF64 (FFTW path) + m_f64 = complex(1.5) + p_f64 = Prism(4, 1.0, 1.5, m_f64) + nₘₐₓ = 4 + Nr = 20 + Nϑ = 20 + Nφ = 16 + 𝐓_ref = calc_T_iitm(p_f64, 2π, nₘₐₓ, Nr, Nϑ, Nφ) + Csca_ref = calc_Csca(𝐓_ref) + Cext_ref = calc_Cext(𝐓_ref) + + # Acb path: Arblib.dft! + prec = 128 + m_acb = Acb(3, 0; prec = prec) / 2 + p_acb = Prism(4, Arb(1; prec = prec), Arb(3; prec = prec) / 2, m_acb) + 𝐓_acb = calc_T_iitm(p_acb, 2 * Arb(π; prec = prec), nₘₐₓ, Nr, Nϑ, Nφ) + # calc_Csca/calc_Cext return Arb for Acb T-matrices + Csca_acb = Float64(calc_Csca(𝐓_acb)) + Cext_acb = Float64(calc_Cext(𝐓_acb)) + + @test abs(Csca_acb - Csca_ref) < 1e-8 * abs(Csca_ref) + @test abs(Cext_acb - Cext_ref) < 1e-8 * abs(Cext_ref) +end + @testitem "Spheroids can be solved by the nfold-shape solver" begin struct FourFoldSpheroid{T, CT} <: AbstractNFoldShape{4, T, CT} s::Spheroid{T, CT} diff --git a/src/TransitionMatrices.jl b/src/TransitionMatrices.jl index 11fc0c7..c03935b 100644 --- a/src/TransitionMatrices.jl +++ b/src/TransitionMatrices.jl @@ -5,6 +5,7 @@ using Arblib: ArbLike, AcbLike, ArbVectorLike, AcbVectorLike, ArbMatrixLike, Acb using DoubleFloats: Double64 using FastGaussQuadrature: FastGaussQuadrature import FFTW +using GenericFFT using ForwardDiff: ForwardDiff using GenericLinearAlgebra: Diagonal, GenericLinearAlgebra, cond, inv using LinearAlgebra: lu, mul!