From 4c63395a2baf8a38b47c03a7483cb2417dfcbab3 Mon Sep 17 00:00:00 2001 From: Subhajit Kar Date: Tue, 11 Aug 2026 11:50:03 +0300 Subject: [PATCH 1/2] Fix correctness and performance defects found by codebase review Two rounds of review over the branch diff and then over the files the branch never touched. Every finding below was reproduced before fixing and re-verified after. Breaking (see CHANGELOG.md for porting instructions): - analysis_axisym / analysis_axisym_l omitted the phi quadrature factor cphi*nlon = 2pi, so every coefficient was 1/2pi too small and they inverted neither synthesis_axisym nor the m=0 column of analysis. - dist_synthesis! / dist_synthesis_sphtor! keep rejecting real_output=false with real-backed output arrays. real_output=false now computes a different function (m>=0 half, no Hermitian mirror), so no tolerance can accept the old usage; the error messages now name the porting path. Correctness: - Batch QST and sphtor transforms derived their output eltype from a single input instead of promoting across all of them, silently truncating double-precision components (2.05e-8 error, now 1.39e-17). - ForwardDiff could not flow through any plan-based batch transform. SHTPlan is FFTW-backed and cannot hold Dual, so the batch entry points now route non-FFTW eltypes through the plan-free cfg-form transforms. - load_config silently rebuilt :driscoll_healy grids as :regular, changing theta nodes and weights (8.7e-16 -> 3e-3 error). save_config now records use_dh_weights. - shtns_set_grid returned success for unrecognized grid codes while producing a mirrored grid; the fallback branch was missing reverse!. - Pole-inclusive grids with nlat == 1 built an all-NaN config from pi/0 and returned NaN from every later transform without raising. - dist_SH_mul_mx! crashed on every mres > 1 config (missing m % mres guard). - dist_SH_Yrotate crashed on mres > 1; a Y-rotation mixes orders and is not representable in an mres-strided layout, so it now says so up front. - device_transfer_arrays rejected :cuda/:amdgpu, the config vocabulary. - Eleven comments left by the orthonormal refactor prescribed conversions the code no longer performs, including the AD extension's only statement of the rrule/primal normalization contract, which stated its inverse. Performance: - Distributed analysis went from three collectives per call to one on first use and none after: the phi_is_local_all / theta_is_distributed predicates are reduced together in a single Allreduce and cached per (pencil, communicator), with a rank-symmetric eviction policy. - dist_synthesis_packed_cplx is single-pass, down from two full distributed syntheses; it now matches the serial reference exactly. - prepare_plm_tables! built NP/NdP tables bit-for-bit identical to plm/dplm. They now alias, halving table memory and build time, and estimate_table_memory finally matches reality (it reported half). - Cached FFT plans are built UNALIGNED, removing a silent fallback to an O(n^2) DFT; batch FFT helpers reuse the cache instead of re-planning. Internal and tests: - 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 needed fixing three separate times across those copies. - The axisym round-trip test fitted and divided out its own scale factor, so it passed while every coefficient was 1/2pi off. It now asserts the absolute identity and pins axisym analysis to the m=0 column of analysis. Verified: serial suite 67279/67279; MPI comprehensive and audit-fix suites pass at 2 and 4 ranks; distributed analysis matches serial to 2.2e-16 on 1/2/4 ranks for both decompositions. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.md | 79 +++++++++++ ext/ParallelLocal.jl | 36 ++--- ext/ParallelTransforms.jl | 203 +++++++++++++++++++++++---- ext/ParallelTransposeTransforms.jl | 17 +-- ext/SHTnsKitAdvancedADExt.jl | 65 ++++----- ext/SHTnsKitGPUExt.jl | 7 +- src/api_compat.jl | 31 +++- src/batch_transforms.jl | 132 +++++++++++++---- src/config.jl | 61 ++++---- src/device_utils.jl | 8 +- src/fftutils.jl | 18 ++- src/layout.jl | 71 ++++++++++ src/parallel_dense.jl | 27 ++-- src/plan.jl | 4 +- src/qst_transforms.jl | 33 ++--- src/transforms.jl | 38 ++--- test/serial/test_basic_transforms.jl | 16 ++- 17 files changed, 613 insertions(+), 233 deletions(-) create mode 100644 CHANGELOG.md 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 From 2007cb0ab5168fab79c41aa68198ebec029b7bf2 Mon Sep 17 00:00:00 2001 From: Subhajit Kar Date: Tue, 11 Aug 2026 12:13:11 +0300 Subject: [PATCH 2/2] Compare planned and non-planned transforms to a tolerance, not bit-exactly MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CI failed on x86_64 at test_plan.jl:158 with a 2e-16 relative difference, while the same assertion passed on arm64. The cause is this branch's own UNALIGNED change: `synthesis` reaches its phi transform through the shared plan cache, which is now built UNALIGNED because it reuses plans across arbitrary caller arrays, while `SHTPlan` plans its own stably-aligned buffers with default flags. FFTW can pick different codelets for the two, so results may differ in the last ulp. That is legitimate floating point, not a convention bug, and it varies by CPU and FFTW build — so bit-exactness was never a property to assert. Keeping the plan's own flags as they are: SHTPlan owns buffers of stable alignment, and forcing UNALIGNED there would permanently give up the aligned SIMD codelets that are the entire point of planning, purely to satisfy an over-specified test. The assertions become isapprox at rtol=1e-12, which is ~4 orders above roundoff and ~10 orders below what this test exists to catch. Verified that a simulated normalization revert still fails it: the relative error would be 0.845 for :schmidt and 4.29 for :fourpi. Co-Authored-By: Claude Opus 5 (1M context) --- test/serial/test_plan.jl | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) 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