Skip to content

IITM: extend n-fold azimuthal FFT acceleration to generic complex types #12

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lucifer1004 merged 1 commit into from
Jun 7, 2026
Merged

IITM: extend n-fold azimuthal FFT acceleration to generic complex types #12

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,10 @@ authors = ["Gabriel Wu <wuzihua@pku.edu.cn> 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"
Expand All @@ -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"
Expand Down
2 changes: 1 addition & 1 deletion src/IITM/arbitrary.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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ₘₐₓ)
Expand Down
159 changes: 137 additions & 22 deletions src/IITM/fourier.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand All @@ -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
4 changes: 2 additions & 2 deletions src/IITM/linearization.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
62 changes: 58 additions & 4 deletions src/IITM/nfold.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Comment on lines +49 to +52

@coderabbitai coderabbitai Bot Jun 7, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

🧩 Analysis chain

🌐 Web query:

Does Arblib.jl version 1.0.0 provide an Arblib.precision function that accepts an Acb argument and returns the working precision in bits?

💡 Result:

Arblib.jl version 1.0.0 does not provide a function named Arblib.precision that specifically accepts an Acb argument to return its working precision in bits [1][2][3]. While Arblib.jl uses Julia's precision function in some contexts—typically via Base.precision for types or by accessing a precision field within its Julia structs (like Arb and Acb)—it does not expose an exported Arblib.precision method for this purpose [1][4][3]. In Arblib.jl, Julia structs such as Arb and Acb contain a precision field that stores the precision used by that specific object [1][4][3]. You can typically access this field directly on an instance (e.g., my_acb.precision). Furthermore, Arblib follows a naming convention where many Arb C-library functions are wrapped and "Juliafied," often mapping to keyword arguments (like prec) rather than dedicated precision-querying functions [1][3]. For querying precision related to the underlying Arb/FLINT library, users may also rely on C-level functions like acb_bits if accessing low-level wrappers [5][6].

Citations:

Fix possible API mismatch in Arblib.precision(s.m) for Acb (src/IITM/nfold.jl:49-52)
When CT === Acb, the code calls Arblib.precision(s.m); Arblib.jl v1.0.0 sources indicate there isn’t an Arblib.precision(::Acb) API, so this can throw a MethodError. Use the correct Arblib precision accessor for Acb (e.g., the precision stored on the Acb object) instead.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/IITM/nfold.jl` around lines 49 - 52, The call Arblib.precision(s.m) can
raise a MethodError for Acb because Arblib does not define precision(::Acb);
update the conditional to obtain the precision directly from the Acb object (or
its field/accessor) when CT === Acb and pass that to
_azimuthal_fourier_workspace instead of Arblib.precision(s.m); locate the
creation of fourier_workspace in the _iitm_fft_capable(CT) branch and replace
the Arblib.precision(s.m) call with the correct accessor on the Acb value (or
use a safe helper that returns 0 for non-Acb types) so
_azimuthal_fourier_workspace receives a valid numeric precision.


it = collect(OrderDegreeIterator(nₘₐₓ))
its = [collect(enumerate([x for x in it if (x[2] % N + N) % N == i]))
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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}
Expand Down
1 change: 1 addition & 0 deletions src/TransitionMatrices.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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!
Expand Down
Loading