diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 00000000..5e9782f0 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,79 @@ +# Changelog + +## Unreleased (v1.2.18) + +### Breaking changes + +**`dist_synthesis!` / `dist_synthesis_sphtor!` reject `real_output=false` with real output arrays.** +`real_output=false` used to return the *real* field wrapped as complex (a real +buffer re-typed), so writing it into a `PencilArray{Float64}` happened to work. +It now performs a genuine complex synthesis — summing the m ≥ 0 half without the +Hermitian mirror — which is a different function, not the same field in a wider +type. On a typical config the complex-path result has `|imag|` up to 1.34 and a +real part differing from the real field by 1.31, against a field magnitude of +2.96, so no tolerance check can bridge the two. + +*Porting:* if you passed `real_output=false` with a real output array and wanted +the real field, pass `real_output=true`. If you want the true complex synthesis, +pass a complex output `PencilArray`. The error message states both. + +**`analysis_axisym` / `analysis_axisym_l` now return different values.** +Both omitted the φ quadrature factor `cfg.cphi * nlon = 2π`, so every returned +coefficient was `1/2π` too small — they inverted neither `synthesis_axisym` nor +the m=0 column of the full `analysis`. They now agree with both. Anything that +compensated for the old scale downstream must drop that compensation. + +### Fixed + +- **Silent precision loss in batch QST/sphtor transforms.** `analysis_qst_batch`, + `_synthesis_qst_batch` and the sphtor batch pair derived their output element + type from one input array instead of promoting across all of them, truncating + double-precision components to a single-precision sibling's type (measured + error 2.05e-8 instead of ~1e-17). +- **`load_config` silently downgraded `:driscoll_healy` grids to `:regular`**, + changing both the θ nodes and the quadrature weights; analysis→synthesis error + degraded from 8.7e-16 to ~3e-3. `save_config` now records `use_dh_weights`. +- **`shtns_set_grid` returned success for unrecognized grid codes** while + producing a south-to-north grid still reported as north-pole-first. The + fallback branch was missing the `reverse!` the real Gauss branch applies. +- **Pole-inclusive grids with `nlat == 1` produced an all-NaN config** (`π/0`) + that returned NaN from every subsequent transform with no error raised. Now + rejected with a message naming the cause. +- **`dist_SH_mul_mx!` crashed on every `mres > 1` config** — it walked all orders + through `LM_index`, which requires multiples of `mres`. +- **`dist_SH_Yrotate` crashed on `mres > 1`.** A Y-rotation mixes orders and so + cannot be represented in an `mres`-strided layout at all; it now says that + up front instead of failing deep inside the rotation. +- **`device_transfer_arrays(cfg, ...)` rejected `:cuda`/`:amdgpu`** — the config + vocabulary — because `to_device` accepted only `:cpu`/`:gpu`. +- **ForwardDiff could not flow through any plan-based batch transform.** + `SHTPlan` is FFTW-backed and cannot hold `ForwardDiff.Dual`; the batch entry + points now route non-FFTW element types through the plan-free `cfg`-form + transforms. +- **Cached FFT plans could silently fall back to an O(n²) DFT.** The plan cache + keys on shape and strides but not alignment, so reuse on a differently-aligned + matrix threw and callers quietly downgraded. Plans are now built `UNALIGNED`. +- Legendre south-pole and normalization-comment corrections carried over from the + orthonormal refactor; eleven comments prescribed conversions the code no longer + performs. + +### Performance + +- **Distributed analysis: 3 collectives per call → 1 on first use, 0 after.** + The `φ_is_local_all` / `θ_is_distributed` predicates are reduced together in a + single `Allreduce` and cached per `(pencil, communicator)`. +- **`dist_synthesis_packed_cplx` is single-pass**, down from two full distributed + syntheses. The negative-m φ bins are filled in the same θ/m traversal as the + positive ones, reusing one Legendre row per `(m, θ)`. It now matches the serial + reference exactly. +- **Legendre table memory halved.** `prepare_plm_tables!` was building + `NP_tables`/`NdP_tables` bit-for-bit identical to `plm_tables`/`dplm_tables`; + they now alias. `estimate_table_memory` previously reported half the true + figure, so jobs sized by it allocated twice their budget. +- Batch FFT helpers reuse the shared plan cache instead of re-planning per call. + +### Internal + +- `pack_lm!`/`pack_lm`/`unpack_lm!`/`unpack_lm` in `src/layout.jl` replace six + open-coded copies of the packed↔dense `(l,m)` mapping. The `m % mres` guard had + to be fixed three separate times across those copies. diff --git a/ext/ParallelLocal.jl b/ext/ParallelLocal.jl index 5612fd98..076f3cc5 100644 --- a/ext/ParallelLocal.jl +++ b/ext/ParallelLocal.jl @@ -252,17 +252,9 @@ end """ function SHTnsKit.dist_analysis_packed(cfg::SHTnsKit.SHTConfig, fθφ::PencilArray) Alm = SHTnsKit.dist_analysis(cfg, fθφ) - # `LM_index` throws unless m is a multiple of mres, so stride like the serial - # twin `analysis_packed` (src/transforms.jl:120) instead of walking every m. - Qlm = zeros(ComplexF64, cfg.nlm) - for m in 0:cfg.mmax - (m % cfg.mres == 0) || continue - for l in m:cfg.lmax - lm = SHTnsKit.LM_index(cfg.lmax, cfg.mres, l, m) + 1 - Qlm[lm] = Alm[l+1, m+1] - end - end - return Qlm + # Shared with the serial twin `analysis_packed`; `pack_lm` carries the + # `m % mres == 0` stride that `LM_index` requires. + return SHTnsKit.pack_lm(cfg, Alm) end """ @@ -270,15 +262,9 @@ end """ function SHTnsKit.dist_synthesis_packed(cfg::SHTnsKit.SHTConfig, Qlm::AbstractVector{<:Complex}; prototype_θφ::PencilArray, real_output::Bool=true) length(Qlm) == cfg.nlm || throw(DimensionMismatch("Qlm length")) - # Same mres stride as `synthesis_packed` (src/transforms.jl:141); without it - # `LM_index` throws on the first order that is not a multiple of mres. - Alm = zeros(ComplexF64, cfg.lmax+1, cfg.mmax+1) - for m in 0:cfg.mmax - (m % cfg.mres == 0) || continue - for l in m:cfg.lmax - Alm[l+1, m+1] = Qlm[SHTnsKit.LM_index(cfg.lmax, cfg.mres, l, m) + 1] - end - end + # Shared with the serial twin `synthesis_packed`; `unpack_lm` carries the + # `m % mres == 0` stride that `LM_index` requires. + Alm = SHTnsKit.unpack_lm(cfg, Qlm) return SHTnsKit.dist_synthesis(cfg, Alm; prototype_θφ, real_output) end @@ -362,7 +348,11 @@ function SHTnsKit.dist_synthesis_packed_cplx(cfg::SHTnsKit.SHTConfig, alm_packed Aminus[l+1, m+1] = conj(alm_packed[SHTnsKit.LM_cplx_index(lmax, mmax, l, -m) + 1]) end end - zp = SHTnsKit.dist_synthesis(cfg, Aplus; prototype_θφ, real_output=false) - zn = SHTnsKit.dist_synthesis(cfg, Aminus; prototype_θφ, real_output=false) - return zp .+ conj.(zn) + # Single pass: `dist_synthesis` fills the −m φ-FFT bins from `Aminus` in the + # same θ/m traversal that fills the +m bins, reusing one Legendre row per + # (m, θ) — P̄_l^{|m|} depends only on |m|. This used to be two full distributed + # syntheses combined as `zp + conj(zn)`, which doubled the Legendre work, the + # inverse FFT and, on a φ-distributed pencil, the communication. Same shape as + # the serial twin `synthesis_packed_cplx` (src/complex_packed.jl:110-130). + return SHTnsKit.dist_synthesis(cfg, Aplus; prototype_θφ, real_output=false, Aminus) end diff --git a/ext/ParallelTransforms.jl b/ext/ParallelTransforms.jl index 75dc7bf1..0b73be00 100644 --- a/ext/ParallelTransforms.jl +++ b/ext/ParallelTransforms.jl @@ -62,6 +62,88 @@ DEBUGGING CHECKLIST ================================================================================ =# +# ===== PER-PENCIL TOPOLOGY PREDICATES ===== +# +# `φ_is_local_all` and `θ_is_distributed` decide which branch a transform takes, +# and the branches enter different collectives, so both must be REDUCED rather +# than evaluated per rank — see the long comments at each use site. But they are +# also fixed properties of the decomposition: for a given pencil and +# communicator the answer never changes. Recomputing them in the body of every +# transform turned the recommended θ-only pencil from one collective per +# `dist_analysis` call into three, which is a latency-bound slowdown and an extra +# serialization point at scale. `DistAnalysisPlan`/`DistSphtorPlan` already +# compute them once at construction (ParallelPlans.jl); this gives the +# non-planned path the same treatment. +# +# Two properties make this safe: +# +# 1. ONE collective, not two. `φ_is_local_all` is stored inverted so both +# predicates are OR-reductions and share a single `Allreduce` on a bitmask. +# Even a cache miss is cheaper than the code it replaces. +# +# 2. Rank-symmetric hit/miss. The cache is keyed on the pencil OBJECT and the +# communicator, so every rank misses on the first use of a given pencil and +# hits on every later use — the same pattern on all ranks, which is exactly +# the condition a collective needs. Do NOT re-key this on anything a rank +# could compute differently from its peers (local sizes, ownership, rank id): +# if one rank hits while another misses, the missing rank enters the +# `Allreduce` alone and the job hangs. +# +# `Pencil` is an immutable struct, so it cannot key a `WeakKeyDict` (that needs a +# finalizer, hence a mutable key). Entries are therefore strong and are bounded +# by `_PENCIL_TOPOLOGY_MAX` instead, so code that rebuilds a pencil per shell or +# timestep cannot grow this without limit. The eviction is a full `empty!` keyed +# on a size the ranks all reach at the same call — deliberately NOT an LRU or +# per-entry policy, so that the hit/miss pattern stays identical on every rank. +const _PENCIL_TOPOLOGY_CACHE = IdDict{Any,IdDict{Any,Tuple{Bool,Bool}}}() +const _PENCIL_TOPOLOGY_LOCK = ReentrantLock() +const _PENCIL_TOPOLOGY_MAX = 256 + +""" + _pencil_topology(key_arr, comm, nθ_local, nφ_local, nlat, nlon) -> (φ_is_local_all, θ_is_distributed) + +Reduced topology predicates for `key_arr`'s pencil over `comm`, computed once per +`(pencil, comm)` pair with a single `Allreduce` and cached thereafter. + +Every rank must call this at the same point in its program, exactly as it must +for the `Allreduce` this replaces. +""" +function _pencil_topology(key_arr, comm, nθ_local::Int, nφ_local::Int, nlat::Int, nlon::Int) + key = pencil(key_arr) + # Look up without holding the lock across the collective below. + lock(_PENCIL_TOPOLOGY_LOCK) + try + inner = get(_PENCIL_TOPOLOGY_CACHE, key, nothing) + if inner !== nothing + hit = get(inner, comm, nothing) + hit === nothing || return hit + end + finally + unlock(_PENCIL_TOPOLOGY_LOCK) + end + + # Both predicates as OR-reductions in one bitmask: bit 0 = "some rank does NOT + # own the full φ range", bit 1 = "some rank owns fewer than all latitudes". + flags = UInt8(0) + nφ_local == nlon || (flags |= 0x01) + nθ_local < nlat && (flags |= 0x02) + allflags = MPI.Allreduce(flags, |, comm) + val = ((allflags & 0x01) == 0, (allflags & 0x02) != 0) + + lock(_PENCIL_TOPOLOGY_LOCK) + try + # Bounded, and evicted wholesale so every rank drops the same entries at + # the same call — a partial eviction could leave one rank hitting while + # another misses, and the miss would enter the `Allreduce` alone. + length(_PENCIL_TOPOLOGY_CACHE) >= _PENCIL_TOPOLOGY_MAX && empty!(_PENCIL_TOPOLOGY_CACHE) + inner = get!(() -> IdDict{Any,Tuple{Bool,Bool}}(), _PENCIL_TOPOLOGY_CACHE, key) + inner[comm] = val + finally + unlock(_PENCIL_TOPOLOGY_LOCK) + end + return val +end + # ===== ENHANCED PACKED STORAGE SYSTEM ===== # Reduces memory usage by ~50% for large spectral arrays by storing only l≥m coefficients # This is optional - dense storage (full lmax×mmax matrix) is the default @@ -516,7 +598,9 @@ Decompose latitude instead: `SHTnsKit.create_spatial_pencil(cfg; comm)` or `Penc # fixed elsewhere in this file, and it hangs the same way. Reduced here, # above the branch, so every rank executes the collective unconditionally # (`use_rfft_effective` is itself uniform: a keyword plus a shared eltype). - φ_is_local_all = MPI.Allreduce(nlon_local == nlon, &, comm) + # `θ_is_distributed` is reduced in the SAME collective and reused below. + φ_is_local_all, θ_is_distributed = + _pencil_topology(fθφ, comm, nθ_local, nlon_local, nlat, nlon) if use_rfft_effective nbins = nlon ÷ 2 + 1 @@ -647,10 +731,10 @@ Decompose latitude instead: `SHTnsKit.create_spatial_pencil(cfg; comm)` or `Penc # IMPORTANT: Only reduce if θ is actually distributed! # - If θ is distributed: each rank has different θ points → need Allreduce # - If only φ is distributed: all ranks have same θ points after gather → skip reduction - # Reduced, not per-rank: `nθ_local < nlat` is not uniform when a pencil has + # `θ_is_distributed` was reduced above alongside `φ_is_local_all` — reduced, + # not per-rank, because `nθ_local < nlat` is not uniform when a pencil has # more θ partitions than rows (nlat=1 on ≥2 θ-ranks), and the lone owner would # then skip the block while the empty ranks enter the collective alone and hang. - θ_is_distributed = MPI.Allreduce(nθ_local < nlat, |, comm) if θ_is_distributed # φ-partners that hold the same θ-slab carry identical post-gather @@ -792,10 +876,26 @@ function dist_analysis_with_scratch_buffers(plan::DistAnalysisPlan, fθφ::Penci end -function SHTnsKit.dist_synthesis(cfg::SHTnsKit.SHTConfig, Alm::AbstractMatrix; prototype_θφ::PencilArray, real_output::Bool=true, use_rfft::Bool=false) +""" +Optional `Aminus` (internal): coefficients for the NEGATIVE-m half of a genuinely +complex field, in the same `(lmax+1, mmax+1)` layout as `Alm`, with column `m+1` +holding `conj(a_{l,-m})` and column 1 unused. + +When given, the negative-m φ-FFT bins are filled from `Aminus` in the SAME θ/m +traversal that fills the positive bins, reusing the one Legendre row per (m, θ). +`dist_synthesis_packed_cplx` used to get this by calling `dist_synthesis` twice +and adding `zp + conj(zn)`, which doubled the Legendre work, the inverse FFT and +(on a φ-distributed pencil) the communication. Requires `real_output=false` and +the complex bin layout — an rfft buffer has no negative-m slots. +""" +function SHTnsKit.dist_synthesis(cfg::SHTnsKit.SHTConfig, Alm::AbstractMatrix; prototype_θφ::PencilArray, real_output::Bool=true, use_rfft::Bool=false, Aminus::Union{Nothing,AbstractMatrix}=nothing) lmax, mmax = cfg.lmax, cfg.mmax nlon = cfg.nlon nlat = cfg.nlat + if Aminus !== nothing + real_output && throw(ArgumentError("dist_synthesis with Aminus requires real_output=false")) + size(Aminus) == size(Alm) || throw(DimensionMismatch("Aminus must match Alm's shape")) + end # Contract: Alm must be replicated identically on every rank. Sample-hash a # bounded prefix + a small tail slice instead of the full matrix so the check @@ -806,7 +906,12 @@ function SHTnsKit.dist_synthesis(cfg::SHTnsKit.SHTConfig, Alm::AbstractMatrix; p k = min(n, 128) probe_head = n == 0 ? UInt64(0) : hash(view(Alm, 1:k)) probe_tail = n <= k ? UInt64(0) : hash(view(Alm, (n - k + 1):n)) - local_sig = hash((size(Alm, 1), size(Alm, 2), probe_head, probe_tail)) + # `Aminus` must be replicated too — it feeds the same collective-free + # local traversal, so a rank-varying copy would silently produce a + # different field per rank. + probe_minus = Aminus === nothing ? UInt64(0) : + hash((hash(view(Aminus, 1:k)), n <= k ? UInt64(0) : hash(view(Aminus, (n - k + 1):n)))) + local_sig = hash((size(Alm, 1), size(Alm, 2), probe_head, probe_tail, probe_minus)) rank0_sig = MPI.bcast(local_sig, 0, comm) if local_sig != rank0_sig throw(ArgumentError("dist_synthesis requires Alm replicated across ranks (signature mismatch on rank $(MPI.Comm_rank(comm))).")) @@ -843,6 +948,11 @@ function SHTnsKit.dist_synthesis(cfg::SHTnsKit.SHTConfig, Alm::AbstractMatrix; p # Compute synthesized values for each local θ for (ii, iglob) in enumerate(θ_globals) # Get normalized Legendre polynomials (P̄ = Nlm·P) at this latitude + # `gm` accumulates the negative-m half from `Aminus` on the SAME + # Legendre row — P̄_l^{|m|} depends only on |m|, so the -m bin costs + # one extra multiply-add per l instead of a whole second traversal. + gm = 0.0 + 0.0im + want_minus = Aminus !== nothing && mval > 0 if cfg.use_plm_tables && !isempty(cfg.NP_tables) # NP_tables[col][l+1, iglob] = P̄_l^m already; no extra Nlm multiply tbl = cfg.NP_tables[col] @@ -850,22 +960,37 @@ function SHTnsKit.dist_synthesis(cfg::SHTnsKit.SHTConfig, Alm::AbstractMatrix; p @inbounds @simd for l in mval:lmax g += tbl[l+1, iglob] * Alm[l+1, col] end + if want_minus + @inbounds @simd for l in mval:lmax + gm += tbl[l+1, iglob] * Aminus[l+1, col] + end + end else SHTnsKit.Plm_norm_row!(P, xv[iglob], lmax, mval) g = 0.0 + 0.0im @inbounds @simd for l in mval:lmax g += P[l+1] * Alm[l+1, col] end + if want_minus + @inbounds @simd for l in mval:lmax + gm += P[l+1] * Aminus[l+1, col] + end + end end # Store in Fourier coefficient array Fθm[ii, mval + 1] = inv_scaleφ * g - # For real output (complex path only), mirror onto negative-m bin. + # Negative-m bin. Two distinct sources: + # - real output: Hermitian mirror of the +m bin. + # - complex output with `Aminus`: the independent -m coefficients, + # conjugated, matching the `zp + conj(zn)` the two-pass form built. # rfft buffer has no slot for negative m — irfft reconstructs implicitly. if real_output && !use_rfft_effective && mval > 0 conj_index = nlon - mval + 1 Fθm[ii, conj_index] = conj(Fθm[ii, mval + 1]) + elseif want_minus + Fθm[ii, nlon - mval + 1] = conj(inv_scaleφ * gm) end end end @@ -921,14 +1046,24 @@ function SHTnsKit.dist_synthesis(cfg::SHTnsKit.SHTConfig, Alm::PencilArray; prot end function SHTnsKit.dist_synthesis!(plan::DistPlan, fθφ_out::PencilArray, Alm::PencilArray; real_output::Bool=true) - # `dist_synthesis` now returns a genuinely complex field for real_output=false - # (it used to hand back a real buffer re-wrapped as complex), so a real output - # array can no longer absorb it. Reject up front, before any collective, as - # the sphtor twin does — the check is on local eltypes, identical on every - # rank, so all ranks throw together instead of deadlocking. + # Rejected up front, before any collective — the test is on local eltypes, + # identical on every rank, so all ranks throw together instead of deadlocking. + # + # This combination cannot be silently accepted, and no "is the imaginary part + # negligible?" check can rescue it. `real_output=false` no longer means "the + # real field, typed complex" (which is what the old code returned, by wrapping + # a real buffer); it now sums only the m ≥ 0 half WITHOUT the Hermitian mirror, + # which is a genuinely different function. Measured on a typical config: the + # complex-path result has |imag| up to 1.34 and its REAL part differs from the + # real field by 1.31, against a field magnitude of 2.96. So a caller who used + # to pass `real_output=false` with a real output array wanted the real field, + # and today that is spelled `real_output=true` — hence the message below. if !real_output && eltype(fθφ_out) <: Real throw(ArgumentError("dist_synthesis! with real_output=false needs a complex output " * - "PencilArray; got eltype=$(eltype(fθφ_out))")) + "PencilArray; got eltype=$(eltype(fθφ_out)). If you are porting code " * + "that used this combination before v1.2.18: it used to return the REAL " * + "field wrapped as complex, so pass real_output=true to keep that result. " * + "Pass a complex output PencilArray to get the true complex synthesis.")) end f = SHTnsKit.dist_synthesis(plan.cfg, Alm; prototype_θφ=plan.prototype_θφ, real_output, use_rfft=plan.use_rfft) copyto!(fθφ_out, f) @@ -960,8 +1095,10 @@ function SHTnsKit.dist_analysis_sphtor(cfg::SHTnsKit.SHTConfig, Vtθφ::PencilAr end # φ-locality must be agreed by ALL ranks (see `dist_analysis_standard`): a # per-rank test lets the sole owner of a short φ dimension take the local - # branch while empty ranks enter the collective alone. - φ_is_local_all = MPI.Allreduce(nlon_local == nlon, &, comm) + # branch while empty ranks enter the collective alone. `θ_is_distributed` is + # reduced in the SAME collective and reused below. + φ_is_local_all, θ_is_distributed = + _pencil_topology(Vtθφ, comm, nθ_local, nlon_local, nlat, nlon) nbins = use_rfft_effective ? (nlon ÷ 2 + 1) : nlon Ftθm = Matrix{ComplexF64}(undef, nθ_local, nbins) @@ -1040,10 +1177,10 @@ function SHTnsKit.dist_analysis_sphtor(cfg::SHTnsKit.SHTConfig, Vtθφ::PencilAr # Only reduce if θ is actually distributed across processes # When φ is distributed but θ is not, all ranks compute identical results after gathering φ - # Reduced, not per-rank: `nθ_local < nlat` is not uniform when a pencil has + # `θ_is_distributed` was reduced above alongside `φ_is_local_all` — reduced, + # not per-rank, because `nθ_local < nlat` is not uniform when a pencil has # more θ partitions than rows (nlat=1 on ≥2 θ-ranks), and the lone owner would # then skip the block while the empty ranks enter the collective alone and hang. - θ_is_distributed = MPI.Allreduce(nθ_local < nlat, |, comm) if θ_is_distributed # Same dedup-then-reduce as the scalar dist_analysis_standard path (see @@ -1321,15 +1458,18 @@ end function SHTnsKit.dist_synthesis_sphtor!(plan::DistSphtorPlan, Vtθφ_out::PencilArray, Vpθφ_out::PencilArray, Slm::AbstractMatrix, Tlm::AbstractMatrix; real_output::Bool=true) - # A complex field cannot be written into real output arrays. Reject that - # combination up front, BEFORE any collective: previously the scratch path - # silently stored only the real part, and routing it to the allocating path - # instead turns the silent truncation into an `InexactError` from `copyto!` - # part-way through a collective region. The check is on local eltypes, which - # every rank agrees on, so all ranks throw together and nothing deadlocks. + # A complex field cannot be written into real output arrays. Rejected up front, + # BEFORE any collective; the test is on local eltypes, which every rank agrees + # on, so all ranks throw together and nothing deadlocks. See the twin in + # `dist_synthesis!` for why no imaginary-part tolerance can accept this + # instead: `real_output=false` computes a different function now, not the same + # field in a wider type, so the porting fix is `real_output=true`. if !real_output && (eltype(Vtθφ_out) <: Real || eltype(Vpθφ_out) <: Real) throw(ArgumentError("dist_synthesis_sphtor! with real_output=false needs complex output " * - "PencilArrays; got eltype(Vt)=$(eltype(Vtθφ_out)), eltype(Vp)=$(eltype(Vpθφ_out))")) + "PencilArrays; got eltype(Vt)=$(eltype(Vtθφ_out)), eltype(Vp)=$(eltype(Vpθφ_out)). " * + "If you are porting code that used this combination before v1.2.18: it used " * + "to return the REAL field wrapped as complex, so pass real_output=true to keep " * + "that result. Pass complex output PencilArrays to get the true complex synthesis.")) end # The scratch spatial buffers are `Matrix{Float64}` (see `_SphtorScratch`), so @@ -1906,8 +2046,10 @@ function dist_analysis_distributed(cfg::SHTnsKit.SHTConfig, fθφ::PencilArray; # reachable with the explicit MPITopology the empty-partition cases need) the # single owner sees `1 < 1 == false` and skips the block while the empty # ranks see `0 < 1 == true` and enter the full-comm Allreduce alone — which - # never completes. One extra small collective buys a matched one below. - θ_is_distributed = MPI.Allreduce(nθ_local < cfg.nlat, |, comm) + # never completes. Cached per (pencil, comm), so the collective is paid once + # for this decomposition rather than on every call. + _, θ_is_distributed = + _pencil_topology(fθφ, comm, nθ_local, size(parent(fθφ), 2), cfg.nlat, cfg.nlon) if θ_is_distributed # A 2D (θ×φ) spatial pencil must sum each θ-slab once, not once per # φ-partner: `_gather_and_fft_phi` hands every partner the FULL longitude @@ -2690,7 +2832,9 @@ function _dist_analysis_2d_safe(cfg::SHTnsKit.SHTConfig, fθφ::PencilArray; # Check if θ is distributed. Reduced, not per-rank: see the same guard in # `dist_analysis_distributed` — an unmatched full-comm Allreduce hangs. - θ_is_distributed = MPI.Allreduce(nθ_local < nlat, |, comm) + # Cached per (pencil, comm); the collective is paid once per decomposition. + _, θ_is_distributed = + _pencil_topology(fθφ, comm, nθ_local, size(parent(fθφ), 2), nlat, nlon) if θ_is_distributed # Same reasoning as `dist_analysis_distributed`: every φ-partner of a @@ -3112,8 +3256,11 @@ function _dist_analysis_2d_aligned(cfg::SHTnsKit.SHTConfig, fθφ::PencilArray; # Reduced, not per-rank: `nθ_local < nlat` is not uniform when a pencil has # more θ partitions than rows (nlat=1 on ≥2 θ-ranks), and the lone owner would # then skip the block while the empty ranks enter the collective alone and hang. Reduced over `l_comm`, - # which is the communicator the guarded Allreduce below actually uses. - θ_is_distributed = MPI.Allreduce(nθ_local < nlat, |, l_comm) + # which is the communicator the guarded Allreduce below actually uses — the + # cache is keyed on (pencil, comm) precisely so this l_comm answer never + # aliases the full-comm one for the same pencil. + _, θ_is_distributed = + _pencil_topology(fθφ, l_comm, nθ_local, size(parent(fθφ), 2), nlat, nlon) if θ_is_distributed # Use packed communication to avoid sending zeros in triangular region diff --git a/ext/ParallelTransposeTransforms.jl b/ext/ParallelTransposeTransforms.jl index d0168fb8..5a590091 100644 --- a/ext/ParallelTransposeTransforms.jl +++ b/ext/ParallelTransposeTransforms.jl @@ -205,11 +205,9 @@ function SHTnsKit.DistTransposePlan( NP[mi] = tbl_NP end - # 7. Per-local-m external↔internal scale. The Legendre tables above are - # orthonormal+CS, but this plan's API exchanges coefficients in cfg's - # convention (as `dist_analysis`/`dist_synthesis` do), so the transforms - # convert with this matrix. Without it a `norm=:schmidt` or - # `cs_phase=false` config silently disagreed with the cfg-form paths. + # No per-m normalization table is built or stored: the Legendre tables above + # are orthonormal+CS and so is this plan's API, matching `dist_analysis`/ + # `dist_synthesis` and every other transform. Nothing to convert. return DistTransposePlan( cfg, nlat, nlon, lmax, mmax, nlev, comm, @@ -308,9 +306,8 @@ function SHTnsKit.dist_synthesis!(plan::DistTransposePlan, f::PencilArray, Alm:: nlat = plan.nlat nlev = plan.nlev - # Incoming coefficients are in cfg's convention; the tables are orthonormal. - # Convert a copy so the caller's array is left untouched (no-op, and no - # allocation, on the default orthonormal+CS config). + # Incoming coefficients are orthonormal+CS, matching the tables — read them + # directly, no copy and no conversion. # Legendre expansion: for each local m, sum over l → F[i, mi, lev] # NP[mi] is (lmax+1, nlat) column-major; iterating i (fast dim of F) is cache-friendly. @@ -429,8 +426,8 @@ function SHTnsKit.dist_synthesis_sphtor!(plan::DistTransposePlan, nlat = plan.nlat nlev = plan.nlev - # Incoming coefficients are in cfg's convention; convert copies (no-op on the - # default orthonormal+CS config) so the caller's arrays are left untouched. + # Incoming coefficients are orthonormal+CS, matching the tables — read them + # directly, no copies and no conversion. # Legendre expansion: for each local m, sum over l → Ft[i,mi,lev], Fp[i,mi,lev] # Kernel (from kernels.jl _sphtor_synthesis_kernel_otf): diff --git a/ext/SHTnsKitAdvancedADExt.jl b/ext/SHTnsKitAdvancedADExt.jl index 2a83d9e8..568930a1 100644 --- a/ext/SHTnsKitAdvancedADExt.jl +++ b/ext/SHTnsKitAdvancedADExt.jl @@ -20,30 +20,30 @@ import SHTnsKit: wigner_d_matrix_deriv # ---- normalization in the adjoint ------------------------------------- # - # The `_adjoint_*` helpers work entirely in the INTERNAL (orthonormal + CS) - # convention. Some primals do not: the sphtor pair converts on the way in - # (`synthesis_sphtor`, src/sphtor_transforms.jl:178) and on the way out - # (`analysis_sphtor`, :275). That conversion is a real diagonal scale `M`, - # so it has to appear in the adjoint too: + # There is NO normalization factor in any adjoint here, and adding one would + # be a bug. Every transform in the package — scalar, sphtor, QST, plan, + # distributed, GPU — now emits and consumes coefficients in the single + # INTERNAL (orthonormal + CS) convention, which is exactly the convention the + # `_adjoint_*` helpers work in. Primal and adjoint therefore agree with no + # scaling on either side, for every `cfg.norm`/`cs_phase`. # - # synthesis-like y = F(M ⊙ a) ⇒ ā = M ⊙ Fᴴ(ȳ) - # analysis-like a = F(x) ⊘ M ⇒ x̄ = Fᴴ(ā ⊘ M) + # This block used to say the opposite, because the sphtor pair once converted + # on the way in and out with a real diagonal scale `M`, which forced a + # matching `M`/`1/M` into the pullbacks. Those conversions are gone. If you + # reintroduce an `M ⊙ ȳ` here to "match the primal", every non-default-norm + # gradient becomes wrong by M[l,m] (finite differences: 40–180% relative + # error on :schmidt and :fourpi) and nothing in the suite will catch it — + # the regression tests assert forward equality only. # - # Omitting it left every non-default `cfg.norm`/`cs_phase` gradient wrong by - # M[l,m] — finite differences showed 40–180% relative error on :schmidt and - # :fourpi, while the dense scalar pair (which never converts) was exact. - # Both are no-ops on the default config, so the hot path is untouched. - # `_ensure_norm_scale_matrix!` lazily BUILDS and caches a constant (l,m) table - # on the config. Its `setindex!` is invisible to a caller but fatal to Zygote - # ("Mutating arrays is not supported") whenever a differentiated function - # reaches it — e.g. `analysis_qst`/`_synthesis_qst`, which have no rrule and - # are traced through. The table does not depend on any differentiated value, - # so declare the whole builder non-differentiable; that covers every traced - # call site at once instead of rewriting each one to dodge the cache. + # `convert_alm_norm!` survives as a standalone public utility for callers who + # want coefficients in some other convention; no transform calls it. It does + # reach `_ensure_norm_scale_matrix!`, which lazily BUILDS and caches a + # constant (l,m) table on the config. That `setindex!` is invisible to a + # caller but fatal to Zygote ("Mutating arrays is not supported") if a + # differentiated function ever reaches it. The table does not depend on any + # differentiated value, so keep the builder declared non-differentiable. ChainRulesCore.@non_differentiable SHTnsKit._ensure_norm_scale_matrix!(::Any) - _needs_norm(cfg) = cfg.norm !== :orthonormal || cfg.cs_phase == false - # A loss that consumes only ONE of a two-output transform hands the other slot # a `ZeroTangent`. The `_adjoint_*` kernels take arrays, so materialise it to # an explicit zero matrix of the right shape rather than letting it reach them. @@ -130,27 +130,10 @@ import SHTnsKit: wigner_d_matrix_deriv # the synthesis adjoint must not, and misses the `wm = 2` doubling for m > 0. # Dense (l+1, m+1) matrix ↔ packed LM-order vector, skipping m % mres ≠ 0. - function _unpack_lm(cfg::SHTnsKit.SHTConfig, Qlm::AbstractVector) - A = zeros(eltype(Qlm), cfg.lmax + 1, cfg.mmax + 1) - @inbounds for m in 0:cfg.mmax - (m % cfg.mres == 0) || continue - for l in m:cfg.lmax - A[l+1, m+1] = Qlm[LM_index(cfg.lmax, cfg.mres, l, m) + 1] - end - end - return A - end - - function _pack_lm(cfg::SHTnsKit.SHTConfig, A::AbstractMatrix) - Qlm = zeros(eltype(A), cfg.nlm) - @inbounds for m in 0:cfg.mmax - (m % cfg.mres == 0) || continue - for l in m:cfg.lmax - Qlm[LM_index(cfg.lmax, cfg.mres, l, m) + 1] = A[l+1, m+1] - end - end - return Qlm - end + # Thin aliases over the canonical pair in src/layout.jl — this file used to + # carry its own copy of both loops. + const _unpack_lm = SHTnsKit.unpack_lm + const _pack_lm = SHTnsKit.pack_lm function ChainRulesCore.rrule(::typeof(SHTnsKit.analysis_packed), cfg::SHTnsKit.SHTConfig, Vr) y = SHTnsKit.analysis_packed(cfg, Vr) diff --git a/ext/SHTnsKitGPUExt.jl b/ext/SHTnsKitGPUExt.jl index ba7738a9..76f1a394 100644 --- a/ext/SHTnsKitGPUExt.jl +++ b/ext/SHTnsKitGPUExt.jl @@ -596,8 +596,7 @@ function gpu_analysis(cfg::SHTConfig, spatial_data; device=get_device(), real_ou Qlm = Array(coeffs) # NO conversion: the kernels emit orthonormal P̄ output and CPU `analysis` is # orthonormal-only, so returning the raw coefficients is what "matching CPU - # analysis" now means. (The sphtor GPU path still converts because its CPU - # twin `analysis_sphtor` still does.) + # analysis" means. The GPU sphtor path does the same, as does its CPU twin. return Qlm end @@ -623,8 +622,8 @@ function gpu_synthesis(cfg::SHTConfig, coeffs; device=get_device(), real_output= size(coeffs, 2) == mmax + 1 || throw(DimensionMismatch("coeffs must have $(mmax+1) columns (mmax+1), got $(size(coeffs, 2))")) # NO conversion: the kernel expects orthonormal input and CPU `synthesis` is - # orthonormal-only, so the coefficients pass straight through. (The sphtor GPU - # path still converts — its CPU twin `synthesis_sphtor` still does.) + # orthonormal-only, so the coefficients pass straight through. The GPU sphtor + # path does the same, as does its CPU twin. coeffs_int = coeffs # Transfer coefficients to GPU diff --git a/src/api_compat.jl b/src/api_compat.jl index 97116ad0..712098f4 100644 --- a/src/api_compat.jl +++ b/src/api_compat.jl @@ -186,14 +186,25 @@ function shtns_set_grid(cfg::SHTConfig, flags::Integer, eps::Real, nlat::Integer x[i+1] = cos(θ[i+1]) end elseif grid_type == 5 # reg_poles, include poles + # Needs both poles: nlat==1 gives π/0 = Inf and θ[1] = 0*Inf = NaN, and + # `_min_nlat_for_grid` clamps only to lmax+1, which is 1 for lmax=0 — so + # the config would come back all-NaN and every transform would silently + # return NaN. Same guard as the `create_regular_config` builder. + nlat >= 2 || throw(ArgumentError("SHT_REGULAR_POLES needs nlat ≥ 2 (got nlat=$nlat)")) for i in 0:(nlat-1) θ[i+1] = i * (π / (nlat-1)) w[i+1] = (π / (nlat-1)) * sin(θ[i+1]) x[i+1] = cos(θ[i+1]) end else - # default to gauss + # default to gauss — including the reverse! pair. Omitting it left an + # unrecognized grid code with south-to-north latitudes while the config + # still reported `is_south_pole_first == false` and `grid_type == :gauss`, + # i.e. a silently mirrored grid returned as success. `_grid_symbol` funnels + # every unknown code here, so nothing else catches a bad flag. x, w = gausslegendre(nlat) + reverse!(x) + reverse!(w) θ = acos.(x) end # Note: south_pole_first will be applied after updating cfg @@ -551,6 +562,12 @@ function save_config(cfg::SHTConfig, filename::String) println(io, "nlat = $(cfg.nlat)") println(io, "nlon = $(cfg.nlon)") println(io, "grid_type = $(cfg.grid_type)") + # `grid_type` alone does not round-trip: `:driscoll_healy` is + # `create_regular_config(...; include_poles=true, use_dh_weights=true)`, + # and without recording that flag `load_config` rebuilt it as a plain + # `:regular` midpoint grid — different θ nodes AND different quadrature + # weights, silently (analysis→synthesis error 8.7e-16 → ~3e-3). + println(io, "use_dh_weights = $(cfg.grid_type === :driscoll_healy)") # Write normalization options println(io, "norm = $(cfg.norm)") @@ -600,7 +617,7 @@ function load_config(filename::String) # Parse value based on expected type if key in ("lmax", "mmax", "mres", "nlat", "nlon", "nlat_padded", "howmany", "spec_dist") params[key] = parse(Int, val) - elseif key in ("cs_phase", "real_norm", "robert_form", "south_pole_first", "allow_padding", "on_the_fly", "use_plm_tables") + elseif key in ("cs_phase", "real_norm", "robert_form", "south_pole_first", "allow_padding", "on_the_fly", "use_plm_tables", "use_dh_weights") params[key] = parse(Bool, val) elseif key in ("grid_type", "norm") params[key] = Symbol(val) @@ -621,7 +638,12 @@ function load_config(filename::String) robert_form=get(params, "robert_form", false) ) else - include_poles = grid_type == :regular_poles + # `:driscoll_healy` is also a pole-inclusive grid, and it needs its + # `use_dh_weights` flag back or it silently degrades to `:regular`. + # Fall back to deriving the flag from grid_type for files written before + # `use_dh_weights` was recorded. + use_dh_weights = get(params, "use_dh_weights", grid_type == :driscoll_healy) + include_poles = grid_type == :regular_poles || grid_type == :driscoll_healy cfg = create_regular_config( params["lmax"], params["nlat"]; mmax=params["mmax"], mres=params["mres"], nlon=params["nlon"], @@ -629,7 +651,8 @@ function load_config(filename::String) cs_phase=get(params, "cs_phase", true), real_norm=get(params, "real_norm", false), robert_form=get(params, "robert_form", false), - include_poles=include_poles + include_poles=include_poles, + use_dh_weights=use_dh_weights ) end diff --git a/src/batch_transforms.jl b/src/batch_transforms.jl index 7e1bb899..553fb827 100644 --- a/src/batch_transforms.jl +++ b/src/batch_transforms.jl @@ -98,12 +98,36 @@ This mirrors the `shtns_set_many` functionality from the SHTns C library. # Batch FFT helpers build one plan from the first field and reuse it for every # sibling slice. Fallback delegates to scalar FFT wrappers so AD element types # keep the same behavior as the non-batch API. +""" + _fftw_planable(T) -> Bool + +Whether `SHTPlan` can carry element type `T`. + +`SHTPlan`'s scratch buffers are `Matrix{ComplexF64}` with FFTW plans built over +them, and FFTW has plans only for `Float32`/`Float64` and their complex forms. +A `ForwardDiff.Dual` therefore cannot go through a plan at all — `analysis!` hits +`plan.Fθk[i,j] = f[i,j]` and raises `MethodError: Float64(::Dual)`. + +The plan-free `cfg`-form transforms have no such limit: their φ transform falls +back to the pure-Julia `_dft_phi` (src/fftutils.jl), which is generic. So the +batch entry points test this and route non-FFTW element types field-by-field +through the `cfg`-form functions — slower than a shared plan, but differentiable, +which is the whole point of passing Duals in. +""" +_fftw_planable(::Type{<:Union{Float32,Float64}}) = true +_fftw_planable(::Type{Complex{Float32}}) = true +_fftw_planable(::Type{Complex{Float64}}) = true +_fftw_planable(::Type) = false + function _batch_fft_phi!(Fφ_batch::AbstractArray{<:Complex,3}, fields::AbstractArray{<:Real,3}) nfields = size(fields, 3) nfields == 0 && return Fφ_batch try - plan = plan_fft!(view(Fφ_batch, :, :, 1), 2; flags=FFTW.ESTIMATE | FFTW.UNALIGNED) + # Reuse the shared plan cache rather than re-planning on every call. + # Every k-slice of a 3D array has identical size and strides, so one cache + # entry serves the whole batch; the cache passes UNALIGNED for us. + plan = _cached_local_fft_plan(:fft, view(Fφ_batch, :, :, 1)) @inbounds for k in 1:nfields Fk = view(Fφ_batch, :, :, k) Xk = view(fields, :, :, k) @@ -125,7 +149,7 @@ function _batch_rfft_phi!(Fφ_batch::AbstractArray{<:Complex,3}, fields::Abstrac nfields == 0 && return Fφ_batch try - plan = plan_rfft(view(fields, :, :, 1), 2; flags=FFTW.ESTIMATE | FFTW.UNALIGNED) + plan = _cached_local_fft_plan(:rfft, view(fields, :, :, 1)) @inbounds for k in 1:nfields mul!(view(Fφ_batch, :, :, k), plan, view(fields, :, :, k)) end @@ -150,7 +174,7 @@ function _batch_ifft_phi!(Fφ_batch::AbstractArray{<:Complex,3}) # unmutated source and are naturally idempotent; only this one corrupts.) k_done = 0 try - plan = plan_ifft!(view(Fφ_batch, :, :, 1), 2; flags=FFTW.ESTIMATE | FFTW.UNALIGNED) + plan = _cached_local_fft_plan(:ifft, view(Fφ_batch, :, :, 1)) @inbounds for k in 1:nfields Fk = view(Fφ_batch, :, :, k) mul!(Fk, plan, Fk) @@ -187,7 +211,7 @@ function _batch_irfft_phi!(f_out::AbstractArray{<:Real,3}, Fφ_batch::AbstractAr nfields == 0 && return f_out try - plan = plan_irfft(view(Fφ_batch, :, :, 1), nlon, 2; flags=FFTW.ESTIMATE | FFTW.UNALIGNED) + plan = _cached_local_fft_plan(:irfft, view(Fφ_batch, :, :, 1), nlon) @inbounds for k in 1:nfields mul!(view(f_out, :, :, k), plan, view(Fφ_batch, :, :, k)) end @@ -689,10 +713,24 @@ function analysis_sphtor_batch(cfg::SHTConfig, Vt_batch::AbstractArray{<:Real,3} size(Vp_batch) == size(Vt_batch) || throw(DimensionMismatch("Vt and Vp must have same shape")) lmax, mmax = cfg.lmax, cfg.mmax - Slm_batch = zeros(ComplexF64, lmax + 1, mmax + 1, nfields) - Tlm_batch = zeros(ComplexF64, lmax + 1, mmax + 1, nfields) - plan = SHTPlan(cfg) + # Follow the inputs, promoted across both components — hardcoding ComplexF64 + # truncated Float32 data and made Dual inputs unwritable. + CT = complex(float(promote_type(eltype(Vt_batch), eltype(Vp_batch)))) + Slm_batch = zeros(CT, lmax + 1, mmax + 1, nfields) + Tlm_batch = zeros(CT, lmax + 1, mmax + 1, nfields) + + if !_fftw_planable(CT) + # No FFTW plan exists for this element type; go through the plan-free + # `cfg`-form transform per field so AD types still work. + for k in 1:nfields + S, T = analysis_sphtor(cfg, view(Vt_batch, :, :, k), view(Vp_batch, :, :, k)) + Slm_batch[:, :, k] .= S + Tlm_batch[:, :, k] .= T + end + return Slm_batch, Tlm_batch + end + plan = SHTPlan(cfg) for k in 1:nfields analysis_sphtor!(plan, view(Slm_batch, :, :, k), view(Tlm_batch, :, :, k), view(Vt_batch, :, :, k), view(Vp_batch, :, :, k)) @@ -735,15 +773,27 @@ function _synthesis_sphtor_batch(cfg::SHTConfig, Slm_batch::AbstractArray{<:Comp nfields = size(Slm_batch, 3) nlat, nlon = cfg.nlat, cfg.nlon - if real_output - Vt_batch = Array{Float64,3}(undef, nlat, nlon, nfields) - Vp_batch = Array{Float64,3}(undef, nlat, nlon, nfields) - else - Vt_batch = Array{ComplexF64,3}(undef, nlat, nlon, nfields) - Vp_batch = Array{ComplexF64,3}(undef, nlat, nlon, nfields) + # Follow the inputs, promoted across both — hardcoding Float64/ComplexF64 + # truncated ComplexF32 data and made Dual coefficients unwritable. + RT = real(float(promote_type(eltype(Slm_batch), eltype(Tlm_batch)))) + CT = complex(RT) + OT = real_output ? RT : CT + Vt_batch = Array{OT,3}(undef, nlat, nlon, nfields) + Vp_batch = Array{OT,3}(undef, nlat, nlon, nfields) + + if !_fftw_planable(CT) + # No FFTW plan for this element type — per-field `cfg`-form transform, + # which stays differentiable. + for k in 1:nfields + Vt, Vp = synthesis_sphtor(cfg, view(Slm_batch, :, :, k), view(Tlm_batch, :, :, k); + real_output=real_output) + Vt_batch[:, :, k] .= Vt + Vp_batch[:, :, k] .= Vp + end + return Vt_batch, Vp_batch end - plan = SHTPlan(cfg) + plan = SHTPlan(cfg) for k in 1:nfields synthesis_sphtor!(plan, view(Vt_batch, :, :, k), view(Vp_batch, :, :, k), view(Slm_batch, :, :, k), view(Tlm_batch, :, :, k); @@ -775,22 +825,37 @@ function analysis_qst_batch(cfg::SHTConfig, Vr_batch::AbstractArray{<:Real,3}, size(Vp_batch) == size(Vr_batch) || throw(DimensionMismatch("Vr and Vp must have same shape")) lmax, mmax = cfg.lmax, cfg.mmax - # Follow the input eltype, as `analysis_batch` does. - CT = complex(float(eltype(Vr_batch))) + # Follow the input eltype, as `analysis_batch` does — but promote across all + # three fields, not just Vr: S/T are computed from Vt/Vp, so keying on Vr + # alone silently truncates them (and makes Dual inputs unwritable) whenever + # the three arrays differ in precision. + CT = complex(float(promote_type(eltype(Vr_batch), eltype(Vt_batch), eltype(Vp_batch)))) Qlm_batch = zeros(CT, lmax + 1, mmax + 1, nfields) Slm_batch = zeros(CT, lmax + 1, mmax + 1, nfields) Tlm_batch = zeros(CT, lmax + 1, mmax + 1, nfields) - plan = SHTPlan(cfg) + if !_fftw_planable(CT) + # No FFTW plan for this element type (e.g. ForwardDiff.Dual) — take the + # plan-free `cfg`-form transform per field so gradients flow. + for k in 1:nfields + Q, S, T = analysis_qst(cfg, view(Vr_batch, :, :, k), + view(Vt_batch, :, :, k), view(Vp_batch, :, :, k)) + Qlm_batch[:, :, k] .= Q + Slm_batch[:, :, k] .= S + Tlm_batch[:, :, k] .= T + end + return Qlm_batch, Slm_batch, Tlm_batch + end + + plan = SHTPlan(cfg) for k in 1:nfields analysis!(plan, view(Qlm_batch, :, :, k), view(Vr_batch, :, :, k)) analysis_sphtor!(plan, view(Slm_batch, :, :, k), view(Tlm_batch, :, :, k), view(Vt_batch, :, :, k), view(Vp_batch, :, :, k)) end - # The scalar plan is orthonormal-only (matching `analysis`/`synthesis`) while - # the sphtor plan converts to cfg's convention, so Q must be converted here - # or this call returns a triple on two normalizations — the same defect - # fixed in `analysis_qst`, which this is the batch form of. + # No normalization conversion here, by design: the scalar plan and the sphtor + # plan are both orthonormal+CS, as is every other transform in the package, so + # Q/S/T come back on one convention. See `analysis_qst`, the non-batch form. return Qlm_batch, Slm_batch, Tlm_batch end @@ -833,8 +898,10 @@ function _synthesis_qst_batch(cfg::SHTConfig, Qlm_batch::AbstractArray{<:Complex nlat, nlon = cfg.nlat, cfg.nlon # Output eltype follows the input, as in `_synthesis_batch` — hardcoding - # Float64/ComplexF64 mismatched the FFTW plan for a ComplexF32 batch. - RT = real(float(eltype(Qlm_batch))) + # Float64/ComplexF64 mismatched the FFTW plan for a ComplexF32 batch. Promote + # across all three spectral arrays: Vt/Vp come from S/T, so keying on Q alone + # would round them down to Q's precision. + RT = real(float(promote_type(eltype(Qlm_batch), eltype(Slm_batch), eltype(Tlm_batch)))) CT = complex(RT) if real_output Vr_batch = Array{RT,3}(undef, nlat, nlon, nfields) @@ -845,11 +912,22 @@ function _synthesis_qst_batch(cfg::SHTConfig, Qlm_batch::AbstractArray{<:Complex Vt_batch = Array{CT,3}(undef, nlat, nlon, nfields) Vp_batch = Array{CT,3}(undef, nlat, nlon, nfields) end - plan = SHTPlan(cfg) - - # Q arrives in cfg's convention (matching S/T and `analysis_qst_batch`), but - # the scalar plan is orthonormal-only — convert, mirroring `_synthesis_qst`. + # Q/S/T all arrive orthonormal+CS — the one convention the whole package uses + # — so nothing is converted here, mirroring `_synthesis_qst`. + + if !_fftw_planable(CT) + # No FFTW plan for this element type — per-field `cfg`-form transform. + for k in 1:nfields + Vr, Vt, Vp = _synthesis_qst(cfg, view(Qlm_batch, :, :, k), view(Slm_batch, :, :, k), + view(Tlm_batch, :, :, k), Val(real_output)) + Vr_batch[:, :, k] .= Vr + Vt_batch[:, :, k] .= Vt + Vp_batch[:, :, k] .= Vp + end + return Vr_batch, Vt_batch, Vp_batch + end + plan = SHTPlan(cfg) for k in 1:nfields synthesis!(plan, view(Vr_batch, :, :, k), view(Qlm_batch, :, :, k); real_output=real_output) diff --git a/src/config.jl b/src/config.jl index 6de67e88..7b84da78 100644 --- a/src/config.jl +++ b/src/config.jl @@ -1076,6 +1076,13 @@ function create_regular_config(lmax::Int, nlat::Int; mmax::Int=lmax, mres::Int=1 else # Use trapezoidal rule with both poles # Poles (θ=0 and θ=π) get half-weight per the trapezoidal rule + # A pole-inclusive grid needs at least the two poles: `nlat == 1` makes + # `h = π/0 = Inf` and `θ[1] = 0*Inf = NaN`, and since the generic + # `nlat ≥ lmax+1` check passes for lmax=0 the whole config would come + # back all-NaN and every later transform would return NaN with no + # exception. Fail here instead, where the cause is visible. + nlat >= 2 || throw(ArgumentError("pole-inclusive grids need nlat ≥ 2 (got nlat=$nlat); " * + "use grid_type=:regular for a single-latitude grid")) h = π / (nlat - 1) for i in 0:(nlat-1) θi = i * h @@ -1226,43 +1233,23 @@ function prepare_plm_tables!(cfg::SHTConfig) end end - # Pre-fuse Nlm so scalar and sphtor kernels can read a single product per - # element on both the P and dP/dx tables. - # zeros (not undef): the l 0 for all i. - dg = Vector{Float64}(undef, lmax + 1) # θ-derivative scratch (reuse P-norm scratch g from above) - for m in 0:mmax - NdP = NdP_tables[m+1] - for i in 1:nlat - s_i = sqrt(max(0.0, 1.0 - cfg.x[i]^2)) - inv_s = s_i == 0 ? 0.0 : 1.0 / s_i # pole guard: avoid 0/0→NaN (see header note) - Plm_norm_and_dPdtheta_row!(g, dg, cfg.x[i], lmax, m) - @inbounds for l in m:lmax - NdP[l+1, i] = -dg[l+1] * inv_s - end - end - end + # The "pre-fused Nlm" NP/NdP tables are, by construction, the tables already + # built above: `Plm_norm_row!` returns P̄ = Nlm·rawP, which is exactly what + # `tbl` holds, and the NdP convention -(dP̄/dθ)/sinθ is exactly what `dtbl` + # holds. Rebuilding them from the same recurrences produced bit-for-bit + # identical arrays at double the build time and double the resident memory, + # while `estimate_table_memory` (which counts (mmax+1)*2 tables) reported + # half the true figure — so a job sized by the estimator allocated twice its + # budget and was OOM-killed. Alias instead. + # + # INVARIANT: these four must stay value-identical. If the NP or NdP + # convention ever diverges from the plm/dplm one, build them separately again + # rather than aliasing — and fix `estimate_table_memory` to match. + # + # Deliberately NOT copies: these are build-once, read-only caches (nothing + # mutates a table after `prepare_plm_tables!` returns). + NP_tables = tables + NdP_tables = dtables # Enable table usage and store in configuration cfg.plm_tables = tables diff --git a/src/device_utils.jl b/src/device_utils.jl index 6b79806d..b2831b76 100644 --- a/src/device_utils.jl +++ b/src/device_utils.jl @@ -281,7 +281,13 @@ cpu_arr = to_device(gpu_arr, :cpu) function to_device(arr::AbstractArray, backend::Symbol=current_backend()) if backend == :cpu return _to_cpu(arr) - elseif backend == :gpu + elseif backend == :gpu || backend == :cuda || backend == :amdgpu + # `:cuda`/`:amdgpu` are aliases for `:gpu`, not extra backends. The config + # vocabulary enforced by `set_config_device!` is `:cpu`/`:cuda`/`:amdgpu`, + # while this function's own is `:cpu`/`:gpu`, and the legacy + # `device_transfer_arrays(cfg, ...)` forwards `cfg.compute_device` + # straight here — so without these a config that `is_gpu_config` reports + # as valid threw `ArgumentError: Unknown backend: cuda` on every transfer. return _to_gpu(arr) else throw(ArgumentError("Unknown backend: $backend")) diff --git a/src/fftutils.jl b/src/fftutils.jl index 5090cfdb..b553989a 100644 --- a/src/fftutils.jl +++ b/src/fftutils.jl @@ -108,14 +108,24 @@ function _cached_local_fft_plan(kind::Symbol, A::AbstractMatrix, nlon::Int=0) try plan = get(_LOCAL_FFT_PLAN_CACHE, key, nothing) if plan === nothing + # UNALIGNED is required, not an optimization choice. The cache key + # covers eltype/size/strides but NOT the base pointer's alignment, so + # a plan built for one array gets reused for another that FFTW may + # consider differently aligned. Without UNALIGNED that reuse throws + # `ArgumentError`, which the callers catch and answer by falling back + # to the pure-Julia O(nlat·nlon²) DFT — an order-of-magnitude + # slowdown with no error surfaced. Forfeiting the aligned SIMD + # codelets is much cheaper than forfeiting the FFT. The batch helpers + # in batch_transforms.jl pass the same flag for the same reason. + flags = FFTW.ESTIMATE | FFTW.UNALIGNED plan = if kind === :fft - plan_fft!(A, 2; flags=FFTW.ESTIMATE) + plan_fft!(A, 2; flags) elseif kind === :ifft - plan_ifft!(A, 2; flags=FFTW.ESTIMATE) + plan_ifft!(A, 2; flags) elseif kind === :rfft - plan_rfft(A, 2; flags=FFTW.ESTIMATE) + plan_rfft(A, 2; flags) elseif kind === :irfft - plan_irfft(A, nlon, 2; flags=FFTW.ESTIMATE) + plan_irfft(A, nlon, 2; flags) else throw(ArgumentError("unknown FFT plan kind: $kind")) end diff --git a/src/layout.jl b/src/layout.jl index 7684a9ed..4cec7b4a 100644 --- a/src/layout.jl +++ b/src/layout.jl @@ -252,3 +252,74 @@ function im_from_lm(lm::Int, lmax::Int, mres::Int) end end + +# ============================================================================== +# Packed ↔ dense (l,m) conversion +# ============================================================================== +# +# The mapping between the packed LM-order vector and the dense (l+1, m+1) matrix +# is needed in several places: the serial `analysis_packed`/`synthesis_packed`, +# their distributed twins in ext/ParallelLocal.jl, and the packed rrules in +# ext/SHTnsKitAdvancedADExt.jl. It used to be open-coded at each of those sites, +# which cost the `m % mres == 0` guard three separate bug fixes — `LM_index` +# throws on any order that is not a multiple of `mres`, so a loop over the full +# `0:mmax` range dies on the first such order. Everything routes through the two +# functions below now, so the guard exists once. + +""" + unpack_lm!(A::AbstractMatrix, cfg, Qlm::AbstractVector) -> A + +Scatter the packed LM-order coefficient vector `Qlm` into the dense +`(lmax+1, mmax+1)` matrix `A`. Orders with `m % cfg.mres != 0` are absent from +the packed layout and are left untouched in `A` (pass a zeroed `A`, or use +[`unpack_lm`](@ref), if those entries must be zero). +""" +function unpack_lm!(A::AbstractMatrix, cfg, Qlm::AbstractVector) + lmax, mmax, mres = cfg.lmax, cfg.mmax, cfg.mres + size(A, 1) == lmax + 1 || throw(DimensionMismatch("A must have $(lmax+1) rows")) + size(A, 2) == mmax + 1 || throw(DimensionMismatch("A must have $(mmax+1) columns")) + length(Qlm) == cfg.nlm || throw(DimensionMismatch("Qlm must have length $(cfg.nlm)")) + @inbounds for m in 0:mmax + (m % mres == 0) || continue + for l in m:lmax + A[l+1, m+1] = Qlm[LM_index(lmax, mres, l, m) + 1] + end + end + return A +end + +""" + unpack_lm(cfg, Qlm::AbstractVector) -> Matrix + +Allocating form of [`unpack_lm!`](@ref); entries with no packed counterpart are +zero. Element type follows `Qlm`. +""" +unpack_lm(cfg, Qlm::AbstractVector) = + unpack_lm!(zeros(eltype(Qlm), cfg.lmax + 1, cfg.mmax + 1), cfg, Qlm) + +""" + pack_lm!(Qlm::AbstractVector, cfg, A::AbstractMatrix) -> Qlm + +Gather the dense `(lmax+1, mmax+1)` matrix `A` into the packed LM-order vector +`Qlm`, skipping orders with `m % cfg.mres != 0`. Inverse of [`unpack_lm!`](@ref). +""" +function pack_lm!(Qlm::AbstractVector, cfg, A::AbstractMatrix) + lmax, mmax, mres = cfg.lmax, cfg.mmax, cfg.mres + size(A, 1) == lmax + 1 || throw(DimensionMismatch("A must have $(lmax+1) rows")) + size(A, 2) == mmax + 1 || throw(DimensionMismatch("A must have $(mmax+1) columns")) + length(Qlm) == cfg.nlm || throw(DimensionMismatch("Qlm must have length $(cfg.nlm)")) + @inbounds for m in 0:mmax + (m % mres == 0) || continue + for l in m:lmax + Qlm[LM_index(lmax, mres, l, m) + 1] = A[l+1, m+1] + end + end + return Qlm +end + +""" + pack_lm(cfg, A::AbstractMatrix) -> Vector + +Allocating form of [`pack_lm!`](@ref). Element type follows `A`. +""" +pack_lm(cfg, A::AbstractMatrix) = pack_lm!(zeros(eltype(A), cfg.nlm), cfg, A) diff --git a/src/parallel_dense.jl b/src/parallel_dense.jl index c9bffa83..0e6991d7 100644 --- a/src/parallel_dense.jl +++ b/src/parallel_dense.jl @@ -45,17 +45,21 @@ function dist_SH_Yrotate(cfg::SHTnsKit.SHTConfig, Alm::AbstractMatrix, beta::Rea lmax, mmax = cfg.lmax, cfg.mmax size(Alm,1)==lmax+1 && size(Alm,2)==mmax+1 || throw(DimensionMismatch("Alm dims")) size(Rlm,1)==lmax+1 && size(Rlm,2)==mmax+1 || throw(DimensionMismatch("Rlm dims")) - Q = Vector{complex(float(eltype(Alm)))}(undef, cfg.nlm) - @inbounds for m in 0:mmax, l in m:lmax - idx = SHTnsKit.LM_index(lmax, cfg.mres, l, m) + 1 - Q[idx] = Alm[l+1, m+1] - end + # A Y-rotation mixes orders, so it cannot be expressed in an mres-strided + # layout at all: rotating an mres=2 field produces m=1 components with nowhere + # to live. `shtns_rotation_apply_real` states the same restriction. Say so + # here rather than failing downstream with `m must be a multiple of mres` + # (what the un-strided loop used to do) or `LM packed size mismatch`. + cfg.mres == 1 || throw(ArgumentError("dist_SH_Yrotate requires mres==1 (got mres=$(cfg.mres)); " * + "a Y-rotation mixes orders and cannot be represented in an mres-strided layout")) + # Canonical packed↔dense pair (src/layout.jl) instead of an open-coded loop. + Q = SHTnsKit.pack_lm(cfg, complex(float(eltype(Alm))).(Alm)) R = similar(Q) SHTnsKit.SH_Yrotate(cfg, Q, beta, R) - @inbounds for m in 0:mmax, l in m:lmax - idx = SHTnsKit.LM_index(lmax, cfg.mres, l, m) + 1 - Rlm[l+1, m+1] = R[idx] - end + # Orders absent from the packed layout have no rotated value to write back; + # zero them rather than leaving whatever the caller's buffer held. + fill!(Rlm, zero(eltype(Rlm))) + SHTnsKit.unpack_lm!(Rlm, cfg, R) return Rlm end @@ -76,6 +80,11 @@ function dist_SH_mul_mx!(cfg::SHTnsKit.SHTConfig, mx::AbstractVector{<:Real}, Al # - Y_{l-1}^m contributes to Y_l^m via b_{l-1}^m (the upward coefficient from l-1) # - Y_{l+1}^m contributes to Y_l^m via a_{l+1}^m (the downward coefficient from l+1) @inbounds for m in 0:mmax, l in m:lmax + # `LM_index` throws unless m is a multiple of mres, so stride like every + # other packed-index loop in the package (`analysis_packed`, + # `synthesis_packed`, `pack_lm!`). Without this the whole dense operator + # path died on the first m=1 for any mres>1 config. + (m % cfg.mres == 0) || continue acc = zero(promote_type(eltype(Alm), eltype(Rlm), complex(eltype(mx)))) # eltype-preserving accumulator (AD-safe) # Contribution from lower degree neighbor Y_{l-1}^m if l > m && l > 0 diff --git a/src/plan.jl b/src/plan.jl index 8ebdc384..d67c389b 100644 --- a/src/plan.jl +++ b/src/plan.jl @@ -450,8 +450,8 @@ function analysis!(plan::SHTPlan, alm_out::AbstractMatrix, f::AbstractMatrix) # converted to cfg's convention, which meant swapping the plan in silently # changed the coefficients — plan∘plan and dense∘dense each round-tripped, # but mixing them (e.g. `synthesis(cfg, analysis!(plan, alm, f))`) was off by - # M[l,m] plus a sign on odd m. The sphtor plan methods above DO convert, and - # correctly so: their non-plan twins convert too. + # M[l,m] plus a sign on odd m. The sphtor plan methods above do not convert + # either, for the same reason: their non-plan twins do not. return alm_out end diff --git a/src/qst_transforms.jl b/src/qst_transforms.jl index 095717a5..d82eee2c 100644 --- a/src/qst_transforms.jl +++ b/src/qst_transforms.jl @@ -141,11 +141,9 @@ Complex version of QST to spatial transform, preserving complex values. """ function synthesis_qst_cplx(cfg::SHTConfig, Qlm::AbstractMatrix, Slm::AbstractMatrix, Tlm::AbstractMatrix) validate_qst_dimensions(Qlm, Slm, Tlm, cfg) - # `synthesis_sphtor_cplx` delegates to `_synthesis_sphtor`, which converts - # cfg→internal, while `synthesis_cplx` is orthonormal-only — so Q needs the - # same conversion the real-output `_synthesis_qst` applies. (Both `_cplx` - # sphtor wrappers are thin delegators to the CONVERTING implementations; - # reading the wrapper bodies alone suggests otherwise.) + # No conversion on either component: `synthesis_cplx` and the `_synthesis_sphtor` + # that `synthesis_sphtor_cplx` delegates to are both orthonormal-only, so Q/S/T + # enter on one convention. Matches the real-output `_synthesis_qst`. Vr = synthesis_cplx(cfg, Qlm) Vt, Vp = synthesis_sphtor_cplx(cfg, Slm, Tlm) @@ -163,9 +161,9 @@ function analysis_qst_cplx(cfg::SHTConfig, Vr::AbstractMatrix{<:Complex}, Vt::Ab # Validate input dimensions validate_vector_spatial_dimensions(Vr, Vt, Vp, cfg) - # Transform each component. `analysis_sphtor_cplx` delegates to - # `analysis_sphtor`, which converts internal→cfg, so Q must match or this - # returns a triple on two normalizations (see `analysis_qst`). + # Transform each component. `analysis` and the `analysis_sphtor` that + # `analysis_sphtor_cplx` delegates to are both orthonormal-only, so the + # returned triple sits on one convention (see `analysis_qst`). Qlm = analysis(cfg, Vr) Slm, Tlm = analysis_sphtor_cplx(cfg, Vt, Vp) @@ -206,10 +204,9 @@ end function _synthesis_qst_l(cfg::SHTConfig, Qlm::AbstractMatrix, Slm::AbstractMatrix, Tlm::AbstractMatrix, ltr::Int, ::Val{real_output}) where {real_output} validate_qst_dimensions(Qlm, Slm, Tlm, cfg) - # Same convention split as `_synthesis_qst`: `_synthesis_sphtor_l` converts - # cfg→internal, `_synthesis_l` does not, and `analysis_qst_l` (which routes - # through `analysis_qst`) returns Q in cfg's convention — so convert Q here - # too, or the degree-limited round trip breaks for any non-default norm. + # Same single convention as `_synthesis_qst`: neither `_synthesis_l` nor + # `_synthesis_sphtor_l` converts, and `analysis_qst_l` returns Q orthonormal, + # so the degree-limited round trip closes for every `cfg.norm`. Vr = _synthesis_l(cfg, Qlm, ltr, Val(real_output)) Vt, Vp = _synthesis_sphtor_l(cfg, Slm, Tlm, ltr, Val(real_output)) return Vr, Vt, Vp @@ -224,9 +221,9 @@ function analysis_qst_ml(cfg::SHTConfig, im::Int, Vr_m::AbstractVector{<:Complex # Transform each component for this specific mode Ql = analysis_packed_ml(cfg, im, Vr_m, ltr) Sl, Tl = analysis_sphtor_ml(cfg, im, Vt_m, Vp_m, ltr) - # `analysis_sphtor_ml` converts internal→cfg but `analysis_packed_ml` does - # not, so without this the returned triple sits on two normalizations — the - # same defect fixed in `analysis_qst`. Mirrors that function's scaling. + # `analysis_packed_ml` and `analysis_sphtor_ml` are both orthonormal-only, so + # the returned triple sits on one convention — no scaling here, matching + # `analysis_qst`. return Ql, Sl, Tl end @@ -237,9 +234,9 @@ end Mode-limited synthesis for specific azimuthal mode im. """ function synthesis_qst_ml(cfg::SHTConfig, im::Int, Ql::AbstractVector{<:Complex}, Sl::AbstractVector{<:Complex}, Tl::AbstractVector{<:Complex}, ltr::Int) - # Synthesize each component for this specific mode - # Inverse of the conversion in `analysis_qst_ml`: Q arrives in cfg's - # convention (matching S/T) but `synthesis_packed_ml` expects internal. + # Synthesize each component for this specific mode. Exact inverse of + # `analysis_qst_ml`: Q/S/T all arrive orthonormal, which is what + # `synthesis_packed_ml` and `synthesis_sphtor_ml` expect — nothing to convert. Vr_m = synthesis_packed_ml(cfg, im, Ql, ltr) Vt_m, Vp_m = synthesis_sphtor_ml(cfg, im, Sl, Tl, ltr) diff --git a/src/transforms.jl b/src/transforms.jl index 64d30f0b..470d9fe3 100644 --- a/src/transforms.jl +++ b/src/transforms.jl @@ -102,7 +102,17 @@ function analysis_axisym(cfg::SHTConfig, Vr::AbstractVector{<:Real}) end end - return Ql # No phi scaling needed for single-mode transform (proper inverse of synthesis_axisym) + # φ quadrature factor. The full `analysis` applies `cfg.cphi` to an FFT output + # whose bin 0 already carries an implicit `nlon` from the DFT sum; an + # axisymmetric profile is constant in φ, so the whole `cphi*nlon = 2π` must be + # applied explicitly here. Without it this is NOT the inverse of + # `synthesis_axisym` (which matches `synthesis` exactly) and disagrees with + # the m=0 column of `analysis` by 1/2π. + scaleφ = cfg.cphi * cfg.nlon + @inbounds for l in 0:lmax + Ql[l+1] *= scaleφ + end + return Ql end """ @@ -114,17 +124,9 @@ function analysis_packed(cfg::SHTConfig, Vr::AbstractVector{<:Real}) length(Vr) == cfg.nspat || throw(DimensionMismatch("Vr must have length $(cfg.nspat)")) f = reshape(Vr, cfg.nlat, cfg.nlon) alm_mat = analysis(cfg, f) - Qlm = Vector{eltype(alm_mat)}(undef, cfg.nlm) # Dense matrix output is converted back to SHTns-compatible packed LM # order, skipping unsupported m values when mres > 1. - @inbounds for m in 0:cfg.mmax - (m % cfg.mres == 0) || continue - for l in m:cfg.lmax - lm = LM_index(cfg.lmax, cfg.mres, l, m) + 1 - Qlm[lm] = alm_mat[l+1, m+1] - end - end - return Qlm + return pack_lm(cfg, alm_mat) end """ @@ -134,16 +136,9 @@ Packed scalar synthesis from Qlm (LM order) to flattened real grid (length nlat* """ function synthesis_packed(cfg::SHTConfig, Qlm::AbstractVector{<:Complex}) length(Qlm) == cfg.nlm || throw(DimensionMismatch("Qlm must have length $(cfg.nlm)")) - alm_mat = zeros(eltype(Qlm), cfg.lmax+1, cfg.mmax+1) # Packed LM order stores only valid (l,m) pairs. Expand to dense # (l+1,m+1) so the core synthesis kernel can be reused. - @inbounds for m in 0:cfg.mmax - (m % cfg.mres == 0) || continue - for l in m:cfg.lmax - lm = LM_index(cfg.lmax, cfg.mres, l, m) + 1 - alm_mat[l+1, m+1] = Qlm[lm] - end - end + alm_mat = unpack_lm(cfg, Qlm) f = synthesis(cfg, alm_mat; real_output=true) return vec(f) end @@ -254,7 +249,12 @@ function analysis_axisym_l(cfg::SHTConfig, Vr::AbstractVector{<:Real}, ltr::Int) end end - return Ql # No phi scaling needed for single-mode transform (proper inverse of synthesis_axisym_l) + # Same φ quadrature factor as `analysis_axisym` — see the comment there. + scaleφ = cfg.cphi * cfg.nlon + @inbounds for l in eachindex(Ql) + Ql[l] *= scaleφ + end + return Ql end """ diff --git a/test/serial/test_basic_transforms.jl b/test/serial/test_basic_transforms.jl index a2f46e7c..d1a551de 100644 --- a/test/serial/test_basic_transforms.jl +++ b/test/serial/test_basic_transforms.jl @@ -297,12 +297,16 @@ using SHTnsKit @test length(Ql_rec) == lmax + 1 # m=0 coefficients should be real (imaginary part ~0) @test maximum(abs.(imag.(Ql_rec))) < 1e-10 - # Roundtrip may differ by a consistent scaling factor. - # Verify proportionality: Ql_rec = scale * Ql for a single scale factor. - # Find scale from first non-negligible coefficient. - idx = findfirst(i -> abs(real(Ql[i])) > 1e-10, 1:length(Ql)) - scale = real(Ql_rec[idx]) / real(Ql[idx]) - @test isapprox(real.(Ql_rec), scale .* real.(Ql); rtol=1e-9, atol=1e-11) + # The round trip is an IDENTITY, not a proportionality. This used to + # divide out a fitted scale factor, which made the assertion blind to the + # missing φ quadrature factor (cphi*nlon = 2π) in `analysis_axisym` — the + # test passed while every returned coefficient was 1/2π too small. Assert + # the absolute values so a revert is caught. + @test isapprox(real.(Ql_rec), real.(Ql); rtol=1e-9, atol=1e-11) + # And pin it to the full transform: axisym analysis must equal the m=0 + # column of `analysis` on the same field. + f2d = repeat(f_lat, 1, cfg.nlon) + @test isapprox(Ql_rec, analysis(cfg, f2d)[:, 1]; rtol=1e-9, atol=1e-11) end @testset "Axisymmetric truncated transforms (analysis_axisym_l/synthesis_axisym_l)" begin diff --git a/test/serial/test_plan.jl b/test/serial/test_plan.jl index a597f54a..34c71f66 100644 --- a/test/serial/test_plan.jl +++ b/test/serial/test_plan.jl @@ -142,7 +142,18 @@ end # `synthesis`, so it must be ORTHONORMAL like them — not merely # self-consistent. A self-roundtrip test cannot tell the two apart: # plan∘plan closes under either convention, and only MIXING them breaks. - # These equality assertions are what would catch a revert. + # These assertions are what would catch a revert. + # + # They compare to a TOLERANCE, not bit-for-bit. The two paths reach the + # same φ transform through different FFTW plans — `synthesis` uses the + # shared cache (built UNALIGNED, since it reuses plans across arbitrary + # caller arrays) while `SHTPlan` plans its own stably-aligned buffers — + # so FFTW may pick different codelets and the results can differ in the + # last ulp. That is not a convention bug and varies by CPU and FFTW + # build: `==` passed on arm64 and failed on x86_64 CI at 2e-16. + # The tolerance below is ~4 orders above roundoff and ~10 orders below a + # normalization revert, which is an O(1) relative error (M[l,m] is 40-180% + # off for :schmidt/:fourpi), so a revert is still caught cleanly. lmax = 6 for (nrm, cs) in ((:orthonormal, true), (:schmidt, true), (:fourpi, false)) cfg = create_gauss_config(lmax, lmax + 2; nlon=2*lmax + 1, @@ -154,12 +165,14 @@ end f_plan = zeros(cfg.nlat, cfg.nlon) synthesis!(plan, f_plan, alm) @test all(isfinite, f_plan) - # identical to the non-planned transform, not just close - @test f_plan == synthesis(cfg, alm; real_output=true) + # the non-planned transform to roundoff — same convention, not just + # self-consistent (see the note above on why this is not `==`) + @test isapprox(f_plan, synthesis(cfg, alm; real_output=true); + rtol=1e-12, atol=1e-14) alm_back = zeros(ComplexF64, lmax + 1, lmax + 1) analysis!(plan, alm_back, f_plan) - @test alm_back == analysis(cfg, f_plan) + @test isapprox(alm_back, analysis(cfg, f_plan); rtol=1e-12, atol=1e-14) # and the plan's own roundtrip still recovers the input @test isapprox(alm_back, alm; rtol=1e-9, atol=1e-11) end