diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 68155489..b98e138f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -204,6 +204,7 @@ jobs: mpiexec -n 4 julia --project=. test/parallel/test_mpi_plan_preflight.jl mpiexec -n 2 julia --project=. test/parallel/test_mpi_transpose_operand_preflight.jl mpiexec -n 2 julia --project=. test/parallel/test_mpi_comm_cleanup.jl + mpiexec -n 4 julia --project=. test/parallel/test_mpi_2d_alignment.jl mpiexec -n 2 julia --project=. test/parallel/test_mpi_ad_tangent_spaces.jl mpiexec -n 8 julia --project=. test/parallel/test_mpi_audit_fixes.jl env: diff --git a/CHANGELOG.md b/CHANGELOG.md index 40aa0694..b2483c30 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -37,8 +37,91 @@ coefficient was `1/2π` too small — they inverted neither `synthesis_axisym` n 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. +**`analysis` and the direct evaluators now honour `phi_scale`; an unset +`phi_scale` resolves to `:dft`.** Two halves of one convention had drifted apart. + +*`analysis` ignored `phi_scale` entirely.* `synthesis` scales its Fourier bins by +`phi_inv_scale(cfg)`, but `analysis` applied a fixed `cfg.cphi`, so under `:quad` +the pair was not mutually inverse: `analysis(cfg, synthesis(cfg, alm))` came back +as `alm / 2π` exactly. Analysis now applies `cphi · nlon / phi_inv_scale(cfg)`, +which is `cphi` under `:dft` and restores the inverse property under `:quad`. The +same factor was threaded through the batch, complex-packed, planned, distributed +and adjoint analysis paths so they all agree. + +*The direct evaluators applied no φ factor at all.* `synthesis_point`, +`synthesis_point_cplx`, `synthesis_axisym`, `synthesis_axisym_l`, `SH_to_lat`, +`SH_to_lat_cplx`, `SHqst_to_point`, `SH_to_grad_point`, `SHqst_to_lat` and the +PencilArray local evaluations each disagreed with the grid they claim to sample by +exactly 2π under `:quad`. They now carry the same `phi_inv_scale(cfg)/nlon` factor +`synthesis` does. + +*An unset `phi_scale` (`:auto`) used to fall back to a grid-type guess* — +`grid_type == :gauss ? nlon : nlon/2π` — so a regular grid built through the +exported `SHTConfig(; …)` keyword constructor disagreed by 2π with the identical +grid from `create_regular_config`, which sets `:dft` explicitly. Every constructor +emits `:dft`, so an unset value now means `:dft` too. + +**Default `:dft` behaviour is unchanged in all three cases**; only `:quad` and +hand-built `:auto` configurations move, and they move to the values that make +`analysis` and `synthesis` inverses. + +**Order-mixing rotations now reject `mmax < lmax` instead of truncating.** +A Wigner-d rotation through a general `β` couples `Y_l^m` to every `Y_l^{m'}` with +`|m'| ≤ l`. When storage stopped at `mmax < lmax`, the `|m'| > mmax` components +were silently dropped — measured at `lmax = 8`, that discarded **14.8 %** of the +field's energy at `mmax = 5` and **24.0 %** at `mmax = 3`, with no error and no +warning. `SH_Yrotate`, `SH_Yrotate90`, `SH_Xrotate90` and the Euler-angle API now +raise an `ArgumentError` on such a configuration. Pure Z-rotations are unaffected: +`β ≡ 0` is diagonal and `β ≡ π` is anti-diagonal (`m' = -m`), so both still work +at any `mmax`. + +*Porting:* use `mmax == lmax` for anything but a Z-rotation. Results that appeared +to work before were missing the truncated energy. + +**Structural `SHTConfig` fields are no longer silently inconsistent.** +Assigning `cfg.lmax = 10` left `size(cfg.Nlm) == (7, 7)` while the transforms index +it as `(lmax+1, mmax+1)` under `@inbounds` — an out-of-bounds read of a live array. +`lmax`, `mmax` and `mres` now rebuild the derived spectral layout (`Nlm`, `nlm`, +`li`, `mi`, cached scale matrix and m-ordering) and drop the now-stale Legendre +tables; `nlat`, `nlon`, `grid_type`, `nlm`, `li`, `mi` and `nspat` raise an +`ArgumentError` pointing at the `create_*_config` constructors, because there is no +grid-type-independent way to regenerate the quadrature in place. + +**The exported `SHTConfig(; …)` keyword constructor validates its invariants.** +It previously checked nothing, so a hand-built configuration could violate +`nlon ≥ 2*mmax+1` and then silently synthesize an all-zero field for any mode it +could not resolve, or hand `use_rfft=true` a raw `BoundsError`. It now enforces the +same constraints the `create_*_config` helpers always have. The exported keyword +signature is otherwise unchanged. + ### Fixed +- **`analysis_turbo` / `synthesis_turbo` ignored `mres`.** Both walked a bare + `0:mmax` instead of `0:mres:mmax`, so `analysis_turbo` populated — and + `synthesis_turbo` consumed — coefficient columns an `mres > 1` transform has no + storage for. The disagreement with `analysis`/`synthesis` was O(1), not + roundoff. Both now share the core's cached `m` ordering. The turbo pair also + nested `@threads :static`, which is illegal inside an outer threaded region; + they now fall back to a serial loop there, using the same predicate the core + orchestrators use. +- **Rotation pullbacks read coefficients the primal had already overwritten.** + `SH_Yrotate`, `shtns_rotation_apply_cplx` and `shtns_rotation_apply_real` + captured their primal *input* and read it lazily, so an in-place call + (`Rlm === Qlm`) — or any caller reusing the buffer before the pullback ran — + silently corrupted the angle gradient. Each rule now snapshots what it needs at + primal time. Applies to both the ChainRules and Zygote adjoints. +- **`rrule`s for `analysis`/`synthesis`/`analysis_sphtor`/`synthesis_sphtor` + declared fewer keyword arguments than their primals.** Passing `use_rfft` or + `fft_scratch` — even at its default — made ChainRules skip the rule entirely and + fall through to source tracing. The keywords select a different FFT + implementation of the same linear operator, so the adjoints are unchanged. +- **`synthesis_qst` and `analysis_qst` had no `rrule` at all**, so + differentiating a QST pipeline fell through to Zygote's source tracing and + crashed inside FFTW. Each adjoint is the existing scalar and sphtor adjoints + side by side. +- **`shtns_rotation_apply_real` reported a bare size mismatch for `mres > 1` + configurations**, leaving the caller to reverse-engineer why. The message now + names `mres` and states the restriction, matching `dist_SH_Yrotate`. - **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 @@ -95,6 +178,43 @@ compensated for the old scale downstream must drop that compensation. ### Internal +- **Removed the dead parallel FFT-plan cache.** The `SHTNSKIT_CACHE_PENCILFFTS` + environment variable and the `fft_plan_cache_enabled` / `set_fft_plan_cache!` / + `enable_fft_plan_cache!` / `disable_fft_plan_cache!` controls forwarded to a + cache in the parallel extension whose only reader, `_get_or_plan`, had no call + sites — the "plans" it stored were `NamedTuple` placeholders the FFT wrappers + ignored, so every knob was a no-op. The cache the transforms actually use is now + in `src/fftutils.jl`, shared by the serial and distributed paths, and the same + four controls address it without requiring the extension to be loaded. + `SHTNSKIT_FFT_PLAN_CACHE` is the current spelling; the old name still works. +- **`DistributedSpectralPlan2D` no longer attaches a finalizer.** `close` frees + the plan's `l_comm` / `m_comm` sub-communicators, and `MPI_Comm_free` is + collective; a finalizer runs at whatever point that rank's garbage collector + fires, which is rank-local and nondeterministic. Cleanup is explicit only — call + `close(plan)` collectively. Leaking two communicators until `MPI_Finalize` is + strictly better than a nondeterministic collective. +- **`spatial_view(cfg, A)` is exported**, the missing bridge in the padding API: + `allocate_padded_spatial` returns an array with `nlat_padded ≥ nlat` rows while + every transform requires exactly `nlat`, so the padded buffer could not be passed + to `analysis` at all. The view keeps the padded column stride, so it preserves + what the padding is for. +- **`set_batch_size!` is documented as advisory.** `howmany` / `spec_dist` mirror + the SHTns C batch descriptors and are stored for interoperability, but the Julia + batch entry points take the field count from `size(fields, 3)`; the old docstring + claimed otherwise. +- **Regular-grid quadrature exactness is documented.** Fejér and Clenshaw–Curtis + rules with `nlat` nodes are exact only through degree `nlat - 1`, and analysis + integrates degree `2*lmax`, so the equiangular grids need `nlat ≥ 2*lmax + 1` — + where Gauss–Legendre needs `nlat = lmax + 1`. Below that threshold nothing warns + and `analysis ∘ synthesis` is not an identity (7.2e-2 relative error at + `lmax = 8, nlat = 10`). Both `create_regular_config` and `docs/src/grids.md` now + say so with measured numbers. +- New regression coverage: planned-vs-`cfg` conformance across the tables / rfft / + Robert-form matrix; the `mmax < lmax` rotation guard; angle-gradient survival + under an in-place primal; evaluator `phi_scale` agreement; turbo `mres`; and a + 4-rank `test_mpi_2d_alignment.jl` for the 2D spectral-plan alignment + preconditions, wired into CI. + - `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/docs/Distributed_SHTnsKit_Guide.md b/docs/Distributed_SHTnsKit_Guide.md index 334a4701..0693d5e1 100644 --- a/docs/Distributed_SHTnsKit_Guide.md +++ b/docs/Distributed_SHTnsKit_Guide.md @@ -148,7 +148,7 @@ SHTnsKit.dist_synthesis!(spln, fθφ_out, PencilArray(Alm)) Enable plan caching across calls (optional) ```bash -export SHTNSKIT_CACHE_PENCILFFTS=1 +export SHTNSKIT_FFT_PLAN_CACHE=1 # legacy alias: SHTNSKIT_CACHE_PENCILFFTS ``` --- @@ -242,7 +242,10 @@ Both use the distributed transform paths internally and return gradients in the ``` - Robert form: for vector transforms, set `robert_form=true` in your config to stabilize polar behavior. - Normalization/phase: match `cfg.norm` and `cfg.cs_phase` to your data; conversions are handled internally on input/output. -- FFT plan caching: `ENV["SHTNSKIT_CACHE_PENCILFFTS"] = "1"` to reuse PencilFFTs plans. +- FFT plan caching: `ENV["SHTNSKIT_FFT_PLAN_CACHE"] = "1"` (default) to reuse the + per-rank φ-FFT plans. Equivalently `SHTnsKit.enable_fft_plan_cache!()` / + `disable_fft_plan_cache!()`. The cache is shared with the serial transforms; + cap it with `SHTnsKit.fft_plan_cache_max!(n)`. - Y-rotation strategy: truncated gather typically reduces bandwidth; switch to allgather for high-m–dominated spectra. --- diff --git a/docs/phi_scaling.md b/docs/phi_scaling.md index 35a9d19c..ce17ddd4 100644 --- a/docs/phi_scaling.md +++ b/docs/phi_scaling.md @@ -20,25 +20,32 @@ The `phi_scale` field in `SHTConfig` controls how the longitude (φ) dimension i - Regular grids with poles (simple trapezoidal rule) - **Rationale**: Adjusts for the φ integration measure `dφ` where ∫₀²ᵖ f dφ ≈ (2π/nlon) Σ f_j -### `:auto` - Automatic Selection -- **Behavior**: Chooses based on `grid_type` - - `:gauss` → `:dft` - - `:driscoll_healy` → `:dft` - - `:regular`, `:regular_poles` → `:quad` +### `:auto` - Unset +- **Behavior**: treated as `:dft`. +- Historically this keyed on `grid_type` and handed every non-Gauss grid `:quad`, + so a regular grid built through the exported `SHTConfig(; ...)` keyword + constructor (which defaulted to `:auto`) disagreed by 2π with the identical + grid from `create_regular_config`, which sets `:dft` explicitly. Both + constructors emit `:dft`, so an unset value now means `:dft` as well. ## Configuration ### In Code ```julia -# Explicit control -cfg = create_gauss_config(lmax, nlat; phi_scale=:dft) -cfg = create_regular_config(lmax, nlat; phi_scale=:quad) +# Every constructor emits :dft; there is no phi_scale keyword on them. +cfg = create_gauss_config(lmax, nlat) # :dft +cfg = create_regular_config(lmax, nlat) # :dft -# Automatic (recommended) -cfg = create_gauss_config(lmax, nlat) # Uses :dft -cfg = create_regular_config(lmax, nlat) # Uses :quad +# To opt into the quadrature convention, set it on the config: +cfg.phi_scale = :quad ``` +!!! note + An earlier version of this page showed `create_regular_config(lmax, nlat; + phi_scale=:quad)` and claimed regular grids default to `:quad`. Neither was + true: those constructors take no `phi_scale` keyword (that call raises) and + both set `:dft`. + ### Via Environment Variable ```bash # Override for all grids @@ -60,18 +67,24 @@ function phi_inv_scale(cfg::SHTConfig) return cfg.nlon end - # 2. Use config-specified mode + # 2. Use config-specified mode. `:dft` and anything unset (`:auto`) both + # mean the DFT convention — every constructor emits `:dft`, so an unset + # value is not a signal to guess from the grid type. if cfg.phi_scale === :quad return cfg.nlon / (2π) - elseif cfg.phi_scale === :dft - return cfg.nlon + else + return Float64(cfg.nlon) end - - # 3. Fall back to grid-type heuristic - return cfg.grid_type == :gauss ? cfg.nlon : cfg.nlon / (2π) end ``` +!!! note "The old grid-type fallback is gone" + An unset `phi_scale` used to fall back to `grid_type == :gauss ? nlon : nlon/2π`. + That made a regular grid built through the exported keyword `SHTConfig(...)` + constructor (whose `phi_scale` defaults to `:auto`) disagree by 2π with the + identical grid from `create_regular_config`, which sets `:dft` explicitly. + Unset now resolves to `:dft`, matching every constructor. + ## Why This Matters Incorrect φ-scaling leads to round-trip errors: @@ -90,3 +103,22 @@ The φ-scaling factor must match the quadrature weight convention to ensure: - Commit fc1d114: Regular grids changed to quadrature scaling (`nlon/(2π)`) - Commit 2441db0: Formalized with auto-detection - Current: Explicit `phi_scale` field for clarity and control + + +## Invariant + +Whichever mode is selected, `analysis` and `synthesis` are mutual inverses: + +```julia +analysis(cfg, synthesis(cfg, alm)) ≈ alm # exact under :dft and :quad +``` + +`synthesis` scales its Fourier bins by `phi_inv_scale(cfg)` and the inverse FFT +divides by `nlon`, a net spatial factor `σ`; `analysis` carries `cphi/σ` so the +two cancel. `:quad` therefore changes the scale of the *spatial* field (by 1/2π) +without changing what a round trip returns. Point and latitude evaluators +(`synthesis_point`, `SH_to_lat`, `SHqst_to_lat`, …) apply the same `σ`, so they +always agree with the grid `synthesis` produces. + +Before this was fixed, `analysis` ignored `phi_scale` entirely, so under `:quad` +a round trip returned `alm/2π` and every evaluator was 2π off from the grid. diff --git a/docs/src/grids.md b/docs/src/grids.md index e79474a3..34ad810c 100644 --- a/docs/src/grids.md +++ b/docs/src/grids.md @@ -37,6 +37,26 @@ Gauss–Legendre quadrature with `nlat` points integrates polynomials in quadrature is exact at the intended band limit when `nlat = 2*(lmax + 1)`; `nlat` must be even. +!!! warning "Equiangular grids need `nlat ≥ 2*lmax + 1`" + Fejér and Clenshaw–Curtis rules with `nlat` nodes are exact only through + degree ``n_\mathrm{lat} - 1``, and analysis integrates a product of two + degree-`lmax` Legendre functions — degree ``2\,l_\mathrm{max}``. So the + two regular grids reproduce `alm` from `analysis(synthesis(alm))` only from + `nlat = 2*lmax + 1` upward, where Gauss–Legendre needs just `nlat = lmax + 1`. + + Below the threshold nothing warns; the answer is simply wrong. Measured + relative round-trip error at `lmax = 8`: + + | `nlat` | `:gauss` | regular midpoint | regular with poles | + |---:|---:|---:|---:| + | 10 | 8e-16 | 7.2e-2 | 3.6e-1 | + | 14 | 8e-16 | 1.3e-2 | 7.1e-3 | + | 16 | 8e-16 | 4.2e-3 | 1.3e-3 | + | 17 (`2*lmax+1`) | 8e-16 | 7e-16 | 7e-16 | + + Size the grid deliberately, or use Gauss–Legendre (or Driscoll–Healy, exact + at `nlat = 2*(lmax+1)`) when round-trip accuracy matters. + !!! tip "Which grid should I choose?" Start with Gauss–Legendre unless you need to exchange data with a particular equiangular layout. Use regular midpoint for cell-centred data, regular with diff --git a/ext/ParallelLocal.jl b/ext/ParallelLocal.jl index 208cf6a5..d2d2a7ea 100644 --- a/ext/ParallelLocal.jl +++ b/ext/ParallelLocal.jl @@ -80,6 +80,8 @@ end `Q_p`/`S_p`/`T_p` use the coefficient convention configured by `cfg` (see [`dist_SH_to_lat`](@ref)). +When `cfg.robert_form` is enabled, the tangential components include the +`sin(θ)` factor used by full-grid vector synthesis; the radial component is unchanged. """ function SHTnsKit.dist_SHqst_to_point(cfg::SHTnsKit.SHTConfig, Q_p::PencilArray, S_p::PencilArray, T_p::PencilArray, cost::Real, phi::Real) return SHTnsKit.SHqst_to_point(cfg, Q_p, S_p, T_p, cost, phi) @@ -91,6 +93,8 @@ end `Q_p`/`S_p`/`T_p` use the coefficient convention configured by `cfg` (see [`dist_SH_to_lat`](@ref)). +When `cfg.robert_form` is enabled, the tangential components include the +`sin(θ)` factor used by full-grid vector synthesis; the radial component is unchanged. """ function SHTnsKit.dist_SHqst_to_lat(cfg::SHTnsKit.SHTConfig, Q_p::PencilArray, S_p::PencilArray, T_p::PencilArray, cost::Real; nphi::Int=cfg.nlon, ltr::Int=cfg.lmax, mtr::Int=cfg.mmax) @@ -313,6 +317,12 @@ function _pencil_local_qst(cfg, Q::PencilArray, S::PencilArray, Vt .*= sinth Vp .*= sinth end + # Same φ convention factor the serial evaluators apply (1 under :dft; + # 1/2π under :quad), so a PencilArray input agrees with the dense one. + sphi_scale = RT(SHTnsKit._evaluator_phi_scale(cfg)) + if sphi_scale != 1 + Vr .*= sphi_scale; Vt .*= sphi_scale; Vp .*= sphi_scale + end combined = vcat(Vr, Vt, Vp) _record_local_payload!(length(combined)) MPI.Allreduce!(combined, +, comm) @@ -1135,7 +1145,7 @@ function _analysis_mode_pencil(cfg::SHTnsKit.SHTConfig, im::Int, RT = typeof(real(zero(CT))) P = Vector{Float64}(undef, ltr + 1) rank = MPI.Comm_rank(comm) - phi_scale = axisymmetric ? cfg.cphi * cfg.nlon : cfg.cphi + phi_scale = axisymmetric ? SHTnsKit._analysis_phi_scale(cfg) * cfg.nlon : SHTnsKit._analysis_phi_scale(cfg) for root in 0:(MPI.Comm_size(comm) - 1) send = zeros(CT, counts[root + 1]) @inbounds for (i, θindex) in pairs(θglobals) diff --git a/ext/ParallelTransforms.jl b/ext/ParallelTransforms.jl index fe160b9b..a9b32dc7 100644 --- a/ext/ParallelTransforms.jl +++ b/ext/ParallelTransforms.jl @@ -1100,7 +1100,7 @@ function _analysis_owned_block!(block::AbstractMatrix{CT}, cfg::SHTnsKit.SHTConf fill!(block, zero(CT)) P = Vector{Float64}(undef, cfg.lmax + 1) RT = typeof(real(zero(CT))) - cphi = RT(cfg.cphi) + cphi = RT(SHTnsKit._analysis_phi_scale(cfg)) use_tbl = cfg.use_plm_tables && !isempty(cfg.NP_tables) @inbounds for (local_m, m_index) in pairs(m_indices) m = m_index - 1 @@ -1374,7 +1374,7 @@ Decompose latitude instead: `SHTnsKit.create_spatial_pencil(cfg; comm)` or `Penc if use_packed_storage # Original inline loop for packed storage (not the hot path). # Uses normalized rows (Plm_norm_row! / NP_tables): Nlm is already baked in. - xv = cfg.x; cphi = cfg.cphi # hoist field reads out of the loops below (cfg is mutable, so not auto-hoisted) + xv = cfg.x; cphi = SHTnsKit._analysis_phi_scale(cfg) # hoisted; inverts synthesis under any phi_scale # Stride by mres: create_packed_storage_info only assigns lm_to_packed for # m % mres == 0, leaving every other entry 0. Walking all m under @inbounds # therefore wrote Alm_local[0] — one element before the buffer. @@ -1441,7 +1441,7 @@ Decompose latitude instead: `SHTnsKit.create_spatial_pencil(cfg; comm)` or `Penc if !use_packed_storage # Apply φ scaling (cphi = 2π/nlon). Nlm is NOT applied here: the # normalized recurrence Plm_norm_row! already bakes Nlm into P̄. - cphi = cfg.cphi # hoist field read out of the normalization loop (cfg is mutable) + cphi = SHTnsKit._analysis_phi_scale(cfg) # hoisted; inverts synthesis under any phi_scale @inbounds for m in 0:cfg.mres:mmax @simd ivdep for l in m:lmax Alm_local[l+1, m+1] *= cphi @@ -1451,7 +1451,7 @@ Decompose latitude instead: `SHTnsKit.create_spatial_pencil(cfg; comm)` or `Penc else # θ is not distributed - no reduction needed, just apply φ scaling if !use_packed_storage - cphi = cfg.cphi # hoist field read out of the normalization loop (cfg is mutable) + cphi = SHTnsKit._analysis_phi_scale(cfg) # hoisted; inverts synthesis under any phi_scale @inbounds for m in 0:cfg.mres:mmax @simd ivdep for l in m:lmax Alm_local[l+1, m+1] *= cphi @@ -1515,7 +1515,7 @@ function _dist_analysis_plan_core!(plan::DistAnalysisPlan, end # φ scaling (Nlm already baked into the normalized Legendre rows) - cphi = cfg.cphi + cphi = SHTnsKit._analysis_phi_scale(cfg) # inverts synthesis under any phi_scale (= cphi under :dft) @inbounds for m in 0:cfg.mres:mmax @simd ivdep for l in m:lmax plan.Alm_work[l+1, m+1] *= cphi @@ -2054,7 +2054,7 @@ function _analysis_sphtor_owned_block!(Sblock::AbstractMatrix{CT}, for l in max(1, m):lcap d = RT(dtheta[l + 1]) term = complex(zero(RT), RT(m * over_sin[l + 1])) - factor = wi * RT(cfg.cphi) / RT(l * (l + 1)) + factor = wi * RT(SHTnsKit._analysis_phi_scale(cfg)) / RT(l * (l + 1)) Sblock[l + 1, local_m] += factor * (Ft * d + conj(term) * Fp) Tblock[l + 1, local_m] += @@ -2383,7 +2383,7 @@ function _analysis_sphtor_mode_pencil(cfg::SHTnsKit.SHTConfig, l >= max(1, physical_m) || continue derivative = RT(dtheta[l + 1]) term = complex(zero(RT), RT(physical_m * over_sin[l + 1])) - coefficient = weight * RT(cfg.cphi) / RT(l * (l + 1)) + coefficient = weight * RT(SHTnsKit._analysis_phi_scale(cfg)) / RT(l * (l + 1)) Ssend[k] += coefficient * (Ftheta * derivative + conj(term) * Fphi) Tsend[k] += coefficient * @@ -2571,7 +2571,7 @@ function _legacy_dist_analysis_sphtor(cfg::SHTnsKit.SHTConfig, dPdtheta = Vector{Float64}(undef, lmax + 1) P_over_sth = Vector{Float64}(undef, lmax + 1) Pbuf = Vector{Float64}(undef, lmax + 2) # scratch for normalized dθ recurrence - scaleφ = cfg.cphi + scaleφ = SHTnsKit._analysis_phi_scale(cfg) # inverts synthesis under any phi_scale (= cphi under :dft) # Main vector analysis loop — via the same function barriers the planned # path uses (concrete argument types; the previous inline loop boxed its @@ -2763,12 +2763,12 @@ function _dist_analysis_sphtor_plan_core!( _sphtor_analysis_loop_tbl!(cfg, plan.Slm_work, plan.Tlm_work, cfg.NP_tables, cfg.NdP_tables, plan.Ftθm, plan.Fpθm, plan.θ_globals, plan.sθ_cache, plan.weights_cache, - cfg.robert_form, cfg.cphi, lmax, mmax) + cfg.robert_form, SHTnsKit._analysis_phi_scale(cfg), lmax, mmax) else _sphtor_analysis_loop_otf!(plan.Slm_work, plan.Tlm_work, plan.P, plan.dPdtheta, plan.P_over_sth, plan.Pbuf, plan.Ftθm, plan.Fpθm, plan.x_cache, plan.sθ_cache, plan.inv_sθ_cache, - plan.weights_cache, cfg.robert_form, cfg.cphi, + plan.weights_cache, cfg.robert_form, SHTnsKit._analysis_phi_scale(cfg), cfg.mres, lmax, mmax) end @@ -3930,7 +3930,7 @@ function dist_analysis_distributed(cfg::SHTnsKit.SHTConfig, fθφ::PencilArray; # Compute local contributions to ALL coefficients (same as standard analysis) local_contrib = zeros(ComplexF64, lmax + 1, mmax + 1) - scaleφ = cfg.cphi + scaleφ = SHTnsKit._analysis_phi_scale(cfg) # inverts synthesis under any phi_scale (= cphi under :dft) # Pre-cache weights weights_cache = Vector{Float64}(undef, nθ_local) @@ -4177,6 +4177,19 @@ end return false end +""" + close(plan::DistributedSpectralPlan2D) + +Free the plan's `l_comm` / `m_comm` sub-communicators. + +!!! warning "Collective" + `MPI_Comm_free` is a collective operation: **every** rank holding this plan + must call `close` on it, and in the same order relative to other collectives. + It is therefore not attached to a finalizer — garbage collection is + rank-local and would call it at divergent points. Closing twice is safe + (the second call is a no-op); never closing only leaks the communicators + until `MPI_Finalize`. +""" function Base.close(plan::DistributedSpectralPlan2D) plan.closed && return nothing plan.closed = true @@ -4520,12 +4533,21 @@ function create_distributed_spectral_plan_2d(lmax::Int, mmax::Int, comm::MPI.Com m_group_nlm, with_scratch, scratch, scratch_context, false ) - finalizer(plan) do p - try - close(p) - catch - end - end + # NO finalizer here, deliberately. + # + # `close(plan)` frees the `l_comm` / `m_comm` sub-communicators, and + # `MPI_Comm_free` is a COLLECTIVE: every rank in the communicator must + # call it, in the same order relative to other collectives. A finalizer + # runs whenever that rank's garbage collector happens to fire, which is + # rank-local and nondeterministic — so attaching one meant ranks could + # enter `Comm_free` at completely different points in the program, or + # some not at all before the next collective. That is a hang or worse, + # and it contradicts the policy `ParallelPlans.jl` states for its own + # `Comm_split` bookkeeping. + # + # Leaking two communicators until `MPI_Finalize` is strictly better than + # a nondeterministic collective, so cleanup is explicit only: call + # `close(plan)` (collectively) when you are done with it. return plan catch _safe_comm_free(l_comm) @@ -4880,7 +4902,7 @@ function _dist_analysis_2d_safe(cfg::SHTnsKit.SHTConfig, fθφ::PencilArray; x_cache[ii] = cfg.x[iglob] end - scaleφ = cfg.cphi + scaleφ = SHTnsKit._analysis_phi_scale(cfg) # inverts synthesis under any phi_scale (= cphi under :dft) # Use NP_tables (already normalized P̄) if available; fall back to OTF normalized rows. use_tbl = use_tables && cfg.use_plm_tables && !isempty(cfg.NP_tables) P = Vector{Float64}(undef, lmax + 1) @@ -5388,7 +5410,7 @@ function _dist_analysis_2d_aligned(cfg::SHTnsKit.SHTConfig, fθφ::PencilArray; copyto!(Fθm, Fθm_temp) end - scaleφ = cfg.cphi + scaleφ = SHTnsKit._analysis_phi_scale(cfg) # inverts synthesis under any phi_scale (= cphi under :dft) # Use NP_tables (already normalized P̄) if available; fall back to OTF normalized rows. use_tbl = use_tables && cfg.use_plm_tables && !isempty(cfg.NP_tables) diff --git a/ext/ParallelTransposeTransforms.jl b/ext/ParallelTransposeTransforms.jl index 528670bc..17d11ce6 100644 --- a/ext/ParallelTransposeTransforms.jl +++ b/ext/ParallelTransposeTransforms.jl @@ -515,7 +515,7 @@ function _dist_transpose_analysis_unchecked!( fill!(A, zero(eltype(A))) w = plan.cfg.w - scaleφ = plan.cfg.cphi # 2π/nlon — converts unnormalized rFFT sum to integral + scaleφ = SHTnsKit._analysis_phi_scale(plan.cfg) # 2π/nlon under :dft; inverts synthesis under any phi_scale lmax = plan.lmax nlat = plan.nlat nlev = plan.nlev @@ -657,7 +657,7 @@ function _dist_transpose_analysis_sphtor_unchecked!( fill!(T, zero(eltype(T))) w = plan.cfg.w - scaleφ = plan.cfg.cphi # 2π/nlon — converts unnormalized rFFT sum to integral + scaleφ = SHTnsKit._analysis_phi_scale(plan.cfg) # 2π/nlon under :dft; inverts synthesis under any phi_scale lmax = plan.lmax nlat = plan.nlat nlev = plan.nlev diff --git a/ext/SHTnsKitAMDGPUExt.jl b/ext/SHTnsKitAMDGPUExt.jl index 7840894a..f16d9ce5 100644 --- a/ext/SHTnsKitAMDGPUExt.jl +++ b/ext/SHTnsKitAMDGPUExt.jl @@ -627,7 +627,7 @@ function _amdgpu_scalar_analysis_direct!(owner, cfg::SHTConfig, end backend = ROCBackend() scalar_analysis_kernel!(backend)( - output, fourier, tables.Plm, tables.weights, RT(cfg.cphi), + output, fourier, tables.Plm, tables.weights, RT(SHTnsKit._analysis_phi_scale(cfg)), cfg.lmax, cfg.mmax, cfg.mres, cfg.lmax; ndrange=(cfg.lmax + 1, cfg.mmax + 1), ) @@ -716,7 +716,7 @@ function _amdgpu_scalar_analysis(cfg::SHTConfig, field::AMDGPU.AnyROCArray; backend = ROCBackend() canonical = AMDGPU.zeros(CT, cfg.lmax + 1, cfg.mmax + 1) analyze! = scalar_analysis_kernel!(backend) - analyze!(canonical, fourier, tables.Plm, tables.weights, RT(cfg.cphi), + analyze!(canonical, fourier, tables.Plm, tables.weights, RT(SHTnsKit._analysis_phi_scale(cfg)), cfg.lmax, cfg.mmax, cfg.mres, lcap; ndrange=(lcap + 1, min(cfg.mmax, lcap) + 1)) @@ -849,7 +849,7 @@ function _amdgpu_vector_analysis_direct!(owner, cfg::SHTConfig, vector_analysis_kernel!(ROCBackend())( Sout, Tout, workspace.Ftheta, workspace.Fphi, tables.dtheta, tables.over_sin, tables.weights, tables.scales, - tables.x, RT(cfg.cphi), lcap, min(cfg.mmax, lcap), cfg.mres, + tables.x, RT(SHTnsKit._analysis_phi_scale(cfg)), lcap, min(cfg.mmax, lcap), cfg.mres, cfg.robert_form; ndrange=(lcap + 1, min(cfg.mmax, lcap) + 1), ) AMDGPU.synchronize() @@ -1044,7 +1044,7 @@ function _amdgpu_vector_mode_analysis(cfg::SHTConfig, stored_im::Integer, S = AMDGPU.zeros(CT, lcap - physical_m + 1); Tlm = similar(S) vector_mode_analysis_kernel!(ROCBackend())( S, Tlm, Vt, Vp, tables.dtheta, tables.over_sin, tables.weights, - tables.scales, tables.x, RT(cfg.cphi), physical_m, lcap, + tables.scales, tables.x, RT(SHTnsKit._analysis_phi_scale(cfg)), physical_m, lcap, cfg.robert_form; ndrange=length(S), ) AMDGPU.synchronize() @@ -1181,7 +1181,7 @@ function _amdgpu_vector_batch_analysis(cfg::SHTConfig, S = AMDGPU.zeros(CT, cfg.lmax + 1, cfg.mmax + 1, nfields); Tlm = similar(S) vector_batch_analysis_kernel!(ROCBackend())( S, Tlm, Ft, Fp, tables.dtheta, tables.over_sin, tables.weights, - tables.scales, tables.x, RT(cfg.cphi), cfg.lmax, cfg.mmax, + tables.scales, tables.x, RT(SHTnsKit._analysis_phi_scale(cfg)), cfg.lmax, cfg.mmax, cfg.mres, cfg.robert_form; ndrange=size(S), ) AMDGPU.synchronize() @@ -1537,7 +1537,7 @@ function analysis_axisym(::SHTnsKit.GPU, cfg::SHTConfig, length(field) == cfg.nlat || throw(DimensionMismatch( "field must have length nlat=$(cfg.nlat)", )) - return _amdgpu_mode_analysis(cfg, 0, field, cfg.lmax, cfg.cphi * cfg.nlon) + return _amdgpu_mode_analysis(cfg, 0, field, cfg.lmax, SHTnsKit._analysis_phi_scale(cfg) * cfg.nlon) end analysis_axisym(cfg::SHTConfig, field::AMDGPU.AnyROCArray{T,1}) where {T<:Real} = analysis_axisym(SHTnsKit.GPU(), cfg, field) @@ -1559,7 +1559,7 @@ function analysis_axisym_l(::SHTnsKit.GPU, cfg::SHTConfig, length(field) == cfg.nlat || throw(DimensionMismatch( "field must have length nlat=$(cfg.nlat)", )) - return _amdgpu_mode_analysis(cfg, 0, field, lcap, cfg.cphi * cfg.nlon) + return _amdgpu_mode_analysis(cfg, 0, field, lcap, SHTnsKit._analysis_phi_scale(cfg) * cfg.nlon) end analysis_axisym_l(cfg::SHTConfig, field::AMDGPU.AnyROCArray{T,1}, ltr::Integer) where {T<:Real} = @@ -1599,7 +1599,7 @@ function analysis_packed_ml(::SHTnsKit.GPU, cfg::SHTConfig, im::Int, length(mode) == cfg.nlat || throw(DimensionMismatch( "mode must have length nlat=$(cfg.nlat)", )) - return _amdgpu_mode_analysis(cfg, physical_m, mode, lcap, cfg.cphi) + return _amdgpu_mode_analysis(cfg, physical_m, mode, lcap, SHTnsKit._analysis_phi_scale(cfg)) end analysis_packed_ml(cfg::SHTConfig, im::Int, mode::AMDGPU.AnyROCArray{T,1}, ltr::Integer) where {T<:Complex} = @@ -1637,7 +1637,7 @@ function _amdgpu_analysis_packed_cplx(cfg::SHTConfig, mcap = min(cfg.mmax, lcap) kernel! = complex_packed_analysis_kernel!(ROCBackend()) kernel!(packed, fourier, tables.Plm, tables.weights, tables.scales, - RT(cfg.cphi), cfg.nlon, lcap, cfg.mmax, mcap; + RT(SHTnsKit._analysis_phi_scale(cfg)), cfg.nlon, lcap, cfg.mmax, mcap; ndrange=(lcap + 1, 2mcap + 1)) AMDGPU.synchronize() return packed @@ -1723,7 +1723,7 @@ function _amdgpu_batch_analysis(cfg::SHTConfig, fields::AMDGPU.AnyROCArray; FFTW.fft!(fourier, 2) canonical = AMDGPU.zeros(CT, cfg.lmax + 1, cfg.mmax + 1, size(fields, 3)) kernel! = scalar_batch_analysis_kernel!(ROCBackend()) - kernel!(canonical, fourier, tables.Plm, tables.weights, RT(cfg.cphi), + kernel!(canonical, fourier, tables.Plm, tables.weights, RT(SHTnsKit._analysis_phi_scale(cfg)), cfg.lmax, cfg.mmax, cfg.mres; ndrange=size(canonical)) configured = canonical ./ reshape( tables.scales, cfg.lmax + 1, cfg.mmax + 1, 1, @@ -1765,7 +1765,7 @@ function _amdgpu_batch_analysis_direct!(cfg::SHTConfig, end backend = ROCBackend() scalar_batch_analysis_kernel!(backend)( - output, fourier, tables.Plm, tables.weights, RT(cfg.cphi), + output, fourier, tables.Plm, tables.weights, RT(SHTnsKit._analysis_phi_scale(cfg)), cfg.lmax, cfg.mmax, cfg.mres; ndrange=size(output), ) coefficient_batch_conversion_kernel!(backend)( diff --git a/ext/SHTnsKitAdvancedADExt.jl b/ext/SHTnsKitAdvancedADExt.jl index a9460050..674b5319 100644 --- a/ext/SHTnsKitAdvancedADExt.jl +++ b/ext/SHTnsKitAdvancedADExt.jl @@ -42,8 +42,13 @@ import SHTnsKit: wigner_d_matrix_deriv A isa ChainRulesCore.AbstractZero ? _coeff_zeros(cfg) : _to_complex(A) - function ChainRulesCore.rrule(::typeof(SHTnsKit.analysis), cfg::SHTnsKit.SHTConfig, f) - y = SHTnsKit.analysis(cfg, f) + # `fft_scratch` / `use_rfft` pick a different FFT implementation of the SAME + # linear operator, so the adjoint is unchanged — but a pullback must still + # ACCEPT them. Declaring fewer kwargs than the primal made ChainRules skip + # this rule entirely the moment a caller passed one, even at its default. + function ChainRulesCore.rrule(::typeof(SHTnsKit.analysis), cfg::SHTnsKit.SHTConfig, f; + fft_scratch=nothing, use_rfft::Bool=false) + y = SHTnsKit.analysis(cfg, f; fft_scratch, use_rfft) project_f = ProjectTo(f) function pullback(ȳ) ȳA = _to_complex(ȳ) @@ -61,8 +66,9 @@ import SHTnsKit: wigner_d_matrix_deriv # `_adjoint_synthesis` helper instead. (See `test_adjoint_consistency` # in the test suite for an FD verification.) function ChainRulesCore.rrule(::typeof(SHTnsKit.synthesis), cfg::SHTnsKit.SHTConfig, - alm; real_output::Bool=true) - y = SHTnsKit.synthesis(cfg, alm; real_output) + alm; real_output::Bool=true, + fft_scratch=nothing, use_rfft::Bool=false) + y = SHTnsKit.synthesis(cfg, alm; real_output, fft_scratch, use_rfft) project_alm = ProjectTo(alm) function pullback(ȳ) ȳ_mat = ChainRulesCore.unthunk(ȳ) # materialize Thunk/InplaceableThunk @@ -171,8 +177,9 @@ import SHTnsKit: wigner_d_matrix_deriv # Keep local alias for any direct callers of the ext symbol. const _adjoint_analysis_sphtor = SHTnsKit._adjoint_analysis_sphtor - function ChainRulesCore.rrule(::typeof(SHTnsKit.analysis_sphtor), cfg::SHTnsKit.SHTConfig, Vt, Vp) - Slm, Tlm = SHTnsKit.analysis_sphtor(cfg, Vt, Vp) + function ChainRulesCore.rrule(::typeof(SHTnsKit.analysis_sphtor), cfg::SHTnsKit.SHTConfig, Vt, Vp; + use_rfft::Bool=false) + Slm, Tlm = SHTnsKit.analysis_sphtor(cfg, Vt, Vp; use_rfft) project_Vt = ProjectTo(Vt) project_Vp = ProjectTo(Vp) function pullback(ṠTl) @@ -194,8 +201,8 @@ import SHTnsKit: wigner_d_matrix_deriv const _adjoint_synthesis_sphtor = SHTnsKit._adjoint_synthesis_sphtor function ChainRulesCore.rrule(::typeof(SHTnsKit.synthesis_sphtor), cfg::SHTnsKit.SHTConfig, - Slm, Tlm; real_output::Bool=true) - Vt, Vp = SHTnsKit.synthesis_sphtor(cfg, Slm, Tlm; real_output) + Slm, Tlm; real_output::Bool=true, use_rfft::Bool=false) + Vt, Vp = SHTnsKit.synthesis_sphtor(cfg, Slm, Tlm; real_output, use_rfft) project_Slm = ProjectTo(Slm) project_Tlm = ProjectTo(Tlm) function pullback(Ṽ) @@ -214,6 +221,45 @@ import SHTnsKit: wigner_d_matrix_deriv return (Vt, Vp), pullback end + # QST (3-component) transforms. `synthesis_qst` is the scalar synthesis of Q + # alongside the sphtor synthesis of (S,T), and `analysis_qst` is the mirror, + # so each adjoint is just the two existing adjoints side by side. Without + # these, differentiating a QST pipeline fell through to Zygote's source + # tracing and crashed inside FFTW. + function ChainRulesCore.rrule(::typeof(SHTnsKit.synthesis_qst), cfg::SHTnsKit.SHTConfig, + Qlm, Slm, Tlm; real_output::Bool=true, + use_rfft::Bool=false) + Vr, Vt, Vp = SHTnsKit.synthesis_qst(cfg, Qlm, Slm, Tlm; real_output, use_rfft) + project_Q = ProjectTo(Qlm); project_S = ProjectTo(Slm); project_T = ProjectTo(Tlm) + function pullback(V̄) + zsp() = zeros(Float64, cfg.nlat, cfg.nlon) + V̄r = ChainRulesCore.unthunk(V̄[1]); V̄r = V̄r isa ChainRulesCore.AbstractZero ? zsp() : V̄r + V̄t = ChainRulesCore.unthunk(V̄[2]); V̄t = V̄t isa ChainRulesCore.AbstractZero ? zsp() : V̄t + V̄p = ChainRulesCore.unthunk(V̄[3]); V̄p = V̄p isa ChainRulesCore.AbstractZero ? zsp() : V̄p + Q̄ = SHTnsKit._adjoint_synthesis(cfg, V̄r; real_output=real_output) + S̄, T̄ = SHTnsKit._adjoint_synthesis_sphtor(cfg, V̄t, V̄p; real_output=real_output) + return NoTangent(), NoTangent(), project_Q(Q̄), project_S(S̄), project_T(T̄), + (; real_output=NoTangent(), use_rfft=NoTangent()) + end + return (Vr, Vt, Vp), pullback + end + + function ChainRulesCore.rrule(::typeof(SHTnsKit.analysis_qst), cfg::SHTnsKit.SHTConfig, + Vr, Vt, Vp; use_rfft::Bool=false) + Qlm, Slm, Tlm = SHTnsKit.analysis_qst(cfg, Vr, Vt, Vp; use_rfft) + project_Vr = ProjectTo(Vr); project_Vt = ProjectTo(Vt); project_Vp = ProjectTo(Vp) + function pullback(Ā) + Q̄ = _materialize_coeff(Ā[1], cfg) + S̄ = _materialize_coeff(Ā[2], cfg) + T̄ = _materialize_coeff(Ā[3], cfg) + V̄r = _adjoint_analysis(cfg, Q̄) + V̄t, V̄p = _adjoint_analysis_sphtor(cfg, S̄, T̄) + return NoTangent(), NoTangent(), project_Vr(V̄r), project_Vt(V̄t), project_Vp(V̄p), + (; use_rfft=NoTangent()) + end + return (Qlm, Slm, Tlm), pullback + end + # Complex packed (LM_cplx layout — both signs of m stored explicitly). # # Both transforms are ℂ-linear in their argument (no `real()` anywhere), so @@ -265,7 +311,7 @@ import SHTnsKit: wigner_d_matrix_deriv CT = complex(float(eltype(ā_int))) F̄ = zeros(CT, nlat, nlon) P = Vector{Float64}(undef, lmax + 1) - scaleφ = cfg.cphi + scaleφ = SHTnsKit._analysis_phi_scale(cfg) xv = cfg.x; wv = cfg.w for am in 0:mmax colp = am + 1 @@ -370,16 +416,20 @@ function ChainRulesCore.rrule(::typeof(SHTnsKit.SH_Zrotate), cfg::SHTnsKit.SHTCo end function ChainRulesCore.rrule(::typeof(SHTnsKit.SH_Yrotate), cfg::SHTnsKit.SHTConfig, Qlm, alpha::Real, Rlm) + # Snapshot BEFORE the primal: an in-place rotation (Rlm === Qlm) overwrites + # Qlm, and callers may reuse either buffer before the pullback runs. The dα + # formula needs the *input* coefficients, so they must be preserved here. + Qlm_saved = copy(Qlm) y = SHTnsKit.SH_Yrotate(cfg, Qlm, alpha, Rlm) function pullback(ȳ) inverse = SHTnsKit.SHTRotation(cfg.lmax, cfg.mmax) SHTnsKit.shtns_rotation_set_angles_ZYZ(inverse, 0.0, -alpha, 0.0) - Q̄ = similar(Qlm) + Q̄ = similar(Qlm_saved) _configured_rotation_adjoint!(cfg, inverse, ȳ, Q̄) # angle gradient via d/dβ of Wigner-d at β=alpha dα = zero(float(alpha)) lmax, mmax = cfg.lmax, cfg.mmax - Qlm_canonical = SHTnsKit._internal_coefficients(Qlm, cfg) + Qlm_canonical = SHTnsKit._internal_coefficients(Qlm_saved, cfg) ȳ_canonical = SHTnsKit._analysis_cotangent_to_canonical(ȳ, cfg) for l in 0:lmax mm = min(l, mmax) @@ -439,10 +489,14 @@ end # Adjoint for complex rotation using conjugate-transpose of Wigner-D function ChainRulesCore.rrule(::typeof(SHTnsKit.shtns_rotation_apply_cplx), r::SHTnsKit.SHTRotation, Zlm, Rlm) + # Snapshot BEFORE the primal — the per-l blocks are read then written, so an + # in-place call (Rlm === Zlm) is legal and overwrites the input the angle + # gradients below depend on. See the same fix on SH_Zrotate / SH_Yrotate. + Zlm_saved = copy(Zlm) y = SHTnsKit.shtns_rotation_apply_cplx(r, Zlm, Rlm) function pullback(ȳ) lmax, mmax = r.lmax, r.mmax - Z̄ = similar(Zlm) + Z̄ = similar(Zlm_saved) fill!(Z̄, zero(eltype(Z̄))) α, β, γ = _rotation_rrule_angles(r) # This pullback reimplements the Wigner engine, which works in the Y_l^m @@ -457,7 +511,7 @@ function ChainRulesCore.rrule(::typeof(SHTnsKit.shtns_rotation_apply_cplx), r::S # `ȳ` is the closure's own argument and would be safe, but is renamed too # so the pair reads the same way. ε = SHTnsKit._lmcplx_ybasis_signs(lmax, mmax) - Zε = ε .* Zlm + Zε = ε .* Zlm_saved ȳε = ε .* ȳ gα = 0.0; gβ = 0.0; gγ = 0.0 for l in 0:lmax @@ -533,11 +587,16 @@ end # Adjoint for real packed rotation: extend to full, apply cplx adjoint, fold back function ChainRulesCore.rrule(::typeof(SHTnsKit.shtns_rotation_apply_real), r::SHTnsKit.SHTRotation, Qlm, Rlm) + # Snapshot BEFORE the primal: the input is fully expanded into a scratch + # LM_cplx buffer before anything is written back, so an in-place call + # (Rlm === Qlm) is legal and clobbers the coefficients the angle gradients + # below reconstruct `b` from. Same fix as SH_Zrotate / SH_Yrotate. + Qlm_saved = copy(Qlm) y = SHTnsKit.shtns_rotation_apply_real(r, Qlm, Rlm) function pullback(ȳ) lmax, mmax = r.lmax, r.mmax # Extend cotangent on packed to full complex - Zbar_full = zeros(eltype(Qlm), SHTnsKit.nlm_cplx_calc(lmax, mmax, 1)) + Zbar_full = zeros(eltype(Qlm_saved), SHTnsKit.nlm_cplx_calc(lmax, mmax, 1)) for l in 0:lmax mm = min(l, mmax) # m = 0 @@ -588,11 +647,11 @@ function ChainRulesCore.rrule(::typeof(SHTnsKit.shtns_rotation_apply_real), r::S b[mp + l + 1] = (SHTnsKit.LM_cplx_index(lmax, mmax, l, mp) >= 0) ? (begin # reconstruct from packed Qlm if mp == 0 - Qlm[SHTnsKit.LM_index(lmax, 1, l, 0) + 1] + Qlm_saved[SHTnsKit.LM_index(lmax, 1, l, 0) + 1] elseif mp > 0 - Qlm[SHTnsKit.LM_index(lmax, 1, l, mp) + 1] + Qlm_saved[SHTnsKit.LM_index(lmax, 1, l, mp) + 1] else - (-1)^(-mp) * conj(Qlm[SHTnsKit.LM_index(lmax, 1, l, -mp) + 1]) + (-1)^(-mp) * conj(Qlm_saved[SHTnsKit.LM_index(lmax, 1, l, -mp) + 1]) end end) : 0 b[mp + l + 1] *= cis(-mp * γ) @@ -614,7 +673,7 @@ function ChainRulesCore.rrule(::typeof(SHTnsKit.shtns_rotation_apply_real), r::S end end # Fold back to packed positive-m: q̄(m) = Z̄(m) + (-1)^m conj(Z̄(-m)) - Q̄ = zeros(eltype(Qlm), length(Qlm)) + Q̄ = zeros(eltype(Qlm_saved), length(Qlm_saved)) for l in 0:lmax mm = min(l, mmax) # m=0 diff --git a/ext/SHTnsKitGPUExt.jl b/ext/SHTnsKitGPUExt.jl index a2151294..4c0621bc 100644 --- a/ext/SHTnsKitGPUExt.jl +++ b/ext/SHTnsKitGPUExt.jl @@ -628,7 +628,7 @@ function _cuda_scalar_analysis_direct!(owner, cfg::SHTConfig, end backend = CUDABackend() scalar_analysis_kernel!(backend)( - output, fourier, tables.Plm, tables.weights, RT(cfg.cphi), + output, fourier, tables.Plm, tables.weights, RT(SHTnsKit._analysis_phi_scale(cfg)), cfg.lmax, cfg.mmax, cfg.mres, cfg.lmax; ndrange=(cfg.lmax + 1, cfg.mmax + 1), ) @@ -716,7 +716,7 @@ function _cuda_scalar_analysis(cfg::SHTConfig, field::CUDA.AnyCuArray; backend = CUDABackend() canonical = CUDA.zeros(CT, cfg.lmax + 1, cfg.mmax + 1) analyze! = scalar_analysis_kernel!(backend) - analyze!(canonical, fourier, tables.Plm, tables.weights, RT(cfg.cphi), + analyze!(canonical, fourier, tables.Plm, tables.weights, RT(SHTnsKit._analysis_phi_scale(cfg)), cfg.lmax, cfg.mmax, cfg.mres, lcap; ndrange=(lcap + 1, min(cfg.mmax, lcap) + 1)) @@ -846,7 +846,7 @@ function _cuda_vector_analysis_direct!(owner, cfg::SHTConfig, vector_analysis_kernel!(CUDABackend())( Sout, Tout, workspace.Ftheta, workspace.Fphi, tables.dtheta, tables.over_sin, tables.weights, tables.scales, - tables.x, RT(cfg.cphi), lcap, min(cfg.mmax, lcap), cfg.mres, + tables.x, RT(SHTnsKit._analysis_phi_scale(cfg)), lcap, min(cfg.mmax, lcap), cfg.mres, cfg.robert_form; ndrange=(lcap + 1, min(cfg.mmax, lcap) + 1), ) CUDA.synchronize() @@ -1052,7 +1052,7 @@ function _cuda_vector_mode_analysis(cfg::SHTConfig, stored_im::Integer, Tout = similar(Sout) vector_mode_analysis_kernel!(CUDABackend())( Sout, Tout, Vt, Vp, tables.dtheta, tables.over_sin, - tables.weights, tables.scales, tables.x, RT(cfg.cphi), physical_m, + tables.weights, tables.scales, tables.x, RT(SHTnsKit._analysis_phi_scale(cfg)), physical_m, lcap, cfg.robert_form; ndrange=length(Sout), ) CUDA.synchronize() @@ -1196,7 +1196,7 @@ function _cuda_vector_batch_analysis(cfg::SHTConfig, Tout = similar(Sout) vector_batch_analysis_kernel!(CUDABackend())( Sout, Tout, Ft, Fp, tables.dtheta, tables.over_sin, - tables.weights, tables.scales, tables.x, RT(cfg.cphi), cfg.lmax, + tables.weights, tables.scales, tables.x, RT(SHTnsKit._analysis_phi_scale(cfg)), cfg.lmax, cfg.mmax, cfg.mres, cfg.robert_form; ndrange=size(Sout), ) CUDA.synchronize() @@ -1532,7 +1532,7 @@ function analysis_axisym(::SHTnsKit.GPU, cfg::SHTConfig, length(field) == cfg.nlat || throw(DimensionMismatch( "field must have length nlat=$(cfg.nlat)", )) - return _cuda_mode_analysis(cfg, 0, field, cfg.lmax, cfg.cphi * cfg.nlon) + return _cuda_mode_analysis(cfg, 0, field, cfg.lmax, SHTnsKit._analysis_phi_scale(cfg) * cfg.nlon) end analysis_axisym(cfg::SHTConfig, field::CUDA.AnyCuArray{T,1}) where {T<:Real} = analysis_axisym(SHTnsKit.GPU(), cfg, field) @@ -1553,7 +1553,7 @@ function analysis_axisym_l(::SHTnsKit.GPU, cfg::SHTConfig, length(field) == cfg.nlat || throw(DimensionMismatch( "field must have length nlat=$(cfg.nlat)", )) - return _cuda_mode_analysis(cfg, 0, field, lcap, cfg.cphi * cfg.nlon) + return _cuda_mode_analysis(cfg, 0, field, lcap, SHTnsKit._analysis_phi_scale(cfg) * cfg.nlon) end analysis_axisym_l(cfg::SHTConfig, field::CUDA.AnyCuArray{T,1}, ltr::Integer) where {T<:Real} = analysis_axisym_l(SHTnsKit.GPU(), cfg, field, ltr) @@ -1591,7 +1591,7 @@ function analysis_packed_ml(::SHTnsKit.GPU, cfg::SHTConfig, im::Int, length(mode) == cfg.nlat || throw(DimensionMismatch( "mode must have length nlat=$(cfg.nlat)", )) - return _cuda_mode_analysis(cfg, physical_m, mode, lcap, cfg.cphi) + return _cuda_mode_analysis(cfg, physical_m, mode, lcap, SHTnsKit._analysis_phi_scale(cfg)) end analysis_packed_ml(cfg::SHTConfig, im::Int, mode::CUDA.AnyCuArray{T,1}, ltr::Integer) where {T<:Complex} = @@ -1628,7 +1628,7 @@ function _cuda_analysis_packed_cplx(cfg::SHTConfig, mcap = min(cfg.mmax, lcap) kernel! = complex_packed_analysis_kernel!(CUDABackend()) kernel!(packed, fourier, tables.Plm, tables.weights, tables.scales, - RT(cfg.cphi), cfg.nlon, lcap, cfg.mmax, mcap; + RT(SHTnsKit._analysis_phi_scale(cfg)), cfg.nlon, lcap, cfg.mmax, mcap; ndrange=(lcap + 1, 2mcap + 1)) CUDA.synchronize() return packed @@ -1710,7 +1710,7 @@ function _cuda_batch_analysis(cfg::SHTConfig, fields::CUDA.AnyCuArray; gpu_fft!(fourier, 2) canonical = CUDA.zeros(CT, cfg.lmax + 1, cfg.mmax + 1, size(fields, 3)) kernel! = scalar_batch_analysis_kernel!(CUDABackend()) - kernel!(canonical, fourier, tables.Plm, tables.weights, RT(cfg.cphi), + kernel!(canonical, fourier, tables.Plm, tables.weights, RT(SHTnsKit._analysis_phi_scale(cfg)), cfg.lmax, cfg.mmax, cfg.mres; ndrange=size(canonical)) configured = canonical ./ reshape( tables.scales, cfg.lmax + 1, cfg.mmax + 1, 1, @@ -1752,7 +1752,7 @@ function _cuda_batch_analysis_direct!(cfg::SHTConfig, end backend = CUDABackend() scalar_batch_analysis_kernel!(backend)( - output, fourier, tables.Plm, tables.weights, RT(cfg.cphi), + output, fourier, tables.Plm, tables.weights, RT(SHTnsKit._analysis_phi_scale(cfg)), cfg.lmax, cfg.mmax, cfg.mres; ndrange=size(output), ) coefficient_batch_conversion_kernel!(backend)( diff --git a/ext/SHTnsKitLoopVecExt.jl b/ext/SHTnsKitLoopVecExt.jl index 4701577f..2fa2a865 100644 --- a/ext/SHTnsKitLoopVecExt.jl +++ b/ext/SHTnsKitLoopVecExt.jl @@ -7,6 +7,41 @@ using Base.Threads: @threads # Turbo-optimized variants live under the SHTnsKit namespace so users can call # SHTnsKit.analysis_turbo, synthesis_turbo, etc., when LoopVectorization is loaded. +""" + _turbo_threads_ok() -> Bool + +Whether this call may start its own threaded loop. `@threads :static` cannot be +nested or started concurrently, so a turbo transform invoked from inside an +outer threaded region (or from a worker task) must run serially on that worker. +Delegates to the same predicate the core orchestrators use, so the two agree. +""" +@inline _turbo_threads_ok() = SHTnsKit._use_internal_mloop_threads() + +"""Run a `for` loop with the given `@threads` schedule when safe, else serially.""" +macro _lv_threads(sched, loop) + loop isa Expr && loop.head === :for || + throw(ArgumentError("@_lv_threads requires a for loop")) + return esc(quote + if $(_turbo_threads_ok)() + Base.Threads.@threads $sched $loop + else + $loop + end + end) +end + +""" + _turbo_m_order(cfg) -> Vector{Int} + +The azimuthal orders this config actually represents: `0, mres, 2mres, …`. +Iterating a bare `0:mmax` here made `analysis_turbo` populate — and +`synthesis_turbo` consume — coefficient columns that an `mres > 1` transform has +no storage for, so the turbo pair silently disagreed with `analysis`/`synthesis` +(the disagreement is O(1), not roundoff). Shares the core's cached, load-balanced +ordering so the two stay in lockstep. +""" +@inline _turbo_m_order(cfg::SHTnsKit.SHTConfig) = SHTnsKit.cached_m_order(cfg) + """ SHTnsKit.analysis_turbo(cfg::SHTnsKit.SHTConfig, f::AbstractMatrix) @@ -18,27 +53,30 @@ function SHTnsKit.analysis_turbo(cfg::SHTnsKit.SHTConfig, f::AbstractMatrix) size(f, 1) == nlat || throw(DimensionMismatch("first dim must be nlat=$(nlat)")) size(f, 2) == nlon || throw(DimensionMismatch("second dim must be nlon=$(nlon)")) - fC = complex.(f) - Fφ = SHTnsKit.fft_phi(fC) + # One buffer + the shared cached FFTW plan, matching the core `analysis`. + # The old `fft_phi(complex.(f))` made a complex copy AND re-planned an + # out-of-place FFT on every call. + Fφ = SHTnsKit.fft_phi!(Matrix{complex(float(eltype(f)))}(undef, nlat, nlon), f) lmax, mmax = cfg.lmax, cfg.mmax CT = eltype(Fφ) alm = Matrix{CT}(undef, lmax + 1, mmax + 1) fill!(alm, 0.0 + 0.0im) - scaleφ = cfg.cphi + scaleφ = SHTnsKit._analysis_phi_scale(cfg) # inverts synthesis under any phi_scale (= cphi under :dft) # Bind cfg fields to locals so the @tturbo loops below operate on plain arrays. # LoopVectorization can't analyze property access (cfg.Nlm) inside @tturbo, and cfg is mutable. xv = cfg.x; wv = cfg.w # Adaptive threading: use nested parallelism for better load balancing n_threads = Threads.nthreads() - if mmax + 1 < n_threads ÷ 2 && nlat > 32 + m_order = _turbo_m_order(cfg) + if length(m_order) < n_threads ÷ 2 && nlat > 32 # Few m modes: parallelize over latitude points with thread-local accumulators n_tid = Threads.maxthreadid() CT = eltype(alm) # eltype-derived accumulators (this branch uses plain loops, not @tturbo) thread_alm = [Vector{CT}(undef, lmax + 1) for _ in 1:n_tid] thread_P_bufs = [Vector{Float64}(undef, lmax + 1) for _ in 1:n_tid] # per-thread Legendre scratch (hoisted out of the latitude loop) - for m in 0:mmax + for m in m_order col = m + 1 for t in 1:n_tid fill!(thread_alm[t], zero(CT)) @@ -46,7 +84,7 @@ function SHTnsKit.analysis_turbo(cfg::SHTnsKit.SHTConfig, f::AbstractMatrix) if cfg.use_plm_tables && length(cfg.NP_tables) == mmax + 1 # NP_tables[col][l+1, i] = P̄_l^m already; no extra Nlm multiply tbl = cfg.NP_tables[m + 1] - @threads :static for i in 1:nlat # :static pins iterations → threadid() stable (no buffer race under task migration) + @_lv_threads :static for i in 1:nlat # :static pins iterations → threadid() stable (no buffer race under task migration) tid = Threads.threadid() local_acc = thread_alm[tid] Fi = Fφ[i, col] @@ -56,7 +94,7 @@ function SHTnsKit.analysis_turbo(cfg::SHTnsKit.SHTConfig, f::AbstractMatrix) end end else - @threads :static for i in 1:nlat # :static pins iterations → threadid() stable (no buffer race under task migration) + @_lv_threads :static for i in 1:nlat # :static pins iterations → threadid() stable (no buffer race under task migration) tid = Threads.threadid() local_acc = thread_alm[tid] thread_P = thread_P_bufs[tid] @@ -79,7 +117,8 @@ function SHTnsKit.analysis_turbo(cfg::SHTnsKit.SHTConfig, f::AbstractMatrix) end else # Standard m-parallel approach with dynamic scheduling - @threads :dynamic for m in 0:mmax + @_lv_threads :dynamic for idx in eachindex(m_order) + m = m_order[idx] col = m + 1 if cfg.use_plm_tables && length(cfg.NP_tables) == mmax + 1 # NP_tables[col][l+1, i] = P̄_l^m already; no extra Nlm multiply @@ -124,9 +163,11 @@ function SHTnsKit.synthesis_turbo(cfg::SHTnsKit.SHTConfig, alm::AbstractMatrix; alm_int = SHTnsKit._internal_coefficients(alm, cfg) nlat, nlon = cfg.nlat, cfg.nlon - CT = eltype(alm_int) + # Fourier bins are complex even when every spectral coefficient is real + # (matches the core `_synthesis`). + CT = complex(float(eltype(alm_int))) Fφ = Matrix{CT}(undef, nlat, nlon) - fill!(Fφ, 0.0 + 0.0im) + fill!(Fφ, zero(CT)) inv_scaleφ = SHTnsKit.phi_inv_scale(cfg) # Bind cfg fields to locals so the @tturbo loops below operate on plain arrays. @@ -135,17 +176,18 @@ function SHTnsKit.synthesis_turbo(cfg::SHTnsKit.SHTConfig, alm::AbstractMatrix; # Adaptive threading: use nested parallelism for better load balancing n_threads = Threads.nthreads() - if mmax + 1 < n_threads ÷ 2 && nlat > 32 + m_order = _turbo_m_order(cfg) + if length(m_order) < n_threads ÷ 2 && nlat > 32 # Few m modes: parallelize over latitude points instead n_tid = Threads.maxthreadid() thread_P_bufs = [Vector{Float64}(undef, lmax + 1) for _ in 1:n_tid] # per-thread Legendre scratch (hoisted out of the latitude loop) G = Vector{CT}(undef, nlat) # shared latitude scratch (few-m branch: outer m-loop is serial) - for m in 0:mmax + for m in m_order col = m + 1 if cfg.use_plm_tables && length(cfg.NP_tables) == mmax + 1 # NP_tables[col][l+1, i] = P̄_l^m already; no extra Nlm multiply tbl = cfg.NP_tables[m + 1] - @threads for i in 1:nlat + @_lv_threads :dynamic for i in 1:nlat g_re = 0.0 g_im = 0.0 @tturbo warn_check_args=false for l in m:lmax @@ -157,7 +199,7 @@ function SHTnsKit.synthesis_turbo(cfg::SHTnsKit.SHTConfig, alm::AbstractMatrix; G[i] = complex(g_re, g_im) end else - @threads :static for i in 1:nlat # :static pins iterations → threadid() stable + @_lv_threads :static for i in 1:nlat # :static pins iterations → threadid() stable thread_P = thread_P_bufs[Threads.threadid()] SHTnsKit.Plm_norm_row!(thread_P, xv[i], lmax, m) g_re = 0.0 @@ -186,7 +228,8 @@ function SHTnsKit.synthesis_turbo(cfg::SHTnsKit.SHTConfig, alm::AbstractMatrix; # Per-m latitude scratch as columns of one pre-allocated matrix (each m owns # a distinct column → race-free, no per-iteration allocation). Gcols = Matrix{CT}(undef, nlat, mmax + 1) - @threads :dynamic for m in 0:mmax + @_lv_threads :dynamic for idx in eachindex(m_order) + m = m_order[idx] col = m + 1 thread_G = view(Gcols, :, col) if cfg.use_plm_tables && length(cfg.NP_tables) == mmax + 1 diff --git a/ext/SHTnsKitParallelAMDGPUExt.jl b/ext/SHTnsKitParallelAMDGPUExt.jl index 0b6d74cf..9a040698 100644 --- a/ext/SHTnsKitParallelAMDGPUExt.jl +++ b/ext/SHTnsKitParallelAMDGPUExt.jl @@ -66,7 +66,7 @@ function ParallelExt._dist_transpose_gpu_analysis!( AMDGPU.ROCBackend(), ) kernel!(parent(output), parent(plan.F_buf), tables.Plm, tables.weights, - tables.scales, RT(plan.cfg.cphi), _first_m(plan), plan.lmax, + tables.scales, RT(SHTnsKit._analysis_phi_scale(plan.cfg)), _first_m(plan), plan.lmax, plan.mmax, plan.cfg.mres, plan.lmax; ndrange=size(parent(output))) AMDGPU.synchronize() @@ -111,7 +111,7 @@ function ParallelExt._dist_transpose_gpu_vector_analysis!( ) kernel!(parent(Sout), parent(Tout), parent(plan.F_buf), parent(plan.F_buf2), tables.dtheta, tables.over_sin, - tables.weights, tables.scales, tables.x, RT(plan.cfg.cphi), + tables.weights, tables.scales, tables.x, RT(SHTnsKit._analysis_phi_scale(plan.cfg)), _first_m(plan), plan.lmax, plan.mmax, plan.cfg.mres, plan.cfg.robert_form; ndrange=size(parent(Sout))) diff --git a/ext/SHTnsKitParallelCUDAExt.jl b/ext/SHTnsKitParallelCUDAExt.jl index cadceadf..569c4af5 100644 --- a/ext/SHTnsKitParallelCUDAExt.jl +++ b/ext/SHTnsKitParallelCUDAExt.jl @@ -58,7 +58,7 @@ function ParallelExt._dist_transpose_gpu_analysis!( CUDA.CUDABackend(), ) kernel!(parent(output), parent(plan.F_buf), tables.Plm, tables.weights, - tables.scales, RT(plan.cfg.cphi), _first_m(plan), plan.lmax, + tables.scales, RT(SHTnsKit._analysis_phi_scale(plan.cfg)), _first_m(plan), plan.lmax, plan.mmax, plan.cfg.mres, plan.lmax; ndrange=size(parent(output))) CUDA.synchronize() @@ -103,7 +103,7 @@ function ParallelExt._dist_transpose_gpu_vector_analysis!( ) kernel!(parent(Sout), parent(Tout), parent(plan.F_buf), parent(plan.F_buf2), tables.dtheta, tables.over_sin, - tables.weights, tables.scales, tables.x, RT(plan.cfg.cphi), + tables.weights, tables.scales, tables.x, RT(SHTnsKit._analysis_phi_scale(plan.cfg)), _first_m(plan), plan.lmax, plan.mmax, plan.cfg.mres, plan.cfg.robert_form; ndrange=size(parent(Sout))) diff --git a/ext/SHTnsKitParallelExt.jl b/ext/SHTnsKitParallelExt.jl index 6e18befa..37ba147b 100644 --- a/ext/SHTnsKitParallelExt.jl +++ b/ext/SHTnsKitParallelExt.jl @@ -63,7 +63,9 @@ DEBUGGING TIPS ENVIRONMENT VARIABLES -------------------- -- SHTNSKIT_CACHE_PENCILFFTS: "1" (default) to cache FFT plans, "0" to disable +- SHTNSKIT_FFT_PLAN_CACHE: "1" (default) to cache φ-FFT plans, "0" to disable. + Read by SHTnsKit proper — the cache is shared with the serial transforms. + (legacy alias: SHTNSKIT_CACHE_PENCILFFTS) - SHTNSKIT_VERBOSE_STORAGE: "1" to print storage optimization info ================================================================================ =# @@ -75,8 +77,9 @@ Parallel extension module providing MPI-distributed spherical harmonic transform See module-level comments for architecture overview and debugging tips. # Module state -The extension keeps its FFT plan caches, locks, and cache controls in direct -module constants. +The extension keeps its locks and caches in direct module constants. The φ-FFT +plan cache itself lives in SHTnsKit proper (src/fftutils.jl) and is shared with +the serial transforms. """ using Base.Threads # Threads.@threads and locks/macros @@ -132,119 +135,20 @@ import SHTnsKit # Core spherical harmonic functionality end # ===== MODULE STATE ===== -const _CACHE_PENCILFFTS = Ref(get(ENV, "SHTNSKIT_CACHE_PENCILFFTS", "1") == "1") -const _pfft_cache = IdDict{Any,Any}() -const _PFFT_CACHE_MAX = Ref(parse(Int, get(ENV, "SHTNSKIT_PFFT_CACHE_MAX", "64"))) -const _cache_lock = Threads.ReentrantLock() -const _fftw_cache_lock = Threads.ReentrantLock() - -""" - pfft_cache_max!(n::Int) -> Int - -Set the maximum number of cached FFT plans (shared across all grids and -communicators). Set `n <= 0` to disable the cap entirely. Returns the previous -value. -""" -function pfft_cache_max!(n::Int) - prev = _PFFT_CACHE_MAX[] - _PFFT_CACHE_MAX[] = n - return prev -end +# The φ-FFT plan cache (and its `enable_fft_plan_cache!` / `disable_fft_plan_cache!` +# / `set_fft_plan_cache!` / `fft_plan_cache_enabled` controls) lives in +# src/fftutils.jl and is shared with the serial transforms. This extension used to +# carry a second, parallel-only cache keyed on the pencil decomposition — but the +# only function that ever consulted it, `_get_or_plan`, had no call sites, and the +# "plans" it stored were `NamedTuple` placeholders that the FFT wrappers ignored. +# Every knob pointed at it was therefore inert. It has been deleted rather than +# left in place; recover it from git history if a PencilFFTs-level cache is ever +# actually needed. # Compat helper: `ceildiv` was added in Julia 1.11 const _ceildiv = isdefined(Base, :ceildiv) ? Base.ceildiv : (a, b) -> cld(a, b) ceildiv(a::Integer, b::Integer) = _ceildiv(a, b) -function _fft_plan_cache_enabled_impl() - return _CACHE_PENCILFFTS[] -end - -function _fft_plan_cache_set_impl(flag::Bool; clear::Bool=true) - _CACHE_PENCILFFTS[] = flag - if !flag && clear - lock(_cache_lock) do - empty!(_pfft_cache) - end - end - return flag -end - -function _fft_plan_cache_enable_impl() - return _fft_plan_cache_set_impl(true) -end - -function _fft_plan_cache_disable_impl(; clear::Bool=true) - return _fft_plan_cache_set_impl(false; clear=clear) -end - -@inline function _decomp_hash(A) - if hasfield(typeof(A), :pencil) - pencil = getfield(A, :pencil) - if hasfield(typeof(pencil), :decomposition) - return hash(getfield(pencil, :decomposition)) - elseif hasfield(typeof(pencil), :plan) - return hash(getfield(pencil, :plan)) - end - end - return hash(size(A)) -end - -# Generate cache key based on array characteristics for FFT plan reuse -function _cache_key(kind::Symbol, A) - # Basic array characteristics - base_key = (kind, size(A,1), size(A,2), eltype(A)) - - # Add communicator size with robust error handling - comm_size = try - MPI.Comm_size(communicator(A)) - catch - 1 # Default to single process - end - - # Decomposition hash — no try/catch to avoid closure-box allocations on hot path. - decomp_hash = _decomp_hash(A) - - return (base_key..., comm_size, decomp_hash) -end - -function _get_or_plan(kind::Symbol, A) - # If caching disabled, create plan directly without storing - if !_CACHE_PENCILFFTS[] - return kind === :fft ? plan_fft(A; dims=2) : # Forward FFT along longitude (dim 2) - kind === :ifft ? plan_ifft(A; dims=2) : # Inverse FFT along longitude - kind === :rfft ? (try plan_rfft(A; dims=2) catch; nothing end) : # Real-to-complex FFT - kind === :irfft ? (try plan_irfft(A; dims=2) catch; nothing end) : # Complex-to-real IFFT - error("unknown plan kind") - end - - # Thread-safe caching with optimized lookup - key = _cache_key(kind, A) - - # Thread-safe plan lookup and creation - return lock(_cache_lock) do - # Double-check pattern: another thread might have created the plan - if haskey(_pfft_cache, key) - return _pfft_cache[key] - end - - # Create new plan and cache it for future use - plan = kind === :fft ? plan_fft(A; dims=2) : # Forward FFT along longitude - kind === :ifft ? plan_ifft(A; dims=2) : # Inverse FFT along longitude - kind === :rfft ? (try plan_rfft(A; dims=2) catch; nothing end) : # Real-to-complex FFT - kind === :irfft ? (try plan_irfft(A; dims=2) catch; nothing end) : # Complex-to-real IFFT - error("unknown plan kind") - - # Enforce the soft cap: flush before inserting so the fresh entry survives. - cap = _PFFT_CACHE_MAX[] - if cap > 0 && length(_pfft_cache) >= cap - empty!(_pfft_cache) - end - _pfft_cache[key] = plan - return plan - end -end - - # ===== PENCIL GRID SUGGESTION ===== @inline function _infer_comm_size(comm_or_nprocs::Any) comm_or_nprocs === nothing && return 1 @@ -404,44 +308,6 @@ end # Use FFTW for 1D FFTs along the longitude dimension (not PencilFFTs which is for multi-D) # PencilArrays provides the distributed array framework, FFTW provides the FFTs -# Cache for FFTW 1D plans (key includes inplace flag) -const _fftw_plan_cache = Dict{Tuple{Symbol, Int, DataType, Bool}, Any}() - -""" - get_fftw_plan(kind, n, T) -> plan - -Get or create a cached FFTW plan for 1D transforms. -""" -function get_fftw_plan(kind::Symbol, n::Int, ::Type{T}; inplace::Bool=false) where T - key = (kind, n, T, inplace) - lock(_fftw_cache_lock) do - if haskey(_fftw_plan_cache, key) - return _fftw_plan_cache[key] - end - - # Create sample array for planning - if kind == :fft - sample = zeros(Complex{real(T)}, n) - plan = inplace ? FFTW.plan_fft!(sample) : FFTW.plan_fft(sample) - elseif kind == :ifft - sample = zeros(Complex{real(T)}, n) - plan = inplace ? FFTW.plan_ifft!(sample) : FFTW.plan_ifft(sample) - elseif kind == :rfft - sample = zeros(real(T), n) - plan = FFTW.plan_rfft(sample) # rfft is always out-of-place - elseif kind == :irfft - # For irfft, input size is n÷2+1 - sample = zeros(Complex{real(T)}, n ÷ 2 + 1) - plan = FFTW.plan_irfft(sample, n) - else - error("Unknown FFT kind: $kind") - end - - _fftw_plan_cache[key] = plan - return plan - end -end - """ fft_along_dim2!(output, input) @@ -530,71 +396,6 @@ function ifft_along_dim2!(output::AbstractMatrix{Complex{T}}, input::AbstractMat return output end -# Local FFT wrappers used by the extension's plan cache. -function plan_fft(A::PencilArray; dims=:) - # Return a placeholder that indicates we'll use FFTW on local data - return (kind=:fft, local_size=size(parent(A))) -end - -function plan_ifft(A::PencilArray; dims=:) - return (kind=:ifft, local_size=size(parent(A))) -end - -function fft(A::PencilArray, p) - local_data = parent(A) - nlat, nlon = size(local_data) - output = similar(local_data, Complex{Float64}) - fft_along_dim2!(output, local_data) - return output -end - -function ifft(A::PencilArray, p) - local_data = parent(A) - nlat, nlon = size(local_data) - output = similar(local_data) - ifft_along_dim2!(output, local_data) - return output -end - -# RFFT/IRFFT variants -function plan_rfft(A::PencilArray; dims=:) - return (kind=:rfft, local_size=size(parent(A))) -end - -function plan_irfft(A::PencilArray; dims=:) - return (kind=:irfft, local_size=size(parent(A))) -end - -function rfft(A::PencilArray, p) - local_data = parent(A) - nlat, nlon = size(local_data) - nk = nlon ÷ 2 + 1 - output = Matrix{ComplexF64}(undef, nlat, nk) - @inbounds for i in 1:nlat - row = Vector{Float64}(collect(view(local_data, i, :))) - fft_result = FFTW.rfft(row) - for j in 1:nk - output[i, j] = fft_result[j] - end - end - return output -end - -function irfft(A::AbstractMatrix{<:Complex}, p) - nlat, nk = size(A) - # Assume original nlon was 2*(nk-1) for even-length arrays - nlon = 2 * (nk - 1) - output = Matrix{Float64}(undef, nlat, nlon) - @inbounds for i in 1:nlat - row = Vector{ComplexF64}(collect(view(A, i, :))) - ifft_result = FFTW.irfft(row, nlon) - for j in 1:nlon - output[i, j] = ifft_result[j] - end - end - return output -end - # ===== OPTIMIZED DISTRIBUTED FFT USING TRANSPOSE ===== # When φ is distributed, use a single all-to-all transpose instead of per-row Allgatherv. # This reduces the number of MPI calls from O(nlat) to O(1). @@ -710,38 +511,6 @@ function distributed_rfft_phi!(Fθm_out::AbstractMatrix{Complex{T}}, return Fθm_out end -""" - distributed_irfft_phi!(local_out, Fθm, θ_range, φ_range, nlon, comm) - -Complex-to-real inverse FFT for distributed synthesis. `Fθm` is `(nlat_local, -nlon÷2+1)` and must be identical on every rank in a given θ-slab (caller -responsibility — typical pattern replicates the Fourier buffer). After local -`irfft` to full `(nlat_local, nlon)` real, the function slices this rank's -local φ window into `local_out`. -""" -function distributed_irfft_phi!(local_out::AbstractMatrix{<:Real}, - Fθm::AbstractMatrix{<:Complex}, - θ_range::AbstractRange, φ_range::AbstractRange, - nlon::Int, comm) - nlat_local = length(θ_range) - nlon_local = length(φ_range) - size(Fθm, 2) == nlon ÷ 2 + 1 || throw(DimensionMismatch("Fθm must have nlon÷2+1 columns")) - size(Fθm, 1) == nlat_local || throw(DimensionMismatch("Fθm must have nlat_local rows")) - size(local_out) == (nlat_local, nlon_local) || throw(DimensionMismatch("local_out must be (nlat_local, nlon_local)")) - - spatial_full = Matrix{eltype(local_out)}(undef, nlat_local, nlon) - spatial_full .= FFTW.irfft(Fθm, nlon, 2) - - φ_start = first(φ_range) - @inbounds for j in 1:nlon_local - for i in 1:nlat_local - local_out[i, j] = spatial_full[i, φ_start + j - 1] - end - end - - return local_out -end - """ distributed_ifft_phi!(local_out, Fθm, θ_range, φ_range, nlon, comm) @@ -801,52 +570,6 @@ function efficient_spectral_reduce!(local_data::AbstractVector, comm) return local_data end -""" - bandwidth_aware_broadcast!(data, root, comm) - -Bandwidth-aware broadcasting that adapts to network topology and data size. -Uses pipeline broadcasting for large data and tree broadcasting for small data. -""" -function bandwidth_aware_broadcast!(data::AbstractArray, root::Int, comm) - nprocs = MPI.Comm_size(comm) - data_size_mb = (sizeof(data)) / (1024 * 1024) - - if nprocs > 32 && data_size_mb > 10.0 - # Use pipeline broadcast for large data on large clusters - pipeline_broadcast!(data, root, comm) - else - # Use standard tree broadcast for smaller cases - MPI.Bcast!(data, root, comm) - end - - return data -end - -""" - pipeline_broadcast!(data, root, comm) - -Pipeline broadcast that overlaps communication with local copying for better bandwidth utilization. -""" -function pipeline_broadcast!(data::AbstractArray, root::Int, comm) - rank = MPI.Comm_rank(comm) - nprocs = MPI.Comm_size(comm) - - # Determine pipeline parameters - pipeline_stages = min(nprocs, 8) # Limit pipeline depth - chunk_size = max(1, length(data) ÷ pipeline_stages) - - for stage in 1:pipeline_stages - start_idx = (stage - 1) * chunk_size + 1 - end_idx = stage == pipeline_stages ? length(data) : stage * chunk_size - chunk_view = view(data, start_idx:end_idx) - - # Pipeline broadcast of this chunk - MPI.Bcast!(chunk_view, root, comm) - end - - return data -end - # Note: Avoid forwarding Base.zeros(Pencil) to PencilArrays.zeros to prevent # potential recursion when PencilArrays.zeros may call Base.zeros internally. diff --git a/ext/SHTnsKitZygoteExt.jl b/ext/SHTnsKitZygoteExt.jl index 70b94a47..7bf35d16 100644 --- a/ext/SHTnsKitZygoteExt.jl +++ b/ext/SHTnsKitZygoteExt.jl @@ -286,16 +286,20 @@ Zygote.@adjoint function SHTnsKit.SH_Zrotate(cfg::SHTnsKit.SHTConfig, Qlm::Abstr end Zygote.@adjoint function SHTnsKit.SH_Yrotate(cfg::SHTnsKit.SHTConfig, Qlm::AbstractVector{<:Complex}, alpha::Real, Rlm::AbstractVector{<:Complex}) + # Snapshot BEFORE the primal: an in-place rotation (Rlm === Qlm) overwrites + # Qlm, and callers may reuse either buffer before the pullback runs. The dα + # formula needs the *input* coefficients, so they must be preserved here. + Qlm_saved = copy(Qlm) y = SHTnsKit.SH_Yrotate(cfg, Qlm, alpha, Rlm) function back(ȳ) inverse = SHTnsKit.SHTRotation(cfg.lmax, cfg.mmax) SHTnsKit.shtns_rotation_set_angles_ZYZ(inverse, 0.0, -alpha, 0.0) - Q̄ = similar(Qlm) + Q̄ = similar(Qlm_saved) _zyg_configured_rotation_adjoint!(cfg, inverse, ȳ, Q̄) # angle gradient via derivative of Wigner-d at beta=alpha dα = zero(float(alpha)) lmax, mmax = cfg.lmax, cfg.mmax - Qlm_canonical = SHTnsKit._internal_coefficients(Qlm, cfg) + Qlm_canonical = SHTnsKit._internal_coefficients(Qlm_saved, cfg) ȳ_canonical = SHTnsKit._analysis_cotangent_to_canonical(ȳ, cfg) for l in 0:lmax mm = min(l, mmax) diff --git a/src/SHTnsKit.jl b/src/SHTnsKit.jl index 6c64747a..5ca9ce1e 100644 --- a/src/SHTnsKit.jl +++ b/src/SHTnsKit.jl @@ -101,7 +101,9 @@ ENVIRONMENT VARIABLES -------------------- - SHTNSKIT_PHI_SCALE: "dft" or "quad" for φ scaling convention - SHTNSKIT_VERBOSE_STORAGE: "1" to print storage optimization info -- SHTNSKIT_CACHE_PENCILFFTS: "0" to disable FFT plan caching (parallel ext) +- SHTNSKIT_FFT_PLAN_CACHE: "0" to disable φ-FFT plan caching (serial and distributed) + (legacy alias: SHTNSKIT_CACHE_PENCILFFTS) +- SHTNSKIT_FFT_PLAN_CACHE_MAX: cap on distinct cached plans (default 64; ≤0 = no cap) DEBUGGING TIPS -------------- @@ -146,13 +148,58 @@ function phi_inv_scale(cfg::SHTConfig) end if cfg.phi_scale === :quad return cfg.nlon / (2π) - elseif cfg.phi_scale === :dft - return Float64(cfg.nlon) else - return cfg.grid_type == :gauss ? Float64(cfg.nlon) : cfg.nlon / (2π) + # `:dft` and anything unset (`:auto`). The old fallback keyed on + # `grid_type` and handed every non-Gauss grid `nlon/2π` — so a regular + # grid built through the exported keyword constructor (which defaulted to + # `:auto`) disagreed by 2π with the identical grid from + # `create_regular_config`, which sets `:dft` explicitly. Both constructors + # emit `:dft`, so `:dft` is the right default for an unset value. + return Float64(cfg.nlon) end end +""" + _analysis_phi_scale(cfg) -> Float64 + +The φ quadrature factor `analysis` must apply so that it inverts `synthesis`. + +`synthesis` scales its Fourier bins by `phi_inv_scale(cfg)` and the inverse FFT +divides by `nlon`, a net spatial factor of `σ = phi_inv_scale(cfg)/nlon`. For the +pair to be mutually inverse, analysis must carry `cphi/σ`. + +Under the default `:dft` mode `σ = 1` and this is just `cphi = 2π/nlon`, exactly +what analysis always used — so nothing changes for any configuration the +`create_*_config` constructors produce. Under `:quad` (`σ = 1/2π`) the old fixed +`cphi` made `analysis(synthesis(alm))` come back as `alm/2π`: the two halves of +the transform pair simply disagreed about the convention, with synthesis honouring +`phi_scale` and analysis ignoring it. +""" +@inline _analysis_phi_scale(cfg::SHTConfig) = cfg.cphi * cfg.nlon / phi_inv_scale(cfg) + +""" + _evaluator_phi_scale(cfg) -> Float64 + +The factor a point/latitude evaluator must apply to reproduce the value +`synthesis` writes on the grid: `phi_inv_scale(cfg)/nlon`, i.e. 1 under `:dft` +and `1/2π` under `:quad`. Omitting it made every direct evaluator disagree with +the very grid it claims to sample by a factor of 2π. +""" +@inline _evaluator_phi_scale(cfg::SHTConfig) = phi_inv_scale(cfg) / cfg.nlon + +""" + _evaluator_phi_scale(cfg, ::Type{T}) -> T + +`_evaluator_phi_scale` narrowed to the evaluator's own real type. + +The evaluators promise their caller the element type their coefficients carry — +`Float32` coefficients give `Float32` values, and a `Dual` stays a `Dual`. The +untyped scale is a `Float64`, so multiplying by it silently widens every +`Float32` result to `Float64`. Convert once, at the boundary. +""" +@inline _evaluator_phi_scale(cfg::SHTConfig, ::Type{T}) where {T} = + convert(real(T), _evaluator_phi_scale(cfg)) + include("buffer_utils.jl") # Common buffer allocation patterns include("kernels.jl") # Legendre accumulation kernels include("plan.jl") # Transform planning and optimization @@ -185,6 +232,7 @@ export set_allow_padding!, disable_padding!, is_padding_enabled # Memory p export get_nlat_padded, get_spat_dist, compute_optimal_padding # Padding queries export allocate_padded_spatial, allocate_padded_spatial_batch # Padded array allocation export copy_to_padded!, copy_from_padded!, estimate_padding_overhead # Padding utilities +export spatial_view # Padded buffer → transform-shaped view # ===== BASIC TRANSFORMS ===== # The `*_cplx` helpers are intentionally separate from `real_output=false` @@ -340,52 +388,15 @@ function DistQstPlan(args...; kwargs...) return getproperty(ext, :DistQstPlan)(args...; kwargs...) end -function fft_plan_cache_enabled() - ext = _parallel_ext_module() - return ext === nothing ? false : getproperty(ext, :_fft_plan_cache_enabled_impl)() -end - -function set_fft_plan_cache!(flag::Bool; clear::Bool=true) - ext = _parallel_ext_module() - ext === nothing && error("Parallel extension not loaded") - return getproperty(ext, :_fft_plan_cache_set_impl)(flag; clear=clear) -end - -function enable_fft_plan_cache!() - ext = _parallel_ext_module() - ext === nothing && error("Parallel extension not loaded") - return getproperty(ext, :_fft_plan_cache_enable_impl)() -end - -function disable_fft_plan_cache!(; clear::Bool=true) - ext = _parallel_ext_module() - ext === nothing && error("Parallel extension not loaded") - return getproperty(ext, :_fft_plan_cache_disable_impl)(; clear=clear) -end - -Base.@doc """ - fft_plan_cache_enabled() -> Bool - -Return whether distributed FFT plan caching is currently enabled. -""" fft_plan_cache_enabled - -Base.@doc """ - set_fft_plan_cache!(flag::Bool; clear::Bool=true) - -Enable or disable caching of distributed FFT plans. When disabling and `clear=true`, cached plans are freed. -""" set_fft_plan_cache! - -Base.@doc """ - enable_fft_plan_cache!() - -Convenience wrapper to enable distributed FFT plan caching. -""" enable_fft_plan_cache! - -Base.@doc """ - disable_fft_plan_cache!(; clear::Bool=true) - -Disable distributed FFT plan caching. Pass `clear=false` to retain existing cache entries. -""" disable_fft_plan_cache! +# NOTE: the φ-FFT plan cache and its `fft_plan_cache_enabled` / +# `set_fft_plan_cache!` / `enable_fft_plan_cache!` / `disable_fft_plan_cache!` +# controls live in src/fftutils.jl. They used to forward to the parallel +# extension, where the cache they addressed had no readers at all — `_get_or_plan` +# was never called from anywhere, so every one of these knobs (and the +# `SHTNSKIT_CACHE_PENCILFFTS` environment variable the distributed guide +# advertises) was a no-op. The cache that every transform actually uses, serial +# and distributed alike, is the one in fftutils.jl, so the controls now address +# that and no longer require the extension to be loaded. # ===== PENCIL GRID SUGGESTION ===== function _suggest_pencil_grid_fallback(comm_or_nprocs::Any, nlat::Integer, nlon::Integer; diff --git a/src/batch_transforms.jl b/src/batch_transforms.jl index 7d65174b..72798e7b 100644 --- a/src/batch_transforms.jl +++ b/src/batch_transforms.jl @@ -228,9 +228,19 @@ end """ set_batch_size!(cfg::SHTConfig, howmany::Int; spec_dist::Int=0) -Configure batch processing for multiple fields. After calling this function, -batch transform functions (`analysis_batch`, `synthesis_batch`, etc.) will -process `howmany` fields simultaneously. +Record the SHTns-style batched-layout descriptors on the configuration. + +!!! note "Advisory metadata" + `howmany` and `spec_dist` mirror the SHTns C API's batch descriptors and are + stored for interoperability, but the Julia batch entry points do **not** + consult them: `analysis_batch`, `synthesis_batch` and their sphtor/QST + siblings take the field count from `size(fields, 3)` of the array you pass. + Setting a batch size neither constrains nor accelerates those calls, and + passing a different number of fields is allowed and works. + +Configure the batch descriptors carried by `cfg` for callers that mirror the +SHTns C layout. See the note above for what the Julia batch transforms actually +use. # Arguments - `cfg`: SHTConfig to modify @@ -245,11 +255,10 @@ process `howmany` fields simultaneously. # Example ```julia cfg = create_gauss_config(32, 34) -set_batch_size!(cfg, 4) # Process 4 fields at once +set_batch_size!(cfg, 4) # recorded on cfg; does not bind the call below -# Now use batch transforms fields = rand(cfg.nlat, cfg.nlon, 4) -alms = analysis_batch(cfg, fields) +alms = analysis_batch(cfg, fields) # count comes from size(fields, 3) ``` """ function set_batch_size!(cfg::SHTConfig, howmany::Int; spec_dist::Int=0) @@ -322,7 +331,7 @@ function analysis_batch(cfg::SHTConfig, fields::AbstractArray{<:Real,3}; use_rff _batch_fft_phi!(Fφ_batch, fields) end - scaleφ = cfg.cphi + scaleφ = _analysis_phi_scale(cfg) # inverts synthesis under any phi_scale (= cphi under :dft) # Hoist cfg field reads to locals: cfg is mutable, so reads inside the m/l loops below aren't auto-hoisted. w = cfg.w @@ -423,7 +432,7 @@ function analysis_batch!(cfg::SHTConfig, alm_out::AbstractArray{<:Complex,3}, # mutating an output view that overlaps the spatial input is safe. fill!(alm_out, zero(eltype(alm_out))) - scaleφ = cfg.cphi + scaleφ = _analysis_phi_scale(cfg) # inverts synthesis under any phi_scale (= cphi under :dft) # Hoist cfg field reads to locals: cfg is mutable, so reads inside the m/l loops below aren't auto-hoisted. w = cfg.w diff --git a/src/complex_packed.jl b/src/complex_packed.jl index 017e9980..6fe1be9f 100644 --- a/src/complex_packed.jl +++ b/src/complex_packed.jl @@ -179,7 +179,7 @@ function _analysis_packed_cplx(cfg::SHTConfig, z::AbstractMatrix{<:Complex}, # fft_phi re-plans each call). eltype preserved for the AD/DFT fallback. Fφ = fft_phi!(Matrix{complex(float(eltype(z)))}(undef, size(z)...), z) P = Vector{Float64}(undef, lmax + 1) - scaleφ = cfg.cphi + scaleφ = _analysis_phi_scale(cfg) # inverts synthesis under any phi_scale (= cphi under :dft) xv = cfg.x; wv = cfg.w # hoist field reads out of the m/l loops (cfg is mutable, so not auto-hoisted) # Read both signs of m from the FFT output and store them in LM_cplx order. @@ -285,7 +285,10 @@ function synthesis_point_cplx(cfg::SHTConfig, alm::AbstractVector{<:Complex}, co acc += gp * cis(convert(PT, am) * ph) am > 0 && (acc += gn * cis(-convert(PT, am) * ph)) end - return acc + # Same φ convention factor `synthesis_packed_cplx` carries (1 under :dft; + # 1/2π under :quad), so a point evaluation matches the grid it samples. + sφ = _evaluator_phi_scale(cfg, typeof(acc)) + return sφ == 1 ? acc : acc * sφ end diff --git a/src/config.jl b/src/config.jl index 67ff205c..55def3e6 100644 --- a/src/config.jl +++ b/src/config.jl @@ -248,6 +248,33 @@ function SHTConfig(; nlat_padded::Integer = 0, spat_dist::Integer = 0, ) + # Validate the invariants the transforms rely on. This constructor is + # exported and was previously unchecked, so a hand-built configuration could + # violate e.g. `nlon >= 2*mmax+1` and then silently synthesize an all-zero + # field for any mode it could not resolve. + lmax >= 0 || throw(ArgumentError("lmax must be ≥ 0, got $lmax")) + mmax >= 0 || throw(ArgumentError("mmax must be ≥ 0, got $mmax")) + mmax <= lmax || throw(ArgumentError("mmax must be ≤ lmax, got mmax=$mmax, lmax=$lmax")) + mres >= 1 || throw(ArgumentError("mres must be ≥ 1, got $mres")) + nlat >= 1 || throw(ArgumentError("nlat must be ≥ 1, got $nlat")) + nlon >= 2*mmax + 1 || throw(ArgumentError( + "nlon must be ≥ 2*mmax+1 = $(2*mmax+1) to resolve every azimuthal order, got nlon=$nlon")) + length(θ) == nlat || throw(DimensionMismatch("θ must have nlat=$nlat entries, got $(length(θ))")) + length(x) == nlat || throw(DimensionMismatch("x must have nlat=$nlat entries, got $(length(x))")) + length(w) == nlat || throw(DimensionMismatch("w must have nlat=$nlat entries, got $(length(w))")) + length(st) == nlat || throw(DimensionMismatch("st must have nlat=$nlat entries, got $(length(st))")) + length(φ) == nlon || throw(DimensionMismatch("φ must have nlon=$nlon entries, got $(length(φ))")) + size(Nlm) == (lmax + 1, mmax + 1) || throw(DimensionMismatch( + "Nlm must be ($(lmax+1), $(mmax+1)), got $(size(Nlm))")) + let expected_nlm = nlm_calc(lmax, mmax, mres) + nlm == expected_nlm || throw(ArgumentError( + "nlm must be nlm_calc(lmax, mmax, mres) = $expected_nlm, got $nlm")) + end + length(li) == nlm || throw(DimensionMismatch("li must have nlm=$nlm entries, got $(length(li))")) + length(mi) == nlm || throw(DimensionMismatch("mi must have nlm=$nlm entries, got $(length(mi))")) + nspat == nlat * nlon || throw(ArgumentError( + "nspat must be nlat*nlon = $(nlat*nlon), got $nspat")) + # Concretize vectors/matrices so SHTGrid/SHTNorm/SHTTables fields always # carry the exact declared types. Lets callers pass ranges, views, etc. grid = SHTGrid(collect(Float64, θ), collect(Float64, φ), @@ -392,11 +419,63 @@ function Base.setproperty!(cfg::SHTConfig, name::Symbol, val) return setfield!(getfield(cfg, :_scratch), :otf_Pb, val) elseif name === :_m_order return setfield!(getfield(cfg, :_scratch), :m_order, val) + # ----- structural fields: keep derived state consistent ----- + elseif name === :lmax || name === :mmax || name === :mres + setfield!(cfg, name, Int(val)) + _rebuild_spectral_layout!(cfg) + return val + elseif name === :nlat || name === :nlon || name === :grid_type || + name === :nlm || name === :li || name === :mi || name === :nspat + throw(ArgumentError( + "`cfg.$(name)` cannot be reassigned: the quadrature nodes, weights and " * + "packed index tables are all derived from it, and there is no grid-type-" * + "independent way to regenerate them in place. Build a new configuration " * + "with `create_gauss_config` / `create_regular_config` / `create_config` " * + "instead. (`lmax`, `mmax` and `mres` may be assigned; the spectral layout " * + "is rebuilt for you.)")) else return setfield!(cfg, name, val) end end +""" + _rebuild_spectral_layout!(cfg::SHTConfig) + +Regenerate everything derived from `lmax`/`mmax`/`mres` after one of them is +reassigned: the normalization table, the packed mode count and its `li`/`mi` +lookups, and the cached norm-scale matrix and m-ordering. + +Without this, assigning `cfg.lmax = 10` left `size(cfg.Nlm) == (7, 7)` while the +transforms indexed it as `(lmax+1, mmax+1)` under `@inbounds` — an out-of-bounds +read of a live array. Precomputed Legendre tables are sized `(lmax+1, nlat)` and +cannot survive the change either, so they are dropped; transforms fall back to +the on-the-fly path until `prepare_plm_tables!` is called again. +""" +function _rebuild_spectral_layout!(cfg::SHTConfig) + lmax = getfield(cfg, :lmax); mmax = getfield(cfg, :mmax); mres = getfield(cfg, :mres) + lmax >= 0 || throw(ArgumentError("lmax must be ≥ 0, got $lmax")) + mmax >= 0 || throw(ArgumentError("mmax must be ≥ 0, got $mmax")) + mmax <= lmax || throw(ArgumentError("mmax must be ≤ lmax, got mmax=$mmax, lmax=$lmax")) + mres >= 1 || throw(ArgumentError("mres must be ≥ 1, got $mres")) + + setfield!(cfg, :nlm, nlm_calc(lmax, mmax, mres)) + li, mi = build_li_mi(lmax, mmax, mres) + setfield!(cfg, :li, li) + setfield!(cfg, :mi, mi) + + nrm = getfield(cfg, :_norm) + nrm.Nlm = Nlm_table(lmax, mmax) + nrm.scale_matrix[] = Matrix{Float64}(undef, 0, 0) # rebuilt lazily at the new size + + tbl = getfield(cfg, :_tables) + tbl.enabled = false + tbl.plm = Matrix{Float64}[]; tbl.dplm = Matrix{Float64}[] + tbl.NP = Matrix{Float64}[]; tbl.NdP = Matrix{Float64}[] + + empty!(getfield(cfg, :_scratch).m_order) # rebuilt lazily by cached_m_order + return cfg +end + function Base.propertynames(::SHTConfig, private::Bool=false) return (fieldnames(SHTConfig)..., :θ, :φ, :x, :w, :st, :cphi, @@ -940,15 +1019,60 @@ The first nlat rows contain the actual data; remaining rows are padding. ```julia cfg = create_gauss_config(64, 66) set_allow_padding!(cfg) -field = allocate_padded_spatial(cfg) -# field has size (nlat_padded, nlon), use field[1:cfg.nlat, :] for data +field = allocate_padded_spatial(cfg) # (nlat_padded, nlon) +copy_to_padded!(field, data, cfg) +alm = analysis(cfg, spatial_view(cfg, field)) # transforms need exactly nlat rows ``` + +The transforms require an array with exactly `nlat` rows, so a padded buffer must +go through [`spatial_view`](@ref) (or an equivalent `view(field, 1:cfg.nlat, :)`). +The view keeps the padded column stride, preserving what the padding is for. """ function allocate_padded_spatial(cfg::SHTConfig, T::Type=Float64) nlat_p = get_nlat_padded(cfg) return zeros(T, nlat_p, cfg.nlon) end +""" + spatial_view(cfg::SHTConfig, A::AbstractArray) -> SubArray + +The `(nlat, nlon)` (or `(nlat, nlon, nfields)`) window of a padded spatial buffer, +in the form the transforms accept. + +`allocate_padded_spatial` deliberately returns an array whose first dimension is +`nlat_padded ≥ nlat`, while every transform requires exactly `nlat` rows — so the +padded buffer could not be handed to `analysis` at all, which left the whole +padding API unusable. This is the missing bridge: + +```julia +cfg = create_gauss_config(64, 66) +set_allow_padding!(cfg) +pad = allocate_padded_spatial(cfg) +copy_to_padded!(pad, field, cfg) +alm = analysis(cfg, spatial_view(cfg, pad)) # padded stride preserved +``` + +The view does not throw the padding away: a `SubArray` over the leading rows keeps +the padded column stride (`nlat_padded`, not `nlat`), which is exactly the +cache-conflict avoidance the padding exists for. Results are bit-identical to the +unpadded layout. +""" +function spatial_view(cfg::SHTConfig, A::AbstractMatrix) + size(A, 1) >= cfg.nlat || throw(DimensionMismatch( + "padded spatial array needs ≥ nlat=$(cfg.nlat) rows, got $(size(A, 1))")) + size(A, 2) == cfg.nlon || throw(DimensionMismatch( + "spatial array second dim must be nlon=$(cfg.nlon), got $(size(A, 2))")) + return view(A, 1:cfg.nlat, 1:cfg.nlon) +end + +function spatial_view(cfg::SHTConfig, A::AbstractArray{<:Any,3}) + size(A, 1) >= cfg.nlat || throw(DimensionMismatch( + "padded spatial batch needs ≥ nlat=$(cfg.nlat) rows, got $(size(A, 1))")) + size(A, 2) == cfg.nlon || throw(DimensionMismatch( + "spatial batch second dim must be nlon=$(cfg.nlon), got $(size(A, 2))")) + return view(A, 1:cfg.nlat, 1:cfg.nlon, :) +end + """ allocate_padded_spatial_batch(cfg::SHTConfig, nfields::Int, T::Type=Float64) -> Array{T} @@ -1018,6 +1142,27 @@ Fejér's first rule on `θ = (i+0.5)π/nlat` nodes by default; set By default associated Legendre tables are precomputed, which mirrors SHTns' regular-grid behaviour and improves performance. +!!! warning "Needs `nlat ≥ 2*lmax + 1` to be exact" + Fejér and Clenshaw–Curtis rules with `nlat` nodes integrate polynomials in + `x = cosθ` through degree `nlat - 1` exactly, and analysis integrates a + product of two degree-`lmax` Legendre functions — degree `2*lmax`. So these + grids are exact only from `nlat = 2*lmax + 1` upward, whereas Gauss–Legendre + needs just `nlat = lmax + 1`. + + Below that threshold `analysis(cfg, synthesis(cfg, alm))` does not return + `alm`. Measured relative error at `lmax = 8`: + + | `nlat` | `:regular` | `:regular_poles` | + |---:|---:|---:| + | 10 | 7.2e-2 | 3.6e-1 | + | 14 | 1.3e-2 | 7.1e-3 | + | 16 | 4.2e-3 | 1.3e-3 | + | 17 (`2*lmax+1`) | 7e-16 | 7e-16 | + + Nothing warns at run time, so size the grid deliberately: pass + `nlat ≥ 2*lmax + 1` here, or use [`create_gauss_config`](@ref) (exact at + `nlat = lmax + 1`) or `use_dh_weights=true` (exact at `nlat = 2*(lmax+1)`). + # Driscoll-Healy Quadrature Set `use_dh_weights=true` to use Driscoll-Healy quadrature for exact spherical @@ -1215,6 +1360,7 @@ function prepare_plm_tables!(cfg::SHTConfig) # Working arrays for computing one row at a time P = Vector{Float64}(undef, lmax + 1) # P̄_l^m(x) normalized values dPdtheta = Vector{Float64}(undef, lmax + 1) # dP̄_l^m/dθ normalized θ-derivatives + Pbuf = Vector{Float64}(undef, lmax + 2) # extended P̄ row reused across all nlat*(mmax+1) calls # Compute tables for each azimuthal order m using bounded normalized recurrence # plm_tables[m+1][l+1, i] = P̄_l^m(x_i) (orthonormal; no overflow at high lmax) @@ -1227,7 +1373,7 @@ function prepare_plm_tables!(cfg::SHTConfig) # Compute normalized Legendre polynomials at each latitude point for i in 1:nlat s_i = sqrt(max(0.0, 1.0 - cfg.x[i]^2)) - Plm_norm_and_dPdtheta_row!(P, dPdtheta, cfg.x[i], lmax, m) + Plm_norm_and_dPdtheta_row!(P, dPdtheta, cfg.x[i], lmax, m, Pbuf) # Store normalized values — no Nlm multiply needed (P̄ already = Nlm * rawP) @inbounds @views tbl[:, i] .= P # P̄_l^m(x_i) for l=0:lmax diff --git a/src/core_transforms.jl b/src/core_transforms.jl index 63876e6d..c77a2d23 100644 --- a/src/core_transforms.jl +++ b/src/core_transforms.jl @@ -591,7 +591,8 @@ function _adjoint_analysis(cfg::SHTConfig, Alm̄::AbstractMatrix; Fφ = Matrix{CT}(undef, nlat_local, nlon) fill!(Fφ, zero(eltype(Fφ))) lmax, mmax = cfg.lmax, cfg.mmax - φadj = 2π # nlon (ifft adjoint) × cphi (2π/nlon) = 2π + # nlon (ifft adjoint) × the analysis φ factor. = 2π under :dft. + φadj = cfg.nlon * _analysis_phi_scale(cfg) use_tbl = has_fused_scalar_tables(cfg) P = use_tbl ? nothing : Vector{Float64}(undef, lmax + 1) for m in 0:cfg.mres:mmax @@ -625,7 +626,7 @@ end """Scalar analysis orchestrator. Parallelizes Legendre integration over m-modes.""" function _analysis_scalar_mloop!(alm::AbstractMatrix, cfg::SHTConfig, Fph::AbstractMatrix) lmax, mmax = cfg.lmax, cfg.mmax - scale_phi = cfg.cphi + scale_phi = _analysis_phi_scale(cfg) # inverts synthesis under any phi_scale (= cphi under :dft) m_order = cached_m_order(cfg) if has_fused_scalar_tables(cfg) _analysis_scalar_mloop_tbl!(alm, cfg, Fph, m_order, scale_phi) diff --git a/src/energy_diagnostics.jl b/src/energy_diagnostics.jl index a526fd3f..8575b853 100644 --- a/src/energy_diagnostics.jl +++ b/src/energy_diagnostics.jl @@ -90,6 +90,9 @@ This represents the L² norm of the field, which is conserved under orthonormal spherical harmonic transforms (Parseval's identity). """ function energy_scalar(cfg::SHTConfig, alm::AbstractMatrix; real_field::Bool=true) + # The accumulation loop is `@inbounds`; without this a too-small matrix + # silently returned a wrong number instead of raising, unlike every sibling. + validate_spectral_dimensions(alm, cfg, "alm") lmax, mmax = cfg.lmax, cfg.mmax scale_matrix = _diagnostic_scale_matrix(cfg) # Type-stable accumulator (stays inferrable for Float32 / ForwardDiff.Dual inputs). @@ -109,6 +112,7 @@ For a vector field V = ∇×(T Y_l^m êᵣ) + ∇ₕ(S Y_l^m), the kinetic energ KE = (1/2) ∫ |V|² dΩ = (1/2) Σ [l(l+1)|S_lm|² + l(l+1)|T_lm|²] """ function energy_vector(cfg::SHTConfig, Slm::AbstractMatrix, Tlm::AbstractMatrix; real_field::Bool=true) + validate_spectral_pair_dimensions(Slm, Tlm, cfg, ("Slm", "Tlm")) lmax, mmax = cfg.lmax, cfg.mmax scale_matrix = _diagnostic_scale_matrix(cfg) E = zero(promote_type(Float64, real(float(eltype(Slm))), real(float(eltype(Tlm))))) diff --git a/src/fftutils.jl b/src/fftutils.jl index b553989a..c4ebc661 100644 --- a/src/fftutils.jl +++ b/src/fftutils.jl @@ -87,10 +87,77 @@ const _FFT_BACKEND_DFT = 2 # planning allocation cost without changing the public `SHTPlan` contract. # Strides are part of the key because views with the same shape may need # different FFTW plans. +# +# This is THE φ-FFT plan cache for the whole package: the serial transforms, the +# batch helpers and the distributed extension's per-rank FFTs all route through +# `fft_phi!`/`rfft_phi!`/`ifft_phi!`/`irfft_phi!` and therefore through here. The +# public `enable_fft_plan_cache!` / `disable_fft_plan_cache!` / `set_fft_plan_cache!` +# / `fft_plan_cache_enabled` knobs control THIS cache. const _LOCAL_FFT_PLAN_CACHE = Dict{Tuple{Symbol,DataType,NTuple{2,Int},NTuple{2,Int},Int}, Any}() const _LOCAL_FFT_PLAN_CACHE_LOCK = ReentrantLock() +# `SHTNSKIT_CACHE_PENCILFFTS` is the legacy spelling kept for scripts that already +# set it; `SHTNSKIT_FFT_PLAN_CACHE` is the current name. +const _FFT_PLAN_CACHE_ENABLED = + Ref(get(ENV, "SHTNSKIT_FFT_PLAN_CACHE", + get(ENV, "SHTNSKIT_CACHE_PENCILFFTS", "1")) != "0") + +# Soft cap on distinct cached plans. Transform sizes are few and stable in +# practice, but a long-lived process that sweeps many grid shapes would otherwise +# grow this dictionary without bound. Flush wholesale when the cap is reached. +const _FFT_PLAN_CACHE_MAX = + Ref(parse(Int, get(ENV, "SHTNSKIT_FFT_PLAN_CACHE_MAX", "64"))) + +""" + fft_plan_cache_enabled() -> Bool + +Whether φ-FFT plans are cached and reused across calls. +""" +fft_plan_cache_enabled() = _FFT_PLAN_CACHE_ENABLED[] + +""" + set_fft_plan_cache!(flag::Bool; clear::Bool=true) -> Bool + +Enable or disable φ-FFT plan caching. When disabling with `clear=true` (the +default) the existing entries are dropped as well. Returns `flag`. +""" +function set_fft_plan_cache!(flag::Bool; clear::Bool=true) + _FFT_PLAN_CACHE_ENABLED[] = flag + if !flag && clear + lock(_LOCAL_FFT_PLAN_CACHE_LOCK) do + empty!(_LOCAL_FFT_PLAN_CACHE) + end + end + return flag +end + +""" + enable_fft_plan_cache!() -> Bool + +Convenience wrapper for `set_fft_plan_cache!(true)`. +""" +enable_fft_plan_cache!() = set_fft_plan_cache!(true) + +""" + disable_fft_plan_cache!(; clear::Bool=true) -> Bool + +Convenience wrapper for `set_fft_plan_cache!(false; clear)`. +""" +disable_fft_plan_cache!(; clear::Bool=true) = set_fft_plan_cache!(false; clear=clear) + +""" + SHTnsKit.fft_plan_cache_max!(n::Int) -> Int + +Set the maximum number of cached φ-FFT plans; `n <= 0` removes the cap. +Returns the previous value. +""" +function fft_plan_cache_max!(n::Int) + prev = _FFT_PLAN_CACHE_MAX[] + _FFT_PLAN_CACHE_MAX[] = n + return prev +end + function fft_phi_backend() v = _FFT_BACKEND[] v == _FFT_BACKEND_FFTW && return :fftw @@ -102,32 +169,43 @@ end return (kind, eltype(A), (size(A, 1), size(A, 2)), (stride(A, 1), stride(A, 2)), nlon) end +"""Build one FFTW plan for `A`; see `_cached_local_fft_plan` for the flag rationale.""" +function _build_local_fft_plan(kind::Symbol, A::AbstractMatrix, nlon::Int) + # 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 + return if kind === :fft + plan_fft!(A, 2; flags) + elseif kind === :ifft + plan_ifft!(A, 2; flags) + elseif kind === :rfft + plan_rfft(A, 2; flags) + elseif kind === :irfft + plan_irfft(A, nlon, 2; flags) + else + throw(ArgumentError("unknown FFT plan kind: $kind")) + end +end + function _cached_local_fft_plan(kind::Symbol, A::AbstractMatrix, nlon::Int=0) + _FFT_PLAN_CACHE_ENABLED[] || return _build_local_fft_plan(kind, A, nlon) key = _local_fft_plan_key(kind, A, nlon) lock(_LOCAL_FFT_PLAN_CACHE_LOCK) 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) - elseif kind === :ifft - plan_ifft!(A, 2; flags) - elseif kind === :rfft - plan_rfft(A, 2; flags) - elseif kind === :irfft - plan_irfft(A, nlon, 2; flags) - else - throw(ArgumentError("unknown FFT plan kind: $kind")) + plan = _build_local_fft_plan(kind, A, nlon) + # Enforce the soft cap: flush before inserting so the fresh entry survives. + cap = _FFT_PLAN_CACHE_MAX[] + if cap > 0 && length(_LOCAL_FFT_PLAN_CACHE) >= cap + empty!(_LOCAL_FFT_PLAN_CACHE) end _LOCAL_FFT_PLAN_CACHE[key] = plan end diff --git a/src/layout.jl b/src/layout.jl index 4cec7b4a..5918eddb 100644 --- a/src/layout.jl +++ b/src/layout.jl @@ -227,17 +227,22 @@ the given packed index. Useful for algorithms that need to determine the azimuthal symmetry of a coefficient from its packed storage location. """ -function im_from_lm(lm::Int, lmax::Int, mres::Int) +function im_from_lm(lm::Int, lmax::Int, mres::Int; mmax::Int=lmax) # Validate packed index lm ≥ 0 || throw(ArgumentError("lm must be ≥ 0")) + (0 ≤ mmax ≤ lmax) || throw(ArgumentError("require 0 ≤ mmax ≤ lmax")) # Search through m-blocks to find the one containing this packed index im = 0 # Current reduced m-index being tested base = 0 # Base offset for current m-block - im_max = lmax ÷ mres # maximum valid reduced m-index + # Bound by the last block the LAYOUT actually has, which is set by mmax, not + # by lmax. Using `lmax ÷ mres` let a past-the-end index resolve to an order + # the configuration does not store (`im_from_lm(30, 8, 1)` returned 4 for an + # mmax=3 layout whose valid indices stop at 29) instead of raising. + im_max = mmax ÷ mres # maximum valid reduced m-index while true - im > im_max && throw(ArgumentError("lm=$lm is out of range for lmax=$lmax, mres=$mres")) + im > im_max && throw(ArgumentError("lm=$lm is out of range for lmax=$lmax, mmax=$mmax, mres=$mres")) # Size of current m-block (number of l-modes for this m) block = lmax - im*mres + 1 diff --git a/src/legendre.jl b/src/legendre.jl index c311a24f..e9a435c9 100644 --- a/src/legendre.jl +++ b/src/legendre.jl @@ -245,17 +245,46 @@ At poles (sinθ → 0): """ function Plm_norm_and_dPdtheta_row!(P::AbstractVector{T}, dPdtheta::AbstractVector{T}, x::T, lmax::Int, m::Int) where {T<:Real} - # Fill P with orthonormal P̄_l^m - Plm_norm_row!(P, x, lmax, m) + # Allocating fallback: provide a local Pbuf and delegate. + Pbuf = zeros(T, lmax + 2) + return Plm_norm_and_dPdtheta_row!(P, dPdtheta, x, lmax, m, Pbuf) +end + +""" + Plm_norm_and_dPdtheta_row!(P, dPdtheta, x, lmax, m, Pbuf) + +Buffer-taking form: `Pbuf` is caller-supplied scratch of length ≥ `lmax+2`, +avoiding a heap allocation per call in hot loops (`prepare_plm_tables!` calls +this `nlat*(mmax+1)` times). +""" +function Plm_norm_and_dPdtheta_row!(P::AbstractVector{T}, dPdtheta::AbstractVector{T}, + x::T, lmax::Int, m::Int, + Pbuf::AbstractVector{T}) where {T<:Real} + length(Pbuf) >= lmax + 2 || throw(ArgumentError("Pbuf must have length ≥ lmax+2")) @inbounds begin fill!(dPdtheta, zero(T)) - lmax < m && return P, dPdtheta sinth = sqrt(max(zero(T), one(T) - x*x)) + interior = lmax >= m && sinth >= POLE_TOLERANCE_FACTOR * eps(T) + + if interior + # The dθ recurrence needs P̄ up to lmax+1, and the first lmax+1 of + # those ARE the row we must return. Computing `P` with its own + # `Plm_norm_row!` call and then `Pbuf` with a second one ran the + # (inherently serial) recurrence twice for one row — about 40 % of + # this function's cost. Run it once into the longer buffer instead. + Plm_norm_row!(Pbuf, x, lmax + 1, m) + fill!(P, zero(T)) + copyto!(P, 1, Pbuf, 1, lmax + 1) + else + Plm_norm_row!(P, x, lmax, m) + end + + lmax < m && return P, dPdtheta # Handle poles (sinθ ≈ 0) - if sinth < POLE_TOLERANCE_FACTOR * eps(T) + if !interior if m == 1 # Analytic limit: dP̄_l^1/dθ|_{θ=0} = N_{l,1} * (−l(l+1)/2) # N_{l,1} = sqrt[(2l+1)/(4π) / (l*(l+1))], so @@ -278,11 +307,8 @@ function Plm_norm_and_dPdtheta_row!(P::AbstractVector{T}, dPdtheta::AbstractVect return P, dPdtheta end - # Standard case: use the normalized dθ recurrence. - # We need P̄_{l+1}^m for the l-term, so compute P̄ to lmax+1. - # Use a local buffer for the extended row. - Pbuf = zeros(T, lmax + 2) # indices 1..lmax+2 → degrees 0..lmax+1 - Plm_norm_row!(Pbuf, x, lmax + 1, m) + # Standard case: `Pbuf` already holds P̄ up to degree lmax+1 from the + # single recurrence pass above. # d(l,m) = sqrt((l^2 - m^2) / (4l^2 - 1)) [0 for l=0] inv_sinth = one(T) / sinth @@ -334,16 +360,29 @@ function Plm_norm_dPdtheta_over_sinth_row!(P::AbstractVector{T}, dPdtheta::Abstr x::T, lmax::Int, m::Int, Pbuf::AbstractVector{T}) where {T<:Real} length(Pbuf) >= lmax + 2 || throw(ArgumentError("Pbuf must have length ≥ lmax+2")) - Plm_norm_row!(P, x, lmax, m) @inbounds begin fill!(dPdtheta, zero(T)) fill!(P_over_sinth, zero(T)) - lmax < m && return P, dPdtheta, P_over_sinth sinth = sqrt(max(zero(T), one(T) - x*x)) + interior = lmax >= m && sinth >= POLE_TOLERANCE_FACTOR * eps(T) + + if interior + # Single recurrence pass: the dθ formula needs P̄ to degree lmax+1 and + # the first lmax+1 entries of that row ARE `P`. Filling `P` and `Pbuf` + # with two separate `Plm_norm_row!` calls ran the serial recurrence + # twice per row (~40 % of this function's cost). + Plm_norm_row!(Pbuf, x, lmax + 1, m) + fill!(P, zero(T)) + copyto!(P, 1, Pbuf, 1, lmax + 1) + else + Plm_norm_row!(P, x, lmax, m) + end - if sinth < POLE_TOLERANCE_FACTOR * eps(T) + lmax < m && return P, dPdtheta, P_over_sinth + + if !interior if m == 1 for l in 1:lmax mag = T(0.5) * T(_INV_SQRT_4PI) * sqrt(T(2l + 1) * T(l) * T(l + 1)) @@ -362,8 +401,8 @@ function Plm_norm_dPdtheta_over_sinth_row!(P::AbstractVector{T}, dPdtheta::Abstr return P, dPdtheta, P_over_sinth end - # Standard case: compute P̄ to lmax+1 for dθ recurrence - Plm_norm_row!(Pbuf, x, lmax + 1, m) + # Standard case: `Pbuf` already holds P̄ to degree lmax+1 from the single + # recurrence pass above. inv_sinth = one(T) / sinth diff --git a/src/local.jl b/src/local.jl index dcfc38c3..5e9ae307 100644 --- a/src/local.jl +++ b/src/local.jl @@ -68,6 +68,8 @@ function SH_to_lat(cfg::SHTConfig, Qlm::AbstractVector{<:Complex}, cost::Real; n vals[j+1] += 2 * real(gm * cis(PT(2π * m * j / nphi))) end end + sφ = _evaluator_phi_scale(cfg, eltype(vals)) # 1 under :dft; 1/2π under :quad + sφ == 1 || (vals .*= sφ) return vals end @@ -126,6 +128,8 @@ function SH_to_lat_cplx(cfg::SHTConfig, alm_packed::AbstractVector{<:Complex}, c vals[j+1] += gm * phase + gn * conj(phase) end end + sφ = _evaluator_phi_scale(cfg, eltype(vals)) # 1 under :dft; 1/2π under :quad + sφ == 1 || (vals .*= sφ) return vals end @@ -203,7 +207,11 @@ function SHqst_to_point(cfg::SHTConfig, Qlm::AbstractVector{<:Complex}, Slm::Abs vt *= sinth vp *= sinth end - return real(vr), real(vt), real(vp) + # Narrow the scale to the accumulator's own real type first: the untyped + # scale is Float64, and multiplying by it would widen a Float32 evaluation + # (and break the element type these functions promise their caller). + sφ = _evaluator_phi_scale(cfg, typeof(real(vr))) + return real(vr) * sφ, real(vt) * sφ, real(vp) * sφ end """ @@ -273,7 +281,11 @@ function SH_to_grad_point(cfg::SHTConfig, DrSlm::AbstractVector{<:Complex}, Slm: vt *= sinth vp *= sinth end - return real(vr), real(vt), real(vp) + # Narrow the scale to the accumulator's own real type first: the untyped + # scale is Float64, and multiplying by it would widen a Float32 evaluation + # (and break the element type these functions promise their caller). + sφ = _evaluator_phi_scale(cfg, typeof(real(vr))) + return real(vr) * sφ, real(vt) * sφ, real(vp) * sφ end """ @@ -367,6 +379,10 @@ function SHqst_to_lat(cfg::SHTConfig, Qlm::AbstractVector{<:Complex}, Slm::Abstr Vt .*= sinth Vp .*= sinth end + sφ = _evaluator_phi_scale(cfg, eltype(Vr)) # 1 under :dft; 1/2π under :quad + if sφ != 1 + Vr .*= sφ; Vt .*= sφ; Vp .*= sφ + end return Vr, Vt, Vp end diff --git a/src/plan.jl b/src/plan.jl index bab074ae..97aa886e 100644 --- a/src/plan.jl +++ b/src/plan.jl @@ -212,7 +212,7 @@ function analysis_sphtor!(plan::SHTPlan, Slm_out::AbstractMatrix, Tlm_out::Abstr size(Tlm_out,1)==cfg.lmax+1 && size(Tlm_out,2)==cfg.mmax+1 || throw(DimensionMismatch("Tlm_out dims")) lmax, mmax = cfg.lmax, cfg.mmax - scaleφ = cfg.cphi + scaleφ = _analysis_phi_scale(cfg) # inverts synthesis under any phi_scale (= cphi under :dft) fill!(Slm_out, zero(eltype(Slm_out))); fill!(Tlm_out, zero(eltype(Tlm_out))) # Two passes over (Vt, Vp): each packs the component into a real/complex @@ -421,7 +421,7 @@ function analysis!(plan::SHTPlan, alm_out::AbstractMatrix, f::AbstractMatrix) )) lmax, mmax = cfg.lmax, cfg.mmax - scaleφ = cfg.cphi + scaleφ = _analysis_phi_scale(cfg) # inverts synthesis under any phi_scale (= cphi under :dft) if plan.use_rfft eltype(f) <: Real || throw(ArgumentError("use_rfft plan requires real-valued f")) diff --git a/src/qst_transforms.jl b/src/qst_transforms.jl index 8610ef84..3f3432d6 100644 --- a/src/qst_transforms.jl +++ b/src/qst_transforms.jl @@ -331,12 +331,12 @@ function _synthesis_qst_l(cfg::SHTConfig, Qlm::AbstractMatrix, Slm::AbstractMatr end """ - analysis_qst_ml(cfg, im, Vr_m, Vt_m, Vp_m, ltr) -> (Ql, Sl, Tl) + analysis_qst_ml(cfg, mval, Vr_m, Vt_m, Vp_m, ltr) -> (Ql, Sl, Tl) -Mode-limited transform for specific azimuthal mode im. +Mode-limited transform for specific azimuthal mode mval. """ -function analysis_qst_ml(cfg::SHTConfig, im::Integer, Vr_m::AbstractVector{<:Complex}, Vt_m::AbstractVector{<:Complex}, Vp_m::AbstractVector{<:Complex}, ltr::Integer) - stored_im, _, lcap = _validate_stored_order(cfg, im, ltr) +function analysis_qst_ml(cfg::SHTConfig, mval::Integer, Vr_m::AbstractVector{<:Complex}, Vt_m::AbstractVector{<:Complex}, Vp_m::AbstractVector{<:Complex}, ltr::Integer) + stored_im, _, lcap = _validate_stored_order(cfg, mval, ltr) # Transform each component for this specific mode Ql = analysis_packed_ml(cfg, stored_im, Vr_m, lcap) Sl, Tl = analysis_sphtor_ml(cfg, stored_im, Vt_m, Vp_m, lcap) @@ -345,23 +345,23 @@ function analysis_qst_ml(cfg::SHTConfig, im::Integer, Vr_m::AbstractVector{<:Com return Ql, Sl, Tl end -function analysis_qst_ml(::CPU, cfg::SHTConfig, im::Integer, +function analysis_qst_ml(::CPU, cfg::SHTConfig, mval::Integer, Vr::AbstractVector{<:Complex}, Vt::AbstractVector{<:Complex}, Vp::AbstractVector{<:Complex}, ltr::Integer) for value in (Vr, Vt, Vp) _require_cpu_storage(:analysis_qst_ml, value) end - return analysis_qst_ml(cfg, im, Vr, Vt, Vp, ltr) + return analysis_qst_ml(cfg, mval, Vr, Vt, Vp, ltr) end """ - synthesis_qst_ml(cfg, im, Ql, Sl, Tl, ltr) -> (Vr_m, Vt_m, Vp_m) + synthesis_qst_ml(cfg, mval, Ql, Sl, Tl, ltr) -> (Vr_m, Vt_m, Vp_m) -Mode-limited synthesis for specific azimuthal mode im. +Mode-limited synthesis for specific azimuthal mode mval. """ -function synthesis_qst_ml(cfg::SHTConfig, im::Integer, Ql::AbstractVector{<:Complex}, Sl::AbstractVector{<:Complex}, Tl::AbstractVector{<:Complex}, ltr::Integer) - stored_im, _, lcap = _validate_stored_order(cfg, im, ltr) +function synthesis_qst_ml(cfg::SHTConfig, mval::Integer, Ql::AbstractVector{<:Complex}, Sl::AbstractVector{<:Complex}, Tl::AbstractVector{<:Complex}, ltr::Integer) + stored_im, _, lcap = _validate_stored_order(cfg, mval, ltr) # Each fixed-mode sub-transform converts its component to canonical once. Vr_m = synthesis_packed_ml(cfg, stored_im, Ql, lcap) Vt_m, Vp_m = synthesis_sphtor_ml(cfg, stored_im, Sl, Tl, lcap) @@ -369,12 +369,12 @@ function synthesis_qst_ml(cfg::SHTConfig, im::Integer, Ql::AbstractVector{<:Comp return Vr_m, Vt_m, Vp_m end -function synthesis_qst_ml(::CPU, cfg::SHTConfig, im::Integer, +function synthesis_qst_ml(::CPU, cfg::SHTConfig, mval::Integer, Q::AbstractVector{<:Complex}, S::AbstractVector{<:Complex}, Tlm::AbstractVector{<:Complex}, ltr::Integer) for value in (Q, S, Tlm) _require_cpu_storage(:synthesis_qst_ml, value) end - return synthesis_qst_ml(cfg, im, Q, S, Tlm, ltr) + return synthesis_qst_ml(cfg, mval, Q, S, Tlm, ltr) end diff --git a/src/rotations.jl b/src/rotations.jl index 05cb307f..1af0e325 100644 --- a/src/rotations.jl +++ b/src/rotations.jl @@ -97,6 +97,22 @@ Currently supports fast rotation around the Z-axis by angle `alpha` in radians. Rotate a real-field SH expansion around the Z-axis by angle `alpha`. Input and output are packed `Qlm` vectors (LM order, m ≥ 0). In-place supported if `Rlm === Qlm`. + +# Sign convention + +`R_lm = Q_lm · exp(-i m α)`. This is the **active** rotation of the field by `+α` +about `ẑ`: the rotated field is `g(θ, φ) = f(θ, φ - α)`. Equivalently, a feature +at longitude `φ₀` moves to `φ₀ + α`. + +Three things pin this sign and would all break if it were flipped to `+imα` +(which is the *passive* convention, `f(θ, φ + α)`): + + * the spatial rotation above, verified directly in + `test/serial/test_rotations.jl` against an FFT-grid φ shift; + * the general Wigner engine — `shtns_rotation_apply_real` with + `ZYZ(α, 0, 0)` (or `ZYZ(0, 0, α)`) must equal this function, and it builds + `diag(e^{-imα}) · d(β) · diag(e^{-imγ})`; + * the distributed twins `dist_SH_Zrotate` and every rotation `rrule`. """ function SH_Zrotate(::CPU, cfg::SHTConfig, Qlm::AbstractVector{<:Complex}, alpha::Real, Rlm::AbstractVector{<:Complex}) @@ -582,12 +598,61 @@ function _rotation_host_blocks(r::SHTRotation, ::Type{T}) where {T<:AbstractFloa return (; offsets, values, input_scales, output_scales, alpha=α, gamma=γ) end +""" + _require_full_m_range(r::SHTRotation, β::Real) + +Reject an order-mixing rotation on a layout that cannot hold every order it +produces. + +A Wigner-d rotation through a general `β` couples `Y_l^m` to every `Y_l^{m′}` +with `|m′| ≤ l`. If the storage stops at `mmax < lmax`, the `|m′| > mmax` +components have nowhere to go and were silently dropped — measured at +`lmax = 8`, that quietly discarded **14.8 %** of the field’s energy at +`mmax = 5` and **24.0 %** at `mmax = 3`, with no error and no warning. + +The two degenerate angles are exempt because their `d^l` is not order-mixing: +`β ≡ 0` is diagonal, and `β ≡ π` is anti-diagonal (`m′ = -m`), so `|m′| = |m|` +and a truncated layout still holds the result. That keeps pure Z-rotations +expressed as `ZYZ(α, 0, γ)` working at any `mmax`. + +Mirrors the `mres > 1` restriction stated by `dist_SH_Yrotate` and the packed +distributed rotations, for the same reason. +""" +function _require_full_m_range(r::SHTRotation, β::Real) + r.mmax >= r.lmax && return nothing + abs(sin(float(β))) <= 1e-12 && return nothing # β ≡ 0 (mod π): no m mixing + throw(ArgumentError( + "rotation with β=$(β) mixes azimuthal orders, but this configuration " * + "stores only m ≤ mmax=$(r.mmax) with lmax=$(r.lmax); the |m| > mmax " * + "components such a rotation generates cannot be represented and would " * + "be silently discarded. Use a configuration with mmax == lmax for " * + "Y/X rotations and general Euler angles. Pure Z-rotations (β ≡ 0 mod π) " * + "are unaffected and still work at any mmax.")) +end + +""" + _rotation_packed_length_check(v, expected, name, r) + +Length guard for the packed rotation inputs, with an `mres`-aware message. +""" +function _rotation_packed_length_check(v::AbstractVector, expected::Int, + name::AbstractString, r::SHTRotation) + length(v) == expected && return nothing + throw(DimensionMismatch( + "$name has length $(length(v)), expected $expected for the packed (mres=1) " * + "layout at lmax=$(r.lmax), mmax=$(r.mmax). A Y/X rotation mixes azimuthal " * + "orders, so it cannot be represented in an mres-strided layout at all — an " * + "mres>1 config produces a shorter packed vector and lands here. Use mres=1 " * + "for rotations other than SH_Zrotate.")) +end + function _rotation_apply_cplx_canonical!(r::SHTRotation, Zlm::AbstractVector{<:Complex}, Rlm::AbstractVector{<:Complex}) r.lmax ≥ 0 || return Rlm RT = typeof(real(zero(eltype(Rlm)))) α, β, γ = _rotation_zyz_angles(r, RT) + _require_full_m_range(r, β) # Pre-allocate working arrays at maximum size to avoid per-l allocations nmax = 2 * r.lmax + 1 @@ -692,8 +757,13 @@ function shtns_rotation_apply_real(::CPU, r::SHTRotation, _require_cpu_storage(:shtns_rotation_apply_real, Qlm) _require_cpu_storage(:shtns_rotation_apply_real, Rlm) expected = nlm_calc(r.lmax, r.mmax, 1) - length(Qlm) == expected || throw(DimensionMismatch("LM packed size mismatch")) - length(Rlm) == expected || throw(DimensionMismatch("LM packed size mismatch")) + # A length mismatch here is almost always an `mres > 1` config reaching a + # rotation that mixes orders, which no mres-strided layout can represent. + # Say that, rather than reporting a bare size mismatch the caller has to + # reverse-engineer. (`dist_SH_Yrotate` and the packed distributed rotations + # state the same restriction up front.) + _rotation_packed_length_check(Qlm, expected, "Qlm", r) + _rotation_packed_length_check(Rlm, expected, "Rlm", r) eltype(Qlm) === eltype(Rlm) || throw(ArgumentError( "rotation input and output element types must match", )) diff --git a/src/sphtor_transforms.jl b/src/sphtor_transforms.jl index 8458ad47..970b07f8 100644 --- a/src/sphtor_transforms.jl +++ b/src/sphtor_transforms.jl @@ -404,7 +404,7 @@ function _adjoint_analysis_sphtor(cfg::SHTConfig, Slm̄::AbstractMatrix, Tlm̄:: dPdtheta = Vector{Float64}(undef, lmax + 1) P_over_sinth = Vector{Float64}(undef, lmax + 1) Pbuf = Vector{Float64}(undef, lmax + 2) # scratch for extended P̄ row (avoids per-call alloc) - φadj = 2π + φadj = cfg.nlon * _analysis_phi_scale(cfg) # = 2π under :dft for m in 0:cfg.mres:mmax col = m + 1 @@ -533,7 +533,7 @@ function _analysis_sphtor_mloop!(Slm::AbstractMatrix, Tlm::AbstractMatrix, ltr::Int=cfg.lmax) lmax, mmax = cfg.lmax, cfg.mmax ltr_eff = min(ltr, lmax) - scale_phi = cfg.cphi + scale_phi = _analysis_phi_scale(cfg) # inverts synthesis under any phi_scale (= cphi under :dft) m_order = cached_m_order(cfg) if has_fused_vector_tables(cfg) @@ -997,39 +997,39 @@ synthesis_tor_l_cplx(::CPU, cfg::SHTConfig, Tlm::AbstractMatrix, synthesis_tor_l(CPU(), cfg, Tlm, ltr; real_output=false) """ - synthesis_sph_ml(cfg::SHTConfig, im::Int, Sl::AbstractVector{<:Complex}, ltr::Int) + synthesis_sph_ml(cfg::SHTConfig, mval::Int, Sl::AbstractVector{<:Complex}, ltr::Int) Mode-limited spheroidal-only synthesis wrapper. """ -function synthesis_sph_ml(cfg::SHTConfig, im::Integer, Sl::AbstractVector{<:Complex}, ltr::Integer) +function synthesis_sph_ml(cfg::SHTConfig, mval::Integer, Sl::AbstractVector{<:Complex}, ltr::Integer) # Mode-limited wrappers use zero-vector views for the missing component; - # this avoids an O(ltr-im) allocation on repeated per-mode calls. + # this avoids an O(ltr-mval) allocation on repeated per-mode calls. Tl_zero = _zero_spectral_vector(eltype(Sl), length(Sl)) - return synthesis_sphtor_ml(cfg, im, Sl, Tl_zero, ltr) + return synthesis_sphtor_ml(cfg, mval, Sl, Tl_zero, ltr) end -function synthesis_sph_ml(::CPU, cfg::SHTConfig, im::Integer, +function synthesis_sph_ml(::CPU, cfg::SHTConfig, mval::Integer, Sl::AbstractVector{<:Complex}, ltr::Integer) _require_cpu_storage(:synthesis_sph_ml, Sl) - return synthesis_sph_ml(cfg, im, Sl, ltr) + return synthesis_sph_ml(cfg, mval, Sl, ltr) end """ - synthesis_tor_ml(cfg::SHTConfig, im::Int, Tl::AbstractVector{<:Complex}, ltr::Int) + synthesis_tor_ml(cfg::SHTConfig, mval::Int, Tl::AbstractVector{<:Complex}, ltr::Int) Mode-limited toroidal-only synthesis wrapper. """ -function synthesis_tor_ml(cfg::SHTConfig, im::Integer, Tl::AbstractVector{<:Complex}, ltr::Integer) +function synthesis_tor_ml(cfg::SHTConfig, mval::Integer, Tl::AbstractVector{<:Complex}, ltr::Integer) # Mode-limited wrappers use zero-vector views for the missing component; - # this avoids an O(ltr-im) allocation on repeated per-mode calls. + # this avoids an O(ltr-mval) allocation on repeated per-mode calls. Sl_zero = _zero_spectral_vector(eltype(Tl), length(Tl)) - return synthesis_sphtor_ml(cfg, im, Sl_zero, Tl, ltr) + return synthesis_sphtor_ml(cfg, mval, Sl_zero, Tl, ltr) end -function synthesis_tor_ml(::CPU, cfg::SHTConfig, im::Integer, +function synthesis_tor_ml(::CPU, cfg::SHTConfig, mval::Integer, Tl::AbstractVector{<:Complex}, ltr::Integer) _require_cpu_storage(:synthesis_tor_ml, Tl) - return synthesis_tor_ml(cfg, im, Tl, ltr) + return synthesis_tor_ml(cfg, mval, Tl, ltr) end """ @@ -1081,7 +1081,7 @@ function analysis_sphtor_ml(cfg::SHTConfig, stored_im::Integer, Vt_m::AbstractVe dPdtheta = Vector{Float64}(undef, ltr + 1) P_over_sinth = Vector{Float64}(undef, ltr + 1) Pbuf = Vector{Float64}(undef, ltr + 2) # scratch for extended P̄ row (avoids per-call alloc) - scaleφ = cfg.cphi + scaleφ = _analysis_phi_scale(cfg) # inverts synthesis under any phi_scale (= cphi under :dft) # Integrate using Legendre polynomials and derivatives (pole-safe) for i in 1:nlat @@ -1229,12 +1229,12 @@ synthesis_grad_l(::CPU, cfg::SHTConfig, Slm::AbstractMatrix, synthesis_sph_l(CPU(), cfg, Slm, ltr; kwargs...) """ - synthesis_grad_ml(cfg::SHTConfig, im::Int, Sl::AbstractVector{<:Complex}, ltr::Int) + synthesis_grad_ml(cfg::SHTConfig, mval::Int, Sl::AbstractVector{<:Complex}, ltr::Int) Mode-limited gradient synthesis alias. """ -function synthesis_grad_ml(cfg::SHTConfig, im::Integer, Sl::AbstractVector{<:Complex}, ltr::Integer) - return synthesis_sph_ml(cfg, im, Sl, ltr) +function synthesis_grad_ml(cfg::SHTConfig, mval::Integer, Sl::AbstractVector{<:Complex}, ltr::Integer) + return synthesis_sph_ml(cfg, mval, Sl, ltr) end synthesis_grad_ml(::CPU, cfg::SHTConfig, im::Integer, diff --git a/src/transforms.jl b/src/transforms.jl index b495d146..08d71d71 100644 --- a/src/transforms.jl +++ b/src/transforms.jl @@ -108,7 +108,7 @@ function analysis_axisym(cfg::SHTConfig, Vr::AbstractVector{<:Real}) # 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 + scaleφ = _analysis_phi_scale(cfg) * cfg.nlon # inverts synthesis under any phi_scale (= cphi under :dft) @inbounds for l in 0:lmax Ql[l+1] *= scaleφ end @@ -239,6 +239,11 @@ function synthesis_axisym(cfg::SHTConfig, Qlm::AbstractVector{<:Complex}) Vr[i] = val end + # Same φ convention factor the full `synthesis` carries (1 under :dft; + # 1/2π under :quad), so this stays the m=0 column of `synthesis` and the + # exact inverse of `analysis_axisym` in either mode. + sφ = _evaluator_phi_scale(cfg) + sφ == 1 || (Vr .*= RT(sφ)) return Vr end @@ -270,7 +275,7 @@ function analysis_axisym_l(cfg::SHTConfig, Vr::AbstractVector{<:Real}, ltr::Inte end # Same φ quadrature factor as `analysis_axisym` — see the comment there. - scaleφ = cfg.cphi * cfg.nlon + scaleφ = _analysis_phi_scale(cfg) * cfg.nlon # inverts synthesis under any phi_scale (= cphi under :dft) @inbounds for l in eachindex(Ql) Ql[l] *= scaleφ end @@ -307,25 +312,28 @@ function synthesis_axisym_l(cfg::SHTConfig, Qlm::AbstractVector{<:Complex}, ltr: Vr[i] = val end + # Same φ convention factor as `synthesis_axisym` — see the comment there. + sφ = _evaluator_phi_scale(cfg) + sφ == 1 || (Vr .*= RT(sφ)) return Vr end """ - analysis_packed_ml(cfg, im, Vr_m, ltr) -> Vector{<:Complex} + analysis_packed_ml(cfg, mval, Vr_m, ltr) -> Vector{<:Complex} Transform spatial field for one stored azimuthal order to spherical harmonic -coefficients. `im` is the zero-based stored-order index, so the physical order -is `m = im * cfg.mres`. `Vr_m` contains complex spatial values for that mode. +coefficients. `mval` is the zero-based stored-order index, so the physical order +is `m = mval * cfg.mres`. `Vr_m` contains complex spatial values for that mode. Returns coefficients Q_l for degrees l = m..ltr. """ -function analysis_packed_ml(cfg::SHTConfig, im::Int, Vr_m::AbstractVector{<:Complex}, ltr::Integer) +function analysis_packed_ml(cfg::SHTConfig, mval::Int, Vr_m::AbstractVector{<:Complex}, ltr::Integer) nlat = cfg.nlat length(Vr_m) == nlat || throw(DimensionMismatch("Vr_m length must be nlat=$(nlat)")) - im >= 0 || throw(ArgumentError("im must be >= 0")) - im <= cfg.mmax ÷ cfg.mres || throw(ArgumentError("im must be <= mmax/mres=$(cfg.mmax ÷ cfg.mres)")) - m = im * cfg.mres + mval >= 0 || throw(ArgumentError("mval must be >= 0")) + mval <= cfg.mmax ÷ cfg.mres || throw(ArgumentError("mval must be <= mmax/mres=$(cfg.mmax ÷ cfg.mres)")) + m = mval * cfg.mres ltr = _validate_degree_limit(cfg, ltr) - ltr >= m || throw(ArgumentError("ltr must be >= im*mres=$(m)")) + ltr >= m || throw(ArgumentError("ltr must be >= mval*mres=$(m)")) num_l = ltr - m + 1 CT = complex(float(real(eltype(Vr_m)))) # AD/Float32-safe output eltype @@ -333,7 +341,7 @@ function analysis_packed_ml(cfg::SHTConfig, im::Int, Vr_m::AbstractVector{<:Comp fill!(Ql, zero(CT)) P = Vector{Float64}(undef, ltr + 1) - scaleφ = cfg.cphi # Match full transform normalization + scaleφ = _analysis_phi_scale(cfg) # Match full transform normalization xv = cfg.x; wv = cfg.w # hoist field reads out of the i/l loops (cfg is mutable, so not auto-hoisted) for i in 1:nlat @@ -352,20 +360,20 @@ function analysis_packed_ml(cfg::SHTConfig, im::Int, Vr_m::AbstractVector{<:Comp end """ - synthesis_packed_ml(cfg, im, Ql, ltr) -> Vector{<:Complex} + synthesis_packed_ml(cfg, mval, Ql, ltr) -> Vector{<:Complex} Transform spherical harmonic coefficients for specific mode m to spatial field. -`im` is the zero-based stored-order index (`m = im * cfg.mres`); `Ql` +`mval` is the zero-based stored-order index (`m = mval * cfg.mres`); `Ql` contains coefficients for degrees l = m..ltr. Returns complex spatial values for that azimuthal mode. """ -function synthesis_packed_ml(cfg::SHTConfig, im::Int, Ql::AbstractVector{<:Complex}, ltr::Integer) +function synthesis_packed_ml(cfg::SHTConfig, mval::Int, Ql::AbstractVector{<:Complex}, ltr::Integer) nlat = cfg.nlat - im >= 0 || throw(ArgumentError("im must be >= 0")) - im <= cfg.mmax ÷ cfg.mres || throw(ArgumentError("im must be <= mmax/mres=$(cfg.mmax ÷ cfg.mres)")) - m = im * cfg.mres + mval >= 0 || throw(ArgumentError("mval must be >= 0")) + mval <= cfg.mmax ÷ cfg.mres || throw(ArgumentError("mval must be <= mmax/mres=$(cfg.mmax ÷ cfg.mres)")) + m = mval * cfg.mres ltr = _validate_degree_limit(cfg, ltr) - ltr >= m || throw(ArgumentError("ltr must be >= im*mres=$(m)")) + ltr >= m || throw(ArgumentError("ltr must be >= mval*mres=$(m)")) expected_len = ltr - m + 1 length(Ql) == expected_len || throw(DimensionMismatch("Ql length must be $(expected_len)")) @@ -432,15 +440,15 @@ function synthesis_axisym_l(::CPU, cfg::SHTConfig, _require_cpu_storage(:synthesis_axisym_l, coefficients) return synthesis_axisym_l(cfg, coefficients, ltr) end -function analysis_packed_ml(::CPU, cfg::SHTConfig, im::Int, +function analysis_packed_ml(::CPU, cfg::SHTConfig, mval::Int, mode::AbstractVector{<:Complex}, ltr::Integer) _require_cpu_storage(:analysis_packed_ml, mode) - return analysis_packed_ml(cfg, im, mode, ltr) + return analysis_packed_ml(cfg, mval, mode, ltr) end -function synthesis_packed_ml(::CPU, cfg::SHTConfig, im::Int, +function synthesis_packed_ml(::CPU, cfg::SHTConfig, mval::Int, coefficients::AbstractVector{<:Complex}, ltr::Integer) _require_cpu_storage(:synthesis_packed_ml, coefficients) - return synthesis_packed_ml(cfg, im, coefficients, ltr) + return synthesis_packed_ml(cfg, mval, coefficients, ltr) end """ @@ -538,7 +546,10 @@ function synthesis_point(cfg::SHTConfig, Qlm::AbstractMatrix{<:Complex}, cost::R result += 2 * real(gm * phase) end - return result + # Same φ convention factor the grid `synthesis` carries (1 under :dft; + # 1/2π under :quad), so a point evaluation matches the grid it samples. + sφ = _evaluator_phi_scale(cfg, typeof(result)) + return sφ == 1 ? result : result * sφ end function synthesis_point(::CPU, cfg::SHTConfig, diff --git a/src/vorticity_diagnostics.jl b/src/vorticity_diagnostics.jl index 2c3a9c35..4b408d5a 100644 --- a/src/vorticity_diagnostics.jl +++ b/src/vorticity_diagnostics.jl @@ -210,7 +210,14 @@ function grad_loss_vorticity_Tlm(cfg::SHTConfig, Tlm::AbstractMatrix, ζ_target: gζlm = analysis(cfg, residual) # Analysis includes the loss's quadrature weights, but the adjoint also # carries synthesis's inverse-FFT scale (1 for :dft, 1/2π for :quad). - synthesis_scale = phi_inv_scale(cfg) / cfg.nlon + # `analysis` is used here as a stand-in for the synthesis adjoint: with + # σ = phi_inv_scale(cfg)/nlon the spatial factor of synthesis, and + # g = ∂L/∂ζ_grid = cphi·w·r, the true adjoint is + # ∂L/∂ζlm = σ · Λᴴ(g) = σ · analysis_dft(r), + # where `analysis_dft` is analysis with the fixed cphi factor. Analysis now + # carries cphi/σ instead, i.e. `analysis(r) = analysis_dft(r)/σ`, so the + # compensation picks up σ a second time. Exactly 1 in the default :dft mode. + synthesis_scale = (phi_inv_scale(cfg) / cfg.nlon)^2 # Apply chain rule: ∂L/∂T_lm = ∂L/∂ζ_lm * ∂ζ_lm/∂T_lm lmax, mmax = cfg.lmax, cfg.mmax @@ -218,7 +225,9 @@ function grad_loss_vorticity_Tlm(cfg::SHTConfig, Tlm::AbstractMatrix, ζ_target: gT = similar(Tlm) fill!(gT, 0.0) - for m in 0:mmax, l in max(1,m):lmax + # Stride by mres like every other diagnostic: orders that are not multiples + # of mres have no packed storage, so a gradient entry there is meaningless. + for m in 0:cfg.mres:mmax, l in max(1,m):lmax L2 = l * (l + 1) # Note: negative sign from ζ = -l(l+1)T gT[l+1, m+1] = -L2 * synthesis_scale * _convention_metric(scale_matrix, l, m) * gζlm[l+1, m+1] end @@ -239,13 +248,13 @@ function loss_and_grad_vorticity_Tlm(cfg::SHTConfig, Tlm::AbstractMatrix, ζ_tar # Backward pass for gradient gζlm = analysis(cfg, residual) - synthesis_scale = phi_inv_scale(cfg) / cfg.nlon + synthesis_scale = (phi_inv_scale(cfg) / cfg.nlon)^2 # see grad_loss_vorticity_Tlm lmax, mmax = cfg.lmax, cfg.mmax scale_matrix = _diagnostic_scale_matrix(cfg) gT = similar(Tlm) fill!(gT, 0.0) - for m in 0:mmax, l in max(1,m):lmax + for m in 0:cfg.mres:mmax, l in max(1,m):lmax L2 = l * (l + 1) gT[l+1, m+1] = -L2 * synthesis_scale * _convention_metric(scale_matrix, l, m) * gζlm[l+1, m+1] end diff --git a/test/fixtures/compatibility/host_transfer_allowlist.toml b/test/fixtures/compatibility/host_transfer_allowlist.toml index b0dab7e9..d000b082 100644 --- a/test/fixtures/compatibility/host_transfer_allowlist.toml +++ b/test/fixtures/compatibility/host_transfer_allowlist.toml @@ -1,5 +1,5 @@ [audit] -entry_count = 685 +entry_count = 676 scanner = "deterministic non-overlapping spelling scanner; longer transfer spellings win" scope = "src/device_utils.jl and every ext/*.jl file" @@ -4229,6 +4229,30 @@ snippet_sha256 = "11154a1e24a1c10bab3991372fe0861ea6f27729071474ecfaf5c7b46aaa64 token = "copy" [[entry]] classification = "cpu_only" +key = "ext/SHTnsKitAdvancedADExt.jl|copy|4f6fc4b2fb87282c2fdcfdb386afaac0c96c5d90f06c74191664013078bf75d7|1" +path = "ext/SHTnsKitAdvancedADExt.jl" +reason = "CPU-only automatic-differentiation storage; vendor PencilArray inputs must be rejected before forward work or materialization" +same_snippet_ordinal = 1 +snippet_sha256 = "4f6fc4b2fb87282c2fdcfdb386afaac0c96c5d90f06c74191664013078bf75d7" +token = "copy" +[[entry]] +classification = "cpu_only" +key = "ext/SHTnsKitAdvancedADExt.jl|copy|fcc1f975f3afbc60122289577da76a1629d323077ea5611f3f89d25ba94c433d|1" +path = "ext/SHTnsKitAdvancedADExt.jl" +reason = "CPU-only automatic-differentiation storage; vendor PencilArray inputs must be rejected before forward work or materialization" +same_snippet_ordinal = 1 +snippet_sha256 = "fcc1f975f3afbc60122289577da76a1629d323077ea5611f3f89d25ba94c433d" +token = "copy" +[[entry]] +classification = "cpu_only" +key = "ext/SHTnsKitAdvancedADExt.jl|copy|fcc1f975f3afbc60122289577da76a1629d323077ea5611f3f89d25ba94c433d|2" +path = "ext/SHTnsKitAdvancedADExt.jl" +reason = "CPU-only automatic-differentiation storage; vendor PencilArray inputs must be rejected before forward work or materialization" +same_snippet_ordinal = 2 +snippet_sha256 = "fcc1f975f3afbc60122289577da76a1629d323077ea5611f3f89d25ba94c433d" +token = "copy" +[[entry]] +classification = "cpu_only" key = "ext/SHTnsKitAdvancedADExt.jl|typed_array|6a3f945a436cfadc7d0d88754bc5e5305471ff778353833a998abc86020f20a3|1" path = "ext/SHTnsKitAdvancedADExt.jl" reason = "CPU-only automatic-differentiation storage; vendor PencilArray inputs must be rejected before forward work or materialization" @@ -4517,6 +4541,14 @@ snippet_sha256 = "0678e972f4053722b1aa5d0b265619b74a559c977aeda46711ae43d0c07bac token = "matrix" [[entry]] classification = "cpu_only" +key = "ext/SHTnsKitLoopVecExt.jl|matrix|2c5b53249ab10c4b6f67ff033f0d9e1e91daed69ccbadac17c7c37c0a6bc697a|1" +path = "ext/SHTnsKitLoopVecExt.jl" +reason = "explicit host workspace in a CPU-only algorithm; never route a vendor-backed mathematical result through this allocation" +same_snippet_ordinal = 1 +snippet_sha256 = "2c5b53249ab10c4b6f67ff033f0d9e1e91daed69ccbadac17c7c37c0a6bc697a" +token = "matrix" +[[entry]] +classification = "cpu_only" key = "ext/SHTnsKitLoopVecExt.jl|matrix|307d2a260f79d8bc361a96501373d64779ff7e0904c01cbef4e95780cb11c72a|1" path = "ext/SHTnsKitLoopVecExt.jl" reason = "explicit host workspace in a CPU-only algorithm; never route a vendor-backed mathematical result through this allocation" @@ -5188,22 +5220,6 @@ same_snippet_ordinal = 1 snippet_sha256 = "a24f0dcc14d951e081d408beb6d1acd0f4b8f86eba91bd1f035e69f9b399a2f1" token = "vector" [[entry]] -classification = "metadata_or_storage_preserving" -key = "ext/SHTnsKitParallelExt.jl|collect|8433b6662c3ee812f9d5d61be28a22a4c30b89340af8dab6e7d989764d414287|1" -path = "ext/SHTnsKitParallelExt.jl" -reason = "metadata collection or storage-preserving parent/copy operation; verify that the destination retains the source storage family" -same_snippet_ordinal = 1 -snippet_sha256 = "8433b6662c3ee812f9d5d61be28a22a4c30b89340af8dab6e7d989764d414287" -token = "collect" -[[entry]] -classification = "metadata_or_storage_preserving" -key = "ext/SHTnsKitParallelExt.jl|collect|ee6ffc0f24762d622b08b06e934ab308c06696b4f9772672f4be640735ef83a5|1" -path = "ext/SHTnsKitParallelExt.jl" -reason = "metadata collection or storage-preserving parent/copy operation; verify that the destination retains the source storage family" -same_snippet_ordinal = 1 -snippet_sha256 = "ee6ffc0f24762d622b08b06e934ab308c06696b4f9772672f4be640735ef83a5" -token = "collect" -[[entry]] classification = "cpu_only" key = "ext/SHTnsKitParallelExt.jl|matrix|023d49d77839bf58bb374b6d6f6731bde879032d5c394bb3724750c7449bc2f0|1" path = "ext/SHTnsKitParallelExt.jl" @@ -5221,22 +5237,6 @@ snippet_sha256 = "1ac610d54c86c71678c376f4d15e96d723b3620896bbbec9914099de0a5527 token = "matrix" [[entry]] classification = "cpu_only" -key = "ext/SHTnsKitParallelExt.jl|matrix|1ac610d54c86c71678c376f4d15e96d723b3620896bbbec9914099de0a552729|2" -path = "ext/SHTnsKitParallelExt.jl" -reason = "explicit host workspace in a CPU-only algorithm; never route a vendor-backed mathematical result through this allocation" -same_snippet_ordinal = 2 -snippet_sha256 = "1ac610d54c86c71678c376f4d15e96d723b3620896bbbec9914099de0a552729" -token = "matrix" -[[entry]] -classification = "cpu_only" -key = "ext/SHTnsKitParallelExt.jl|matrix|5ded71272574f235927bbe03113b5064dc30aef85f7c63746d6cf16394d8c7eb|1" -path = "ext/SHTnsKitParallelExt.jl" -reason = "explicit host workspace in a CPU-only algorithm; never route a vendor-backed mathematical result through this allocation" -same_snippet_ordinal = 1 -snippet_sha256 = "5ded71272574f235927bbe03113b5064dc30aef85f7c63746d6cf16394d8c7eb" -token = "matrix" -[[entry]] -classification = "cpu_only" key = "ext/SHTnsKitParallelExt.jl|matrix|ee8d49b69190631fea6f53ce92fca583041b30b039bf9433f864b4b17fe0e0cb|1" path = "ext/SHTnsKitParallelExt.jl" reason = "explicit host workspace in a CPU-only algorithm; never route a vendor-backed mathematical result through this allocation" @@ -5244,22 +5244,6 @@ same_snippet_ordinal = 1 snippet_sha256 = "ee8d49b69190631fea6f53ce92fca583041b30b039bf9433f864b4b17fe0e0cb" token = "matrix" [[entry]] -classification = "cpu_only" -key = "ext/SHTnsKitParallelExt.jl|matrix|f93b14319dca66cd64aefa6f03aef33dd45906d1efec458970c7d8b7686b8c45|1" -path = "ext/SHTnsKitParallelExt.jl" -reason = "explicit host workspace in a CPU-only algorithm; never route a vendor-backed mathematical result through this allocation" -same_snippet_ordinal = 1 -snippet_sha256 = "f93b14319dca66cd64aefa6f03aef33dd45906d1efec458970c7d8b7686b8c45" -token = "matrix" -[[entry]] -classification = "metadata_or_storage_preserving" -key = "ext/SHTnsKitParallelExt.jl|parent|58349167721d1a9bba7d82a45f001f1713623da1b1d945cd7fd7b5d2220272b0|1" -path = "ext/SHTnsKitParallelExt.jl" -reason = "metadata collection or storage-preserving parent/copy operation; verify that the destination retains the source storage family" -same_snippet_ordinal = 1 -snippet_sha256 = "58349167721d1a9bba7d82a45f001f1713623da1b1d945cd7fd7b5d2220272b0" -token = "parent" -[[entry]] classification = "metadata_or_storage_preserving" key = "ext/SHTnsKitParallelExt.jl|parent|6c8ce9fa30f5cb7ab5f6604ca8adc56134475f81939910b6f8673263e5727b5d|1" path = "ext/SHTnsKitParallelExt.jl" @@ -5276,54 +5260,6 @@ same_snippet_ordinal = 1 snippet_sha256 = "82071ef968af13d92eff3e6d2972b5e8cfb8a6c508d32f3442eed291de6c3608" token = "parent" [[entry]] -classification = "metadata_or_storage_preserving" -key = "ext/SHTnsKitParallelExt.jl|parent|afdf86c69aa5206e89b5435a67f6b02a00fd99b8f9037bc58414169e05fef202|1" -path = "ext/SHTnsKitParallelExt.jl" -reason = "metadata collection or storage-preserving parent/copy operation; verify that the destination retains the source storage family" -same_snippet_ordinal = 1 -snippet_sha256 = "afdf86c69aa5206e89b5435a67f6b02a00fd99b8f9037bc58414169e05fef202" -token = "parent" -[[entry]] -classification = "metadata_or_storage_preserving" -key = "ext/SHTnsKitParallelExt.jl|parent|afdf86c69aa5206e89b5435a67f6b02a00fd99b8f9037bc58414169e05fef202|2" -path = "ext/SHTnsKitParallelExt.jl" -reason = "metadata collection or storage-preserving parent/copy operation; verify that the destination retains the source storage family" -same_snippet_ordinal = 2 -snippet_sha256 = "afdf86c69aa5206e89b5435a67f6b02a00fd99b8f9037bc58414169e05fef202" -token = "parent" -[[entry]] -classification = "metadata_or_storage_preserving" -key = "ext/SHTnsKitParallelExt.jl|parent|afdf86c69aa5206e89b5435a67f6b02a00fd99b8f9037bc58414169e05fef202|3" -path = "ext/SHTnsKitParallelExt.jl" -reason = "metadata collection or storage-preserving parent/copy operation; verify that the destination retains the source storage family" -same_snippet_ordinal = 3 -snippet_sha256 = "afdf86c69aa5206e89b5435a67f6b02a00fd99b8f9037bc58414169e05fef202" -token = "parent" -[[entry]] -classification = "metadata_or_storage_preserving" -key = "ext/SHTnsKitParallelExt.jl|parent|b026f8b420c9cf2cebde304c52634b3b660d8980c84d292bcb38f0ef35a921fb|1" -path = "ext/SHTnsKitParallelExt.jl" -reason = "metadata collection or storage-preserving parent/copy operation; verify that the destination retains the source storage family" -same_snippet_ordinal = 1 -snippet_sha256 = "b026f8b420c9cf2cebde304c52634b3b660d8980c84d292bcb38f0ef35a921fb" -token = "parent" -[[entry]] -classification = "metadata_or_storage_preserving" -key = "ext/SHTnsKitParallelExt.jl|parent|da2163864051cb7dbb4972619469a9e138391299cb05766ebabaacd422148cc8|1" -path = "ext/SHTnsKitParallelExt.jl" -reason = "metadata collection or storage-preserving parent/copy operation; verify that the destination retains the source storage family" -same_snippet_ordinal = 1 -snippet_sha256 = "da2163864051cb7dbb4972619469a9e138391299cb05766ebabaacd422148cc8" -token = "parent" -[[entry]] -classification = "metadata_or_storage_preserving" -key = "ext/SHTnsKitParallelExt.jl|parent|e7146bfbd323b0a2294e29c6c4cd684d22e89b8cabb06f79d782705e4548fc65|1" -path = "ext/SHTnsKitParallelExt.jl" -reason = "metadata collection or storage-preserving parent/copy operation; verify that the destination retains the source storage family" -same_snippet_ordinal = 1 -snippet_sha256 = "e7146bfbd323b0a2294e29c6c4cd684d22e89b8cabb06f79d782705e4548fc65" -token = "parent" -[[entry]] classification = "cpu_only" key = "ext/SHTnsKitParallelExt.jl|vector|4fa452b93b779a5d536c482a3ffcf627ff0134c6afb925a707fc312f872d483a|1" path = "ext/SHTnsKitParallelExt.jl" @@ -5349,14 +5285,6 @@ snippet_sha256 = "7f628a4fb8031eb06efa9ec7c386e111cfe896aabecc8c954a393f02c5b975 token = "vector" [[entry]] classification = "cpu_only" -key = "ext/SHTnsKitParallelExt.jl|vector|8433b6662c3ee812f9d5d61be28a22a4c30b89340af8dab6e7d989764d414287|1" -path = "ext/SHTnsKitParallelExt.jl" -reason = "explicit host workspace in a CPU-only algorithm; never route a vendor-backed mathematical result through this allocation" -same_snippet_ordinal = 1 -snippet_sha256 = "8433b6662c3ee812f9d5d61be28a22a4c30b89340af8dab6e7d989764d414287" -token = "vector" -[[entry]] -classification = "cpu_only" key = "ext/SHTnsKitParallelExt.jl|vector|a93a2bfa0355e3c7cd001f93b48c1ff7a08709afb7fb408782aedb91c19250dc|1" path = "ext/SHTnsKitParallelExt.jl" reason = "explicit host workspace in a CPU-only algorithm; never route a vendor-backed mathematical result through this allocation" @@ -5381,19 +5309,19 @@ snippet_sha256 = "a93a2bfa0355e3c7cd001f93b48c1ff7a08709afb7fb408782aedb91c19250 token = "vector" [[entry]] classification = "cpu_only" -key = "ext/SHTnsKitParallelExt.jl|vector|ee6ffc0f24762d622b08b06e934ab308c06696b4f9772672f4be640735ef83a5|1" -path = "ext/SHTnsKitParallelExt.jl" -reason = "explicit host workspace in a CPU-only algorithm; never route a vendor-backed mathematical result through this allocation" +key = "ext/SHTnsKitZygoteExt.jl|copy|11154a1e24a1c10bab3991372fe0861ea6f27729071474ecfaf5c7b46aaa6421|1" +path = "ext/SHTnsKitZygoteExt.jl" +reason = "CPU-only automatic-differentiation storage; vendor PencilArray inputs must be rejected before forward work or materialization" same_snippet_ordinal = 1 -snippet_sha256 = "ee6ffc0f24762d622b08b06e934ab308c06696b4f9772672f4be640735ef83a5" -token = "vector" +snippet_sha256 = "11154a1e24a1c10bab3991372fe0861ea6f27729071474ecfaf5c7b46aaa6421" +token = "copy" [[entry]] classification = "cpu_only" -key = "ext/SHTnsKitZygoteExt.jl|copy|11154a1e24a1c10bab3991372fe0861ea6f27729071474ecfaf5c7b46aaa6421|1" +key = "ext/SHTnsKitZygoteExt.jl|copy|fcc1f975f3afbc60122289577da76a1629d323077ea5611f3f89d25ba94c433d|1" path = "ext/SHTnsKitZygoteExt.jl" reason = "CPU-only automatic-differentiation storage; vendor PencilArray inputs must be rejected before forward work or materialization" same_snippet_ordinal = 1 -snippet_sha256 = "11154a1e24a1c10bab3991372fe0861ea6f27729071474ecfaf5c7b46aaa6421" +snippet_sha256 = "fcc1f975f3afbc60122289577da76a1629d323077ea5611f3f89d25ba94c433d" token = "copy" [[entry]] classification = "cpu_only" diff --git a/test/fixtures/compatibility/task16_gate.toml b/test/fixtures/compatibility/task16_gate.toml index 0d3398ae..b36bfb0e 100644 --- a/test/fixtures/compatibility/task16_gate.toml +++ b/test/fixtures/compatibility/task16_gate.toml @@ -6,13 +6,13 @@ platform = "aarch64-apple-darwin" [local_gate] evidence_kind = "local_readiness" certifying = false -audited_tree_digest = "1bae4f3a0acc933d30369fe3e11d2333379b40dbc55e4dbadf9c705bb860cbd3" +audited_tree_digest = "34289cc58102e077f0a420f00eafab9f7c3629255ebfc5ab9c749e972fc10f44" audited_scope = ["Project.toml", ".github/workflows", "docs/src/shtns37-parity.md", "ext", "src", "test"] audited_exclusions = ["test/fixtures/compatibility/task16_gate.toml", "test/fixtures/compatibility/task16_local_commands.txt", "test/fixtures/compatibility/task16_local_summary.log"] commands_file = "test/fixtures/compatibility/task16_local_commands.txt" -commands_sha256 = "524f0421f4f3aa013f8e4805b5ae2a0d5ffe5c858ff8140f062a77e6c8447712" +commands_sha256 = "a9c74fca5943848f977389ff4519e17125037f1b0cae8b3e20f59e8942573cf5" summary_log = "test/fixtures/compatibility/task16_local_summary.log" -summary_log_sha256 = "c41fe35538fca5cd7b927bbff92c39412816c5f46339fba4a3ab5da518a59818" +summary_log_sha256 = "eaa7e8adbbed059f6604bfa1dfc4e4c45c86eef17703869b1aade77b3aa3794d" # Historical readiness counts; current review runs are recorded in the summary. pkg_exit_code = 0 serial_pass = 76347 @@ -26,7 +26,7 @@ mpi_cpu_fixture_pass_per_rank = 66 [host_transfer_audit] allowlist = "test/fixtures/compatibility/host_transfer_allowlist.toml" scanner = "test/support/host_transfer_inventory.jl" -entry_count = 685 +entry_count = 676 allowscalar_match_count = 0 legacy_cuda_host_result_bridges = 2 diff --git a/test/fixtures/compatibility/task16_local_commands.txt b/test/fixtures/compatibility/task16_local_commands.txt index b74eb75a..861917fd 100644 --- a/test/fixtures/compatibility/task16_local_commands.txt +++ b/test/fixtures/compatibility/task16_local_commands.txt @@ -5,7 +5,7 @@ certifying=false manual_cuda_tool_command=python3 -B -m unittest discover -s test/support -p test_manual_cuda.py -v manual_cuda_contract_command=/Users/subha/.julia/juliaup/julia-1.12.4+0.aarch64.apple.darwin14/bin/julia --startup-file=no --project=. test/serial/test_shtns37_contract.jl manual_cuda_numerical_source_unchanged_from=3b8b2fe118900f2431208088932631db4a162908 -audited_tree_digest=1bae4f3a0acc933d30369fe3e11d2333379b40dbc55e4dbadf9c705bb860cbd3 +audited_tree_digest=34289cc58102e077f0a420f00eafab9f7c3629255ebfc5ab9c749e972fc10f44 audited_scope=Project.toml,.github/workflows,docs/src/shtns37-parity.md,ext,src,test audited_exclusions=test/fixtures/compatibility/task16_gate.toml,test/fixtures/compatibility/task16_local_commands.txt,test/fixtures/compatibility/task16_local_summary.log pkg_command=/usr/bin/env JULIA_DEPOT_PATH=/private/tmp/julia_depot_shtnskit:/Users/subha/.julia SHTNSKIT_RUN_JET_TESTS=1 SHTNSKIT_RUN_AQUA_TESTS=1 /Users/subha/.julia/juliaup/julia-1.12.4+0.aarch64.apple.darwin14/bin/julia --startup-file=no --project=. -e 'using Pkg; Pkg.test()' @@ -67,3 +67,26 @@ qst_communicator_full_parity_file=test/parity/runtests_mpi.jl with no ARGS qst_communicator_mpi_timeout_seconds=900 per test file, including JIT compilation qst_communicator_audit_command=using SHTnsKit; include("test/serial/test_final_parity_gate.jl") qst_communicator_full_parity_ranks=2,4 + +# Convention/AD/dead-cache audit integration, 2026-09-13. Historical and prior +# review runs above are NOT re-runs; the keys below record only what was +# executed for this change. +audit2_date=2026-09-13 +audit2_base_commit=f06d13d1 +audit2_julia=/Users/subha/.julia/juliaup/julia-1.11.1+0.aarch64.apple.darwin14/bin/julia +audit2_julia_version=1.11.1 +audit2_platform=arm64-apple-darwin22.4.0 +audit2_threads=1 +audit2_serial_env=/private/tmp/shtkit_testenv (Test, Random, ForwardDiff, Zygote, ChainRulesCore, LoopVectorization; SHTnsKit dev'd from the worktree) +audit2_mpi_env=/private/tmp/shtkit_mpienv (MPI, PencilArrays, PencilFFTs; SHTnsKit dev'd from the worktree) +audit2_serial_command=julia --startup-file=no --project=/private/tmp/shtkit_testenv test/serial/runtests.jl +audit2_parallel_grid_command=julia --startup-file=no --project=/private/tmp/shtkit_mpienv test/parallel/runtests.jl +audit2_mpi_launcher=MPI.mpiexec() from the MPI env with the same Julia executable and project +audit2_mpi_files_4rank=test/parallel/test_mpi_comprehensive.jl,test/parallel/test_mpi_extended.jl,test/parallel/test_mpi_audit_fixes.jl,test/parallel/test_mpi_2d_alignment.jl +audit2_mpi_files_2rank=test/parallel/test_mpi_parallel_local_correctness.jl,test/parallel/test_mpi_comm_cleanup.jl,test/parallel/test_mpi_plan_preflight.jl,test/parallel/test_mpi_transpose_operand_preflight.jl,test/parallel/test_transpose_sht.jl,test/parallel/test_disttranspose_dealiased.jl,test/parallel/test_dist_plan_alloc.jl +audit2_host_transfer_regen_command=julia --startup-file=no --project=/private/tmp/shtkit_testenv test/support/generate_host_transfer_allowlist.jl +audit2_pkg_test_rerun=false +audit2_jet_rerun=false +audit2_aqua_rerun=false +audit2_gpu_rerun=false +audit2_parity_runtests_rerun=false diff --git a/test/fixtures/compatibility/task16_local_summary.log b/test/fixtures/compatibility/task16_local_summary.log index 0524c64e..16e75ea0 100644 --- a/test/fixtures/compatibility/task16_local_summary.log +++ b/test/fixtures/compatibility/task16_local_summary.log @@ -9,7 +9,7 @@ manual_cuda_no_hardware_guard_rejected=true manual_cuda_julia_syntax_pass=3 manual_cuda_github_status_published=false manual_cuda_numerical_source_unchanged_from=3b8b2fe118900f2431208088932631db4a162908 -audited_tree_digest=1bae4f3a0acc933d30369fe3e11d2333379b40dbc55e4dbadf9c705bb860cbd3 +audited_tree_digest=34289cc58102e077f0a420f00eafab9f7c3629255ebfc5ab9c749e972fc10f44 platform=aarch64-apple-darwin julia=1.12.4 pkg_exit_code=0 @@ -107,3 +107,41 @@ qst_communicator_full_parity_ranks=2,4 qst_communicator_full_parity_exit_code=0 qst_communicator_qst_native_pass_per_rank=55 qst_communicator_shtns37_fixture_pass_per_rank=66 + +# Convention/AD/dead-cache audit integration results, 2026-09-13. Results above +# are historical and were not re-run. +audit2_date=2026-09-13 +audit2_base_commit=f06d13d1 +audit2_julia=1.11.1 +audit2_platform=arm64-apple-darwin22.4.0 +audit2_threads=1 +audit2_serial_pass_threads_4=76848 +audit2_serial_threads_4_exit_code=0 +audit2_serial_pass=76848 +audit2_serial_fail=0 +audit2_serial_exit_code=0 +audit2_parallel_grid_pass=1714 +audit2_parallel_grid_exit_code=0 +audit2_mpi_ranks_4=4 +audit2_mpi_2d_alignment_pass_per_rank=15 +audit2_mpi_audit_fixes_pass_per_rank=41 +audit2_mpi_comprehensive_result=passed +audit2_mpi_extended_result=passed +audit2_mpi_ranks_2=2 +audit2_mpi_parallel_local_pass_per_rank=172 +audit2_mpi_comm_cleanup_pass_per_rank=4 +audit2_mpi_plan_preflight_pass_per_rank=8 +audit2_mpi_transpose_operand_preflight_pass_per_rank=25 +audit2_mpi_exit_code=0 +audit2_host_transfer_entry_count=676 +audit2_host_transfer_regenerated=true +audit2_pkg_test_rerun=false +audit2_jet_rerun=false +audit2_aqua_rerun=false +audit2_gpu_rerun=false +audit2_physical_gpu_execution=false +audit2_parity_cpu_command=julia --startup-file=no --project= test/parity/runtests_cpu.jl +audit2_parity_cpu_local_evaluation_pass=2212 +audit2_parity_cpu_shtns37_manifest_pass=1082 +audit2_parity_mpi_rerun=false +audit2_note=Run on Julia 1.11.1, not the 1.12.4 of the historical readiness runs; no CUDA or ROCm hardware is present on this machine. diff --git a/test/parallel/runtests.jl b/test/parallel/runtests.jl index 7a80e252..7cb96935 100644 --- a/test/parallel/runtests.jl +++ b/test/parallel/runtests.jl @@ -13,6 +13,7 @@ # - test_disttranspose_dealiased.jl : DistTransposePlan on dealiased nlon>2*mmax+1 (mpiexec) # - test_dist_plan_alloc.jl : DistAnalysisPlan correctness + per-call allocation budget (mpiexec) # - test_mpi_audit_fixes.jl : Regressions for the 2026-08 distributed audit fixes (mpiexec) +# - test_mpi_2d_alignment.jl : 2D spectral-plan alignment preconditions (mpiexec, 4 ranks) # # To run MPI tests: # mpiexec -n 4 julia --project test/parallel/test_mpi_comprehensive.jl @@ -21,6 +22,7 @@ # mpiexec -n 2 julia --project test/parallel/test_disttranspose_dealiased.jl # mpiexec -n 2 julia --project test/parallel/test_dist_plan_alloc.jl # mpiexec -n 4 julia --project test/parallel/test_mpi_audit_fixes.jl +# mpiexec -n 4 julia --project test/parallel/test_mpi_2d_alignment.jl # mpiexec -n 1 julia --project test/parallel/test_parallel_ad_storage.jl using Test diff --git a/test/parallel/test_mpi_2d_alignment.jl b/test/parallel/test_mpi_2d_alignment.jl new file mode 100644 index 00000000..92180eb7 --- /dev/null +++ b/test/parallel/test_mpi_2d_alignment.jl @@ -0,0 +1,112 @@ +# SHTnsKit.jl - 2D spectral-plan alignment contracts (run with mpiexec -n 4) +# +# The optimized 2D routines reduce within `l_comm` (ranks sharing an `m_rank`) +# instead of over the whole communicator. That is only valid when the spatial θ +# split lines up with the spectral l split, so that every rank in such a group +# owns a DISTINCT θ slab. Where the precondition holds these paths agree with +# the safe ones to ~1e-14; where it does not they used to return silent garbage: +# +# dist_synthesis_distributed_2d_optimized max|err| = 84.4 (field of O(10)) +# dist_analysis_distributed_2d(assume_aligned=true) rel err = 1.15 +# +# Both now validate collectively and raise. This file pins both directions. +# +# (This is also why `_dist_analysis_2d_aligned` does NOT call +# `_keep_one_phi_partner!` the way its full-comm siblings do: within `l_comm` +# the alignment precondition already guarantees there are no duplicate slabs.) + +using Test, MPI, PencilArrays, PencilFFTs, SHTnsKit, LinearAlgebra, Random + +MPI.Initialized() || MPI.Init() +const COMM = MPI.COMM_WORLD +const RANK = MPI.Comm_rank(COMM) +const NP = MPI.Comm_size(COMM) +const EXT = Base.get_extension(SHTnsKit, :SHTnsKitParallelExt) + +NP == 4 || error("test_mpi_2d_alignment.jl expects 4 ranks, got $NP") + +const LMAX = 8 +const NLAT = LMAX + 2 +const NLON = 2 * LMAX + 2 +const CFG = create_gauss_config(LMAX, NLAT; nlon=NLON) + +Random.seed!(1234) +const FGLOB = randn(NLAT, NLON) +const AREF = analysis(CFG, FGLOB) + +const ALM = let a = zeros(ComplexF64, LMAX + 1, CFG.mmax + 1) + for m in 0:CFG.mmax, l in m:LMAX + a[l+1, m+1] = m == 0 ? Float64(l + 1) : ComplexF64(l + 1, m) + end + a +end +const FREF = synthesis(CFG, ALM) + +"""Build a spatial PencilArray on a `pθ × pφ` process grid, filled from `FGLOB`.""" +function spatial(pθ, pφ) + dims = pφ == 1 ? (1,) : (pθ == 1 ? (2,) : (1, 2)) + pen = Pencil((NLAT, NLON), dims, COMM) + f = PencilArray{Float64}(undef, pen) + lr = PencilArrays.range_local(pen) + loc = parent(f) + for (jj, jg) in enumerate(lr[2]), (ii, ig) in enumerate(lr[1]) + loc[ii, jj] = FGLOB[ig, jg] + end + f +end + +"""Largest spatial error against the serial reference, reduced over all ranks.""" +function spatial_err(out, prototype) + loc = out isa PencilArray ? parent(out) : out + lr = PencilArrays.range_local(PencilArrays.pencil(prototype)) + e = 0.0 + for (jj, jg) in enumerate(lr[2]), (ii, ig) in enumerate(lr[1]) + e = max(e, abs(loc[ii, jj] - FREF[ig, jg])) + end + MPI.Allreduce(e, max, COMM) +end + +@testset "2D plan alignment contracts (4 ranks)" begin + for (pθ, pφ) in ((4, 1), (1, 4), (2, 2)) + f = spatial(pθ, pφ) + plan = EXT.create_distributed_spectral_plan_2d(LMAX, CFG.mmax, COMM; p_l=pθ, p_m=pφ) + try + aligned, _ = EXT.validate_2d_distribution_alignment(plan, f) + # The verdict must be identical on every rank, or the guards below + # would throw on some ranks and enter a collective on others. + @test MPI.Allreduce(aligned, &, COMM) == MPI.Allreduce(aligned, |, COMM) + + @testset "pθ=$pθ × pφ=$pφ (aligned=$aligned)" begin + # The safe paths are correct on every decomposition. + Asafe = EXT.dist_analysis_distributed_2d(CFG, f; plan=plan, assume_aligned=false) + @test maximum(abs, EXT.gather_to_full_dense_2d(Asafe) .- AREF) < + 1e-12 * maximum(abs, AREF) + + dsa = EXT.create_distributed_spectral_array_2d(plan) + EXT.scatter_from_dense_2d!(dsa, ALM) + @test spatial_err(EXT.dist_synthesis_distributed_2d(CFG, dsa; prototype_θφ=f), f) < 1e-11 + + if aligned + # Optimized paths agree with the safe ones where the + # precondition holds. + Aal = EXT.dist_analysis_distributed_2d(CFG, f; plan=plan, assume_aligned=true) + @test maximum(abs, EXT.gather_to_full_dense_2d(Aal) .- AREF) < + 1e-12 * maximum(abs, AREF) + @test spatial_err( + EXT.dist_synthesis_distributed_2d_optimized(CFG, dsa; prototype_θφ=f), f) < 1e-11 + else + # ...and raise, rather than returning garbage, where it does not. + @test_throws ArgumentError EXT.dist_analysis_distributed_2d( + CFG, f; plan=plan, assume_aligned=true) + @test_throws ArgumentError EXT.dist_synthesis_distributed_2d_optimized( + CFG, dsa; prototype_θφ=f) + end + end + finally + close(plan) + end + end +end + +MPI.Barrier(COMM) +RANK == 0 && println("\nAll 2D alignment contract tests PASSED") diff --git a/test/serial/test_basic_transforms.jl b/test/serial/test_basic_transforms.jl index 4f18e287..a65db1f9 100644 --- a/test/serial/test_basic_transforms.jl +++ b/test/serial/test_basic_transforms.jl @@ -8,6 +8,17 @@ using SHTnsKit @isdefined(VERBOSE) || (const VERBOSE = get(ENV, "SHTNSKIT_TEST_VERBOSE", "0") == "1") +# Allocation budgets below are calibrated for a SINGLE-THREADED run, where they +# are exact. The shared m-loop orchestrators start `@threads` tasks whenever +# threads are available, and each threaded region costs a few hundred bytes to +# spawn — a constant independent of problem size, and not the kind of regression +# these budgets exist to catch (a per-row temporary or a dense zero spectrum +# costs tens of KB and grows with lmax). Allow a per-thread slack rather than +# letting the whole suite go red on any multi-core machine; note this makes the +# budgets coarse at `nthreads > 1`, so the tight check is the 1-thread run. +@isdefined(_thread_alloc_slack) || + (_thread_alloc_slack() = Threads.nthreads() > 1 ? 4_096 * Threads.nthreads() : 0) + @testset "Basic Scalar Transforms" begin @testset "Real coefficient matrices" begin for T in (Float32, Float64, Int), norm in (:orthonormal, :schmidt) @@ -295,7 +306,7 @@ using SHTnsKit # Julia/FFTW patch versions can impose a small constant allocation # floor around plan execution. Keep the budget below scratch-sized # allocations while allowing the observed Julia 1.10 Linux floor. - rfft_alloc_budget = 2_048 + rfft_alloc_budget = 2_048 + _thread_alloc_slack() @test @allocated(analysis!(cfg, alm_out, f; fft_scratch=rfft_scratch, use_rfft=true)) <= rfft_alloc_budget @test @allocated(synthesis!(cfg, f_out, alm_c; real_output=true, fft_scratch=rfft_scratch, use_rfft=true)) <= rfft_alloc_budget @@ -374,4 +385,71 @@ using SHTnsKit @test isapprox(Ql_rec_l, Ql_rec_full[1:ltr+1]; rtol=1e-10, atol=1e-12) end + + @testset "analysis inverts synthesis under every phi_scale" begin + # `synthesis` honoured `phi_scale` while `analysis` always applied a fixed + # `cphi`, so the two halves of the pair disagreed about the convention: + # under :quad `analysis(synthesis(alm))` came back as `alm/2π` exactly. + rng = MersenneTwister(7731) + for mode in (:dft, :quad) + cfg = create_gauss_config(6, 8) + cfg.phi_scale = mode + alm = zeros(ComplexF64, cfg.lmax + 1, cfg.mmax + 1) + for m in 0:cfg.mmax, l in m:cfg.lmax + alm[l+1, m+1] = m == 0 ? randn(rng) : complex(randn(rng), randn(rng)) + end + S = copy(alm); T = 0.5 .* alm; S[1,1] = 0; T[1,1] = 0 + + @test analysis(cfg, synthesis(cfg, alm)) ≈ alm rtol=1e-10 + Vt, Vp = synthesis_sphtor(cfg, S, T) + S2, T2 = analysis_sphtor(cfg, Vt, Vp) + @test S2 ≈ S rtol=1e-10 + @test T2 ≈ T rtol=1e-10 + + f = synthesis(cfg, alm) + @test analysis_batch(cfg, reshape(f, size(f)..., 1))[:, :, 1] ≈ alm rtol=1e-10 + + # the planned path must agree with the cfg form in both modes + plan = SHTPlan(cfg) + out = similar(alm) + analysis!(plan, out, f) + @test out ≈ alm rtol=1e-10 + + # a hand-built config must not disagree with the constructor's + # for the same grid: `:auto` used to mean `:quad` for non-Gauss grids + @test SHTnsKit.phi_inv_scale(create_regular_config(6, 10; nlon=14)) == + Float64(14) + end + end + + @testset "padded spatial buffers reach the transforms" begin + # `allocate_padded_spatial` returns nlat_padded rows and every transform + # demands exactly nlat, so the padding API had no usable path into a + # transform at all. `spatial_view` is that path, and it preserves the + # padded column stride the padding exists for. + cfg = create_gauss_config(16, 18) + set_allow_padding!(cfg) + @test get_nlat_padded(cfg) > cfg.nlat + + rng = MersenneTwister(4477) + f = randn(rng, cfg.nlat, cfg.nlon) + pad = allocate_padded_spatial(cfg) + copy_to_padded!(pad, f, cfg) + v = spatial_view(cfg, pad) + + @test size(v) == (cfg.nlat, cfg.nlon) + @test stride(v, 2) == get_nlat_padded(cfg) # padding retained + @test analysis(cfg, v) == analysis(cfg, f) # bit-identical + + batch = allocate_padded_spatial_batch(cfg, 3) + fb = randn(rng, cfg.nlat, cfg.nlon, 3) + for k in 1:3 + copy_to_padded!(view(batch, :, :, k), view(fb, :, :, k), cfg) + end + @test analysis_batch(cfg, spatial_view(cfg, batch)) == analysis_batch(cfg, fb) + + @test_throws DimensionMismatch spatial_view(cfg, zeros(cfg.nlat - 1, cfg.nlon)) + @test_throws DimensionMismatch spatial_view(cfg, zeros(cfg.nlat, cfg.nlon + 1)) + end + end diff --git a/test/serial/test_batch_qst.jl b/test/serial/test_batch_qst.jl index 9b2269fd..c55a624b 100644 --- a/test/serial/test_batch_qst.jl +++ b/test/serial/test_batch_qst.jl @@ -8,6 +8,17 @@ using SHTnsKit @isdefined(VERBOSE) || (const VERBOSE = get(ENV, "SHTNSKIT_TEST_VERBOSE", "0") == "1") +# Allocation budgets below are calibrated for a SINGLE-THREADED run, where they +# are exact. The shared m-loop orchestrators start `@threads` tasks whenever +# threads are available, and each threaded region costs a few hundred bytes to +# spawn — a constant independent of problem size, and not the kind of regression +# these budgets exist to catch (a per-row temporary or a dense zero spectrum +# costs tens of KB and grows with lmax). Allow a per-thread slack rather than +# letting the whole suite go red on any multi-core machine; note this makes the +# budgets coarse at `nthreads > 1`, so the tight check is the 1-thread run. +@isdefined(_thread_alloc_slack) || + (_thread_alloc_slack() = Threads.nthreads() > 1 ? 4_096 * Threads.nthreads() : 0) + function _real_alm(rng, lmax, mmax) a = randn(rng, ComplexF64, lmax + 1, mmax + 1) a[:, 1] .= real.(a[:, 1]) @@ -51,7 +62,7 @@ end synthesis_qst_batch(cfg, Qb, Sb, Tb) GC.gc() - @test @allocated(synthesis_qst_batch(cfg, Qb, Sb, Tb)) <= 25_000 + @test @allocated(synthesis_qst_batch(cfg, Qb, Sb, Tb)) <= 25_000 + _thread_alloc_slack() end @testset "analysis_qst_batch matches per-field analysis_qst" begin @@ -79,7 +90,7 @@ end analysis_qst_batch(cfg, Vr_b, Vt_b, Vp_b) GC.gc() - @test @allocated(analysis_qst_batch(cfg, Vr_b, Vt_b, Vp_b)) <= 28_000 + @test @allocated(analysis_qst_batch(cfg, Vr_b, Vt_b, Vp_b)) <= 28_000 + _thread_alloc_slack() end @testset "Batch QST roundtrip" begin diff --git a/test/serial/test_batch_transforms.jl b/test/serial/test_batch_transforms.jl index 654e7f9a..a322808a 100644 --- a/test/serial/test_batch_transforms.jl +++ b/test/serial/test_batch_transforms.jl @@ -7,6 +7,17 @@ using SHTnsKit @isdefined(VERBOSE) || (const VERBOSE = get(ENV, "SHTNSKIT_TEST_VERBOSE", "0") == "1") +# Allocation budgets below are calibrated for a SINGLE-THREADED run, where they +# are exact. The shared m-loop orchestrators start `@threads` tasks whenever +# threads are available, and each threaded region costs a few hundred bytes to +# spawn — a constant independent of problem size, and not the kind of regression +# these budgets exist to catch (a per-row temporary or a dense zero spectrum +# costs tens of KB and grows with lmax). Allow a per-thread slack rather than +# letting the whole suite go red on any multi-core machine; note this makes the +# budgets coarse at `nthreads > 1`, so the tight check is the 1-thread run. +@isdefined(_thread_alloc_slack) || + (_thread_alloc_slack() = Threads.nthreads() > 1 ? 4_096 * Threads.nthreads() : 0) + @testset "Batch Transforms" begin @testset "Scalar batch analysis" begin lmax = 6 @@ -105,10 +116,10 @@ using SHTnsKit synthesis_batch!(cfg, fields_out, alm_batch; fft_batch=rfft_batch, use_rfft=true) GC.gc() - @test @allocated(analysis_batch!(cfg, alm_batch, fields; fft_batch=fft_batch)) <= 8_000 - @test @allocated(synthesis_batch!(cfg, fields_out, alm_batch; fft_batch=fft_batch)) <= 8_000 - @test @allocated(analysis_batch!(cfg, alm_batch, fields; fft_batch=rfft_batch, use_rfft=true)) <= 8_000 - @test @allocated(synthesis_batch!(cfg, fields_out, alm_batch; fft_batch=rfft_batch, use_rfft=true)) <= 8_000 + @test @allocated(analysis_batch!(cfg, alm_batch, fields; fft_batch=fft_batch)) <= 8_000 + _thread_alloc_slack() + @test @allocated(synthesis_batch!(cfg, fields_out, alm_batch; fft_batch=fft_batch)) <= 8_000 + _thread_alloc_slack() + @test @allocated(analysis_batch!(cfg, alm_batch, fields; fft_batch=rfft_batch, use_rfft=true)) <= 8_000 + _thread_alloc_slack() + @test @allocated(synthesis_batch!(cfg, fields_out, alm_batch; fft_batch=rfft_batch, use_rfft=true)) <= 8_000 + _thread_alloc_slack() @inferred synthesis_batch(cfg, alm_batch) fields_kw = synthesis_batch(cfg, alm_batch; real_output=false) @@ -223,8 +234,8 @@ using SHTnsKit synthesis_sphtor_batch(cfg, Slm_batch, Tlm_batch) analysis_sphtor_batch(cfg, Vt_batch, Vp_batch) GC.gc() - @test @allocated(synthesis_sphtor_batch(cfg, Slm_batch, Tlm_batch)) <= 16_000 - @test @allocated(analysis_sphtor_batch(cfg, Vt_batch, Vp_batch)) <= 18_000 + @test @allocated(synthesis_sphtor_batch(cfg, Slm_batch, Tlm_batch)) <= 16_000 + _thread_alloc_slack() + @test @allocated(analysis_sphtor_batch(cfg, Vt_batch, Vp_batch)) <= 18_000 + _thread_alloc_slack() end @testset "QST batch transforms" begin diff --git a/test/serial/test_configuration.jl b/test/serial/test_configuration.jl index 7fe5e5cd..0c374f00 100644 --- a/test/serial/test_configuration.jl +++ b/test/serial/test_configuration.jl @@ -2,6 +2,7 @@ # Tests for grid configuration, indexing, and normalization using Test +using Random using SHTnsKit @isdefined(VERBOSE) || (const VERBOSE = get(ENV, "SHTNSKIT_TEST_VERBOSE", "0") == "1") @@ -329,33 +330,97 @@ using SHTnsKit @test size(spatial_scratch) == (nlat, nlon) end - @testset "FFT plan cache control" begin - # FFT plan cache requires the parallel extension - # Skip if not available - try - # Save initial state - initial_state = fft_plan_cache_enabled() - @test typeof(initial_state) == Bool - - # Test enable/disable - enable_fft_plan_cache!() - @test fft_plan_cache_enabled() == true - - disable_fft_plan_cache!() - @test fft_plan_cache_enabled() == false + @testset "structural field assignment keeps derived state consistent" begin + # `cfg.lmax = 10` used to leave `size(cfg.Nlm) == (7,7)` while the + # transforms index it as (lmax+1, mmax+1) under `@inbounds` — a live + # out-of-bounds read. lmax/mmax/mres now rebuild the spectral layout; + # the fields that cannot be made consistent are rejected outright. + cfg = create_gauss_config(6, 8) + prepare_plm_tables!(cfg) + @test SHTnsKit.has_fused_scalar_tables(cfg) + + cfg.lmax = 10 + @test size(cfg.Nlm) == (cfg.lmax + 1, cfg.mmax + 1) + @test cfg.nlm == SHTnsKit.nlm_calc(cfg.lmax, cfg.mmax, cfg.mres) + @test length(cfg.li) == cfg.nlm && length(cfg.mi) == cfg.nlm + @test !SHTnsKit.has_fused_scalar_tables(cfg) # stale tables dropped + + cfg2 = create_gauss_config(6, 8; mres=1) + cfg2.mres = 2 + @test cfg2.nlm == SHTnsKit.nlm_calc(6, cfg2.mmax, 2) + @test all(m -> m % 2 == 0, cfg2.mi) + + for bad in (:nlat, :nlon, :grid_type, :nlm, :nspat) + @test_throws ArgumentError setproperty!(create_gauss_config(4, 6), bad, + bad === :grid_type ? :regular : 99) + end + # An inconsistent spectral triple is rejected rather than stored. + @test_throws ArgumentError (c = create_gauss_config(6, 8); c.mmax = 99) + end - # Test set function - set_fft_plan_cache!(true) - @test fft_plan_cache_enabled() == true + @testset "kwarg SHTConfig constructor validates its invariants" begin + # This constructor is exported and used to check nothing, so a hand-built + # config could violate nlon >= 2*mmax+1 and then silently synthesize an + # all-zero field for any mode it could not resolve. + lmax = mmax = 6 + mk(; nlon=2*mmax + 2, nlat=8, kw...) = begin + θ = collect(range(0.1, 3.0; length=nlat)) + base = (; lmax, mmax, mres=1, nlat, nlon, + θ, φ=collect(range(0, 2π; length=nlon+1))[1:nlon], + x=cos.(θ), w=fill(2/nlat, nlat), st=sin.(θ), + Nlm=SHTnsKit.Nlm_table(lmax, mmax), cphi=2π/nlon, + nlm=SHTnsKit.nlm_calc(lmax, mmax, 1), + li=SHTnsKit.build_li_mi(lmax, mmax, 1)[1], + mi=SHTnsKit.build_li_mi(lmax, mmax, 1)[2], + nspat=nlat*nlon, norm=:orthonormal, cs_phase=true, + real_norm=false, robert_form=false) + SHTnsKit.SHTConfig(; base..., kw...) + end + @test mk() isa SHTnsKit.SHTConfig # the valid case still builds + @test_throws ArgumentError mk(nlon=8) # nlon < 2*mmax+1 + @test_throws ArgumentError mk(mres=0) + @test_throws ArgumentError mk(nspat=1) # nspat != nlat*nlon + @test_throws DimensionMismatch mk(w=fill(0.25, 4)) # w shorter than nlat + @test_throws DimensionMismatch mk(Nlm=zeros(2, 2)) # Nlm not (lmax+1, mmax+1) + end - set_fft_plan_cache!(false) - @test fft_plan_cache_enabled() == false + @testset "FFT plan cache control" begin + # These knobs used to forward to a cache in the parallel extension that + # nothing ever read (`_get_or_plan` had no call sites), so the whole + # documented feature was a no-op and this testset was wrapped in a + # try/catch that skipped it whenever MPI was absent. They now control the + # φ-FFT plan cache every transform actually goes through, serial or + # distributed, so assert the behaviour rather than just the flag. + initial_state = fft_plan_cache_enabled() + @test typeof(initial_state) == Bool - # Restore initial state - set_fft_plan_cache!(initial_state) - catch e - @info "Skipping FFT plan cache tests (requires parallel extension)" exception=e - end + enable_fft_plan_cache!() + @test fft_plan_cache_enabled() == true + disable_fft_plan_cache!() + @test fft_plan_cache_enabled() == false + set_fft_plan_cache!(true) + @test fft_plan_cache_enabled() == true + set_fft_plan_cache!(false) + @test fft_plan_cache_enabled() == false + + enable_fft_plan_cache!() + cfg_cache = create_gauss_config(4, 6) + field = randn(MersenneTwister(77), cfg_cache.nlat, cfg_cache.nlon) + ref = analysis(cfg_cache, field) + @test !isempty(SHTnsKit._LOCAL_FFT_PLAN_CACHE) + + disable_fft_plan_cache!() # clears by default + @test isempty(SHTnsKit._LOCAL_FFT_PLAN_CACHE) + @test analysis(cfg_cache, field) ≈ ref # same answer, unplanned + @test isempty(SHTnsKit._LOCAL_FFT_PLAN_CACHE) # and still bypassed + + enable_fft_plan_cache!() + @test analysis(cfg_cache, field) ≈ ref + @test !isempty(SHTnsKit._LOCAL_FFT_PLAN_CACHE) + + @test SHTnsKit.fft_plan_cache_max!(SHTnsKit.fft_plan_cache_max!(8)) == 8 + + set_fft_plan_cache!(initial_state) end @testset "Pencil grid suggestion" begin diff --git a/test/serial/test_energy_diagnostics.jl b/test/serial/test_energy_diagnostics.jl index 2c60a4a0..af2e5570 100644 --- a/test/serial/test_energy_diagnostics.jl +++ b/test/serial/test_energy_diagnostics.jl @@ -283,4 +283,17 @@ using SHTnsKit @test energy_vector(cfg, Slm, Tlm) >= 0 end end + + @testset "spectral energy rejects a mis-sized spectrum" begin + # The accumulation loops are `@inbounds`; a too-small matrix used to + # return a silently wrong number instead of raising. + cfg = create_gauss_config(6, 8) + small = zeros(ComplexF64, 2, 2) + ok = zeros(ComplexF64, cfg.lmax + 1, cfg.mmax + 1) + @test_throws DimensionMismatch energy_scalar(cfg, small) + @test_throws DimensionMismatch energy_vector(cfg, small, ok) + @test_throws DimensionMismatch energy_vector(cfg, ok, small) + @test energy_scalar(cfg, ok) == 0.0 + end + end diff --git a/test/serial/test_gradients.jl b/test/serial/test_gradients.jl index 354833b7..a43ee3de 100644 --- a/test/serial/test_gradients.jl +++ b/test/serial/test_gradients.jl @@ -402,4 +402,75 @@ end @test isapprox(dE_ad, dE_fd; rtol=5e-4, atol=1e-8) end + +end + + +if _HAS_CHAINRULES +@testset "QST rrules and implementation-detail kwargs" begin + # The QST family had no rrule at all, so differentiating a QST pipeline fell + # through to source tracing and crashed inside FFTW. And every rrule that + # declared fewer kwargs than its primal was skipped the moment a caller + # passed one, even at its default value. + cfg = create_gauss_config(5, 7; nlon=12) + rng = MersenneTwister(9412) + rnd() = begin + a = zeros(ComplexF64, cfg.lmax + 1, cfg.mmax + 1) + for m in 0:cfg.mmax, l in m:cfg.lmax + a[l+1, m+1] = m == 0 ? randn(rng) : complex(randn(rng), randn(rng)) + end + a + end + ε = 1e-6 + + @testset "synthesis_qst pullback vs finite differences" begin + Q, S, T = rnd(), rnd(), rnd() + Cr, Ct, Cp = randn(rng, cfg.nlat, cfg.nlon), randn(rng, cfg.nlat, cfg.nlon), randn(rng, cfg.nlat, cfg.nlon) + loss(q, s, t) = (V = synthesis_qst(cfg, q, s, t); + sum(Cr .* V[1]) + sum(Ct .* V[2]) + sum(Cp .* V[3])) + _, back = rrule(synthesis_qst, cfg, Q, S, T) + _, _, Q̄, S̄, T̄, _ = back((Cr, Ct, Cp)) + for (slot, ḡ, perturb) in ((:Q, Q̄, (h, q, s, t) -> (q .+ h, s, t)), + (:S, S̄, (h, q, s, t) -> (q, s .+ h, t)), + (:T, T̄, (h, q, s, t) -> (q, s, t .+ h))) + h = rnd() + fd = (loss(perturb(ε .* h, Q, S, T)...) - loss(perturb(-ε .* h, Q, S, T)...)) / (2ε) + @test real(sum(conj(ḡ) .* h)) ≈ fd rtol=1e-6 atol=1e-8 + end + end + + @testset "analysis_qst pullback vs finite differences" begin + Vr, Vt, Vp = randn(rng, cfg.nlat, cfg.nlon), randn(rng, cfg.nlat, cfg.nlon), randn(rng, cfg.nlat, cfg.nlon) + G1, G2, G3 = rnd(), rnd(), rnd() + loss(vr, vt, vp) = (A = analysis_qst(cfg, vr, vt, vp); + real(sum(conj(G1) .* A[1]) + sum(conj(G2) .* A[2]) + sum(conj(G3) .* A[3]))) + _, back = rrule(analysis_qst, cfg, Vr, Vt, Vp) + _, _, V̄r, V̄t, V̄p = back((G1, G2, G3)) + for (ḡ, perturb) in ((V̄r, (h, a, b, c) -> (a .+ h, b, c)), + (V̄t, (h, a, b, c) -> (a, b .+ h, c)), + (V̄p, (h, a, b, c) -> (a, b, c .+ h))) + h = randn(rng, cfg.nlat, cfg.nlon) + fd = (loss(perturb(ε .* h, Vr, Vt, Vp)...) - loss(perturb(-ε .* h, Vr, Vt, Vp)...)) / (2ε) + @test sum(real(ḡ) .* h) ≈ fd rtol=1e-6 atol=1e-8 + end + end + + @testset "rrules accept the primal's implementation-detail kwargs" begin + alm = rnd() + f = randn(rng, cfg.nlat, cfg.nlon) + Vt, Vp = randn(rng, cfg.nlat, cfg.nlon), randn(rng, cfg.nlat, cfg.nlon) + @test rrule(synthesis, cfg, alm; use_rfft=false) !== nothing + @test rrule(synthesis, cfg, alm; real_output=true, fft_scratch=nothing) !== nothing + @test rrule(analysis, cfg, f; use_rfft=false) !== nothing + @test rrule(analysis_sphtor, cfg, Vt, Vp; use_rfft=false) !== nothing + @test rrule(synthesis_sphtor, cfg, alm, alm; use_rfft=false) !== nothing + # and the rfft path gives the same adjoint as the complex one + ȳ = randn(rng, cfg.nlat, cfg.nlon) + _, b1 = rrule(synthesis, cfg, alm) + _, b2 = rrule(synthesis, cfg, alm; use_rfft=true) + @test b1(ȳ)[3] ≈ b2(ȳ)[3] rtol=1e-10 + end +end +else + @info "Skipping QST rrule tests (ChainRulesCore not available)" end diff --git a/test/serial/test_indexing.jl b/test/serial/test_indexing.jl index 5350283c..9fc918a0 100644 --- a/test/serial/test_indexing.jl +++ b/test/serial/test_indexing.jl @@ -210,4 +210,22 @@ using SHTnsKit cfg = create_gauss_config(lmax, lmax + 2) @test cfg.nlm == nlm_calc(lmax, lmax, mres) end + + @testset "im_from_lm is bounded by mmax, not lmax" begin + # The bound was `lmax ÷ mres`, so for an mmax < lmax layout a + # past-the-end index resolved to an order the configuration does not + # store instead of raising. + lmax, mmax = 8, 3 + n = nlm_calc(lmax, mmax, 1) + # every in-range index still maps to its own order + li, mi = build_li_mi(lmax, mmax, 1) + for k in 1:n + @test SHTnsKit.im_from_lm(k - 1, lmax, 1; mmax=mmax) == mi[k] + end + # one past the end now raises rather than inventing m = 4 + @test_throws ArgumentError SHTnsKit.im_from_lm(n, lmax, 1; mmax=mmax) + # default mmax=lmax keeps the old, correct behaviour for full layouts + @test SHTnsKit.im_from_lm(nlm_calc(lmax, lmax, 1) - 1, lmax, 1) == lmax + end + end diff --git a/test/serial/test_local.jl b/test/serial/test_local.jl index 9a8840b0..58da8940 100644 --- a/test/serial/test_local.jl +++ b/test/serial/test_local.jl @@ -48,6 +48,69 @@ using Random end end + @testset "point/latitude evaluators honour phi_scale" begin + # `synthesis` scales its Fourier bins by `phi_inv_scale(cfg)` and the + # inverse FFT divides by nlon, a net factor of 1 under :dft but 1/2π + # under :quad. The direct evaluators applied no factor at all, so under + # :quad every one of them disagreed with the grid it claims to sample by + # exactly 2π. + for mode in (:dft, :quad) + cfg = create_gauss_config(6, 8) + cfg.phi_scale = mode + rng = MersenneTwister(5150) + A = zeros(ComplexF64, cfg.lmax + 1, cfg.mmax + 1) + for m in 0:cfg.mmax, l in m:cfg.lmax + A[l+1, m+1] = m == 0 ? randn(rng) : complex(randn(rng), randn(rng)) + end + packed = SHTnsKit.pack_lm(cfg, A) + zero_packed = zeros(ComplexF64, cfg.nlm) + + f = synthesis(cfg, A) + Vr, Vt, Vp = synthesis_qst(cfg, A, zero(A), zero(A)) + i, j = 2, 3 + + @test synthesis_point(cfg, A, cfg.x[i], cfg.φ[j]) ≈ f[i, j] rtol=1e-10 + @test SH_to_lat(cfg, packed, cfg.x[i]) ≈ f[i, :] rtol=1e-10 + @test SHqst_to_lat(cfg, packed, zero_packed, zero_packed, cfg.x[i])[1] ≈ + Vr[i, :] rtol=1e-10 + @test SHTnsKit.SHqst_to_point(cfg, packed, zero_packed, zero_packed, + cfg.x[i], cfg.φ[j])[1] ≈ Vr[i, j] rtol=1e-10 + + # axisymmetric pair + a0 = zeros(ComplexF64, cfg.lmax + 1); a0[3] = 0.7 + am = zeros(ComplexF64, cfg.lmax + 1, cfg.mmax + 1); am[3, 1] = 0.7 + @test synthesis_axisym(cfg, a0) ≈ synthesis(cfg, am)[:, 1] rtol=1e-10 + end + end + + @testset "evaluators preserve the coefficient element type" begin + # The φ convention factor is a Float64. Multiplying a Float32 evaluation + # by it silently widens every result to Float64, breaking the element + # type these functions promise their caller (and the Dual types AD needs + # to see through). Narrow the scale at the boundary instead. + for mode in (:dft, :quad), T in (Float32, Float64) + cfg = create_gauss_config(4, 6; nlon=9) + cfg.phi_scale = mode + CT = Complex{T} + A = zeros(CT, cfg.lmax + 1, cfg.mmax + 1) + A[2, 1] = CT(0.6); A[3, 2] = CT(0.4, 0.2) + packed = SHTnsKit.pack_lm(cfg, A) + zpack = zeros(CT, cfg.nlm) + x = T(cfg.x[2]); φ = T(cfg.φ[3]) + + @test synthesis_point(cfg, A, x, φ) isa T + @test eltype(SH_to_lat(cfg, packed, x)) === T + @test eltype(SH_to_lat_cplx(cfg, zeros(CT, SHTnsKit.nlm_cplx_calc(cfg.lmax, cfg.mmax, 1)), x)) === CT + @test all(v -> v isa T, SHTnsKit.SHqst_to_point(cfg, packed, zpack, zpack, x, φ)) + @test all(v -> v isa T, SHTnsKit.SH_to_grad_point(cfg, packed, zpack, x, φ)) + @test all(V -> eltype(V) === T, SHqst_to_lat(cfg, packed, zpack, zpack, x)) + @test eltype(synthesis_axisym(cfg, A[:, 1])) === T + @test eltype(SHTnsKit.synthesis_axisym_l(cfg, A[:, 1], cfg.lmax)) === T + @test SHTnsKit.synthesis_point_cplx( + cfg, zeros(CT, SHTnsKit.nlm_cplx_calc(cfg.lmax, cfg.mmax, 1)), x, φ) isa CT + end + end + @testset "SH_to_lat matches synthesis at grid latitudes" begin lmax = 8 nlat = lmax + 2 diff --git a/test/serial/test_plan.jl b/test/serial/test_plan.jl index f3269036..62c237e4 100644 --- a/test/serial/test_plan.jl +++ b/test/serial/test_plan.jl @@ -8,6 +8,15 @@ using SHTnsKit @isdefined(VERBOSE) || (const VERBOSE = get(ENV, "SHTNSKIT_TEST_VERBOSE", "0") == "1") +# Allocation budgets below are calibrated for a SINGLE-THREADED run, where they +# are exact. The shared m-loop orchestrators start `@threads` tasks whenever +# threads are available, and each threaded region costs a few hundred bytes to +# spawn — a constant independent of problem size, and not the kind of regression +# these budgets exist to catch. Allow a per-thread slack rather than letting the +# suite go red on any multi-core machine; the tight check is the 1-thread run. +@isdefined(_thread_alloc_slack) || + (_thread_alloc_slack() = Threads.nthreads() > 1 ? 4_096 * Threads.nthreads() : 0) + function _rand_real_alm(rng, lmax, mmax) alm = randn(rng, ComplexF64, lmax + 1, mmax + 1) alm[:, 1] .= real.(alm[:, 1]) @@ -161,10 +170,15 @@ end # runners. Keep the strict production-path ceiling while still ruling # out any field-sized allocation in coverage-enabled CI. allocation_limit = Base.JLOptions().code_coverage == 0 ? 128 : 2048 + # The planned forms own their scratch and never spawn tasks, so they stay + # at the strict ceiling regardless of thread count. The `cfg` form routes + # through the threaded m-loop orchestrator, so it pays the task-spawn + # constant; that is not a field-sized allocation and not what this guards. @test @allocated(synthesis!(plan, f, alm)) <= allocation_limit @test @allocated(synthesis!(plan_r, f, alm)) <= allocation_limit @test @allocated(synthesis_sphtor!(plan, Vt, Vp, Slm, Tlm)) <= allocation_limit - @test @allocated(synthesis!(cfg, f, alm; fft_scratch)) <= allocation_limit + @test @allocated(synthesis!(cfg, f, alm; fft_scratch)) <= + allocation_limit + _thread_alloc_slack() end @testset "Planned scalar matches the non-planned path exactly" begin @@ -262,4 +276,51 @@ end @test all(isfinite, f_cplx) @test eltype(f_cplx) <: Complex end + + @testset "planned transforms match the cfg form in every buffer mode" begin + # The planned transforms exist to be a drop-in, faster replacement for + # the `cfg` forms. Pin that they agree numerically across the whole + # matrix of buffer modes: Legendre tables on/off, rfft on/off, and + # Robert form on/off — a divergence in any cell is a silent wrong + # answer for anyone who reaches for a plan. + rng = MersenneTwister(4821) + for tables in (false, true), use_rfft in (false, true), robert in (false, true) + cfgp = create_gauss_config(10, 12) + cfgp.robert_form = robert + tables ? SHTnsKit.prepare_plm_tables!(cfgp) : SHTnsKit.disable_plm_tables!(cfgp) + + alm = zeros(ComplexF64, cfgp.lmax + 1, cfgp.mmax + 1) + for m in 0:cfgp.mmax, l in m:cfgp.lmax + alm[l + 1, m + 1] = m == 0 ? randn(rng) : complex(randn(rng), randn(rng)) + end + S = copy(alm); T = 0.5 .* alm + S[1, 1] = 0; T[1, 1] = 0 + + f = synthesis(cfgp, alm) + Vt, Vp = synthesis_sphtor(cfgp, S, T) + alm_ref = analysis(cfgp, f) + S_ref, T_ref = analysis_sphtor(cfgp, Vt, Vp) + + plan = SHTPlan(cfgp; use_rfft=use_rfft) + tol = 1e-12 + + alm_out = similar(alm_ref) + analysis!(plan, alm_out, f) + @test maximum(abs, alm_out .- alm_ref) <= tol + + f_out = similar(f) + synthesis!(plan, f_out, alm) + @test maximum(abs, f_out .- f) <= tol + + S_out = similar(S_ref); T_out = similar(T_ref) + SHTnsKit.analysis_sphtor!(plan, S_out, T_out, Vt, Vp) + @test maximum(abs, S_out .- S_ref) <= tol + @test maximum(abs, T_out .- T_ref) <= tol + + Vt_out = similar(Vt); Vp_out = similar(Vp) + SHTnsKit.synthesis_sphtor!(plan, Vt_out, Vp_out, S, T) + @test maximum(abs, Vt_out .- Vt) <= tol + @test maximum(abs, Vp_out .- Vp) <= tol + end + end end diff --git a/test/serial/test_qst_transforms.jl b/test/serial/test_qst_transforms.jl index 0e0e9a4e..e2ae8f1f 100644 --- a/test/serial/test_qst_transforms.jl +++ b/test/serial/test_qst_transforms.jl @@ -7,6 +7,17 @@ using SHTnsKit @isdefined(VERBOSE) || (const VERBOSE = get(ENV, "SHTNSKIT_TEST_VERBOSE", "0") == "1") +# Allocation budgets below are calibrated for a SINGLE-THREADED run, where they +# are exact. The shared m-loop orchestrators start `@threads` tasks whenever +# threads are available, and each threaded region costs a few hundred bytes to +# spawn — a constant independent of problem size, and not the kind of regression +# these budgets exist to catch (a per-row temporary or a dense zero spectrum +# costs tens of KB and grows with lmax). Allow a per-thread slack rather than +# letting the whole suite go red on any multi-core machine; note this makes the +# budgets coarse at `nthreads > 1`, so the tight check is the 1-thread run. +@isdefined(_thread_alloc_slack) || + (_thread_alloc_slack() = Threads.nthreads() > 1 ? 4_096 * Threads.nthreads() : 0) + @testset "QST (3D Vector) Transforms" begin @testset "QST roundtrip" begin lmax = 6 @@ -90,7 +101,7 @@ using SHTnsKit # Array header/alignment overhead differs across Julia patch versions # and platforms (notably Windows CI). Keep this below an extra # component-sized scratch allocation while allowing that fixed floor. - qst_l_alloc_budget = 13_000 + qst_l_alloc_budget = 13_000 + _thread_alloc_slack() @test @allocated(synthesis_qst_l(cfg, Qlm, Slm, Tlm, ltr; real_output=true)) <= qst_l_alloc_budget end diff --git a/test/serial/test_rotation_gradients.jl b/test/serial/test_rotation_gradients.jl index 1ac3d149..273c01c0 100644 --- a/test/serial/test_rotation_gradients.jl +++ b/test/serial/test_rotation_gradients.jl @@ -83,6 +83,102 @@ end end end +@testset "Angle gradients survive an in-place primal" begin + # `SH_Zrotate` is covered above. Its siblings capture the primal INPUT and + # read it lazily in the pullback, so the same hazard applies to them: an + # in-place call (Rlm === Qlm) overwrites those coefficients, and so does any + # caller that reuses the buffer before the pullback runs. Each rrule must + # snapshot what it needs at primal time. + rng = MersenneTwister(20931) + ε = 1e-6 + + @testset "SH_Yrotate dα" begin + cfg = create_gauss_config(4, 6; nlon=9) + Q = randn(rng, ComplexF64, cfg.nlm) + C = randn(rng, ComplexF64, cfg.nlm) + α = 0.7 + loss(q, a) = real(sum(conj(C) .* SH_Yrotate(cfg, copy(q), a, similar(q)))) + fd = (loss(Q, α + ε) - loss(Q, α - ε)) / (2ε) + + q = copy(Q) + _, back_sep = ChainRulesCore.rrule(SH_Yrotate, cfg, q, α, similar(q)) + @test back_sep(C)[4] ≈ fd rtol=1e-6 atol=1e-8 + + inplace = copy(Q) # Rlm === Qlm + _, back_ip = ChainRulesCore.rrule(SH_Yrotate, cfg, inplace, α, inplace) + @test back_ip(C)[4] ≈ fd rtol=1e-6 atol=1e-8 + + reused = copy(Q) # caller clobbers the input afterwards + _, back_reuse = ChainRulesCore.rrule(SH_Yrotate, cfg, reused, α, similar(reused)) + fill!(reused, 0) + @test back_reuse(C)[4] ≈ fd rtol=1e-6 atol=1e-8 + + if _HAS_ZYGOTE_ROT + zq = copy(Q) + _, zback = Zygote.pullback(SH_Yrotate, cfg, zq, α, zq) + @test zback(C)[3] ≈ fd rtol=1e-6 atol=1e-8 + end + end + + @testset "shtns_rotation_apply_cplx dβ" begin + lmax = mmax = 4 + mkrot(β) = (r = SHTnsKit.SHTRotation(lmax, mmax); + SHTnsKit.shtns_rotation_set_angles_ZYZ(r, 0.3, β, 0.2); r) + n = SHTnsKit.nlm_cplx_calc(lmax, mmax, 1) + Z = randn(rng, ComplexF64, n) + C = randn(rng, ComplexF64, n) + β = 0.7 + loss(r, z) = (R = similar(z); + SHTnsKit.shtns_rotation_apply_cplx(r, copy(z), R); + real(sum(conj(C) .* R))) + fd = (loss(mkrot(β + ε), Z) - loss(mkrot(β - ε), Z)) / (2ε) + + z = copy(Z) + _, back_sep = ChainRulesCore.rrule(SHTnsKit.shtns_rotation_apply_cplx, mkrot(β), z, similar(z)) + @test back_sep(C)[2].β ≈ fd rtol=1e-6 atol=1e-8 + + zi = copy(Z) # Rlm === Zlm + _, back_ip = ChainRulesCore.rrule(SHTnsKit.shtns_rotation_apply_cplx, mkrot(β), zi, zi) + @test back_ip(C)[2].β ≈ fd rtol=1e-6 atol=1e-8 + end + + @testset "shtns_rotation_apply_real dβ" begin + cfg = create_gauss_config(4, 7; nlon=11) + mkrot(β) = (r = SHTnsKit.SHTRotation(cfg.lmax, cfg.mmax); + SHTnsKit.shtns_rotation_set_angles_ZYZ(r, 0.3, β, 0.2); r) + Q = randn(rng, ComplexF64, cfg.nlm) + C = randn(rng, ComplexF64, cfg.nlm) + β = 0.7 + loss(r, q) = (R = similar(q); + SHTnsKit.shtns_rotation_apply_real(r, copy(q), R); + real(sum(conj(C) .* R))) + fd = (loss(mkrot(β + ε), Q) - loss(mkrot(β - ε), Q)) / (2ε) + + q = copy(Q) + _, back_sep = ChainRulesCore.rrule(SHTnsKit.shtns_rotation_apply_real, mkrot(β), q, similar(q)) + @test back_sep(C)[2].β ≈ fd rtol=1e-6 atol=1e-8 + + qi = copy(Q) # Rlm === Qlm + _, back_ip = ChainRulesCore.rrule(SHTnsKit.shtns_rotation_apply_real, mkrot(β), qi, qi) + @test back_ip(C)[2].β ≈ fd rtol=1e-6 atol=1e-8 + end +end + +@testset "Order-mixing rotations reject mres > 1 with a usable message" begin + cfg = create_gauss_config(4, 6; mres=2) + Q = randn(MersenneTwister(5), ComplexF64, cfg.nlm) + err = try + SH_Yrotate(cfg, Q, 0.3, similar(Q)); nothing + catch e + e + end + # Either guard is fine — SH_Yrotate's own mres check or the packed-length + # check inside the Wigner engine — as long as the message names `mres`, so + # the caller is not left reverse-engineering a bare size mismatch. + @test err isa Union{ArgumentError,DimensionMismatch} + @test occursin("mres", sprint(showerror, err)) +end + @testset "Complex-packed analysis rrule respects configured convention" begin lmax = 4 cfg = create_gauss_config(lmax, 7; nlon=11, norm=:schmidt, diff --git a/test/serial/test_rotations.jl b/test/serial/test_rotations.jl index 4387d0d6..4cdc9b76 100644 --- a/test/serial/test_rotations.jl +++ b/test/serial/test_rotations.jl @@ -33,6 +33,56 @@ using SHTnsKit end end + @testset "Z-rotation sign convention is pinned to physical ground truth" begin + # The `exp(-imα)` sign is not a free choice: flipping it to `exp(+imα)` + # (the passive convention) silently desynchronises SH_Zrotate from the + # actual spatial rotation, from the general Wigner engine, and from the + # distributed twins. Pin all three so it cannot drift again. + lmax = 6 + nlon = 2 * lmax + 2 # even, so an exact grid shift exists + cfg = create_gauss_config(lmax, lmax + 2; nlon=nlon) + rng = MersenneTwister(4242) + + A = zeros(ComplexF64, lmax + 1, cfg.mmax + 1) + for m in 0:cfg.mmax, l in m:lmax + A[l + 1, m + 1] = m == 0 ? randn(rng) : complex(randn(rng), randn(rng)) + end + Qlm = SHTnsKit.pack_lm(cfg, A) + + shift = 3 # α is an exact multiple of the φ spacing + α = 2π * shift / nlon + Rlm = similar(Qlm) + SH_Zrotate(cfg, Qlm, α, Rlm) + + # (1) Ground truth: rotating the FIELD by +α about ẑ is g(θ,φ) = f(θ,φ-α), + # which on this grid is a pure column shift. + f = synthesis(cfg, A) + g = similar(f) + for j in 1:nlon + g[:, j] .= @view f[:, mod(j - 1 - shift, nlon) + 1] + end + spatial = SHTnsKit.pack_lm(cfg, analysis(cfg, g)) + @test isapprox(Rlm, spatial; rtol=1e-10, atol=1e-12) + + # ...and the opposite sign is the OTHER rotation, not a rounding detail. + flipped = [Qlm[k] * cis(+cfg.mi[k] * α) for k in 1:cfg.nlm] + @test !isapprox(flipped, spatial; rtol=1e-3, atol=1e-6) + + # (2) The general ZYZ engine must agree for both the α and γ slots. + for angles in ((α, 0.0, 0.0), (0.0, 0.0, α)) + r = SHTnsKit.SHTRotation(cfg.lmax, cfg.mmax) + SHTnsKit.shtns_rotation_set_angles_ZYZ(r, angles...) + Rzyz = similar(Qlm) + SHTnsKit.shtns_rotation_apply_real(r, Qlm, Rzyz) + @test isapprox(Rzyz, Rlm; rtol=1e-10, atol=1e-12) + end + + # (3) The dense distributed twin must agree coefficient for coefficient. + Rdense = zeros(ComplexF64, lmax + 1, cfg.mmax + 1) + SHTnsKit.dist_SH_Zrotate(cfg, A, α, Rdense) + @test isapprox(SHTnsKit.pack_lm(cfg, Rdense), Rlm; rtol=1e-12, atol=1e-14) + end + @testset "Z-axis rotation in-place" begin lmax = 6 cfg = create_gauss_config(lmax, lmax + 2; nlon=2*lmax + 1) @@ -128,6 +178,104 @@ using SHTnsKit @test isapprox(Rlm4, Qlm; rtol=1e-8, atol=1e-10) end + @testset "X-rotation 90 is Rx(+pi/2), not its inverse" begin + # `SH_Xrotate90` must be the +90° turn its name promises. The existing + # "four 90° rotations = identity" check passes for either sign, so pin + # the direction explicitly against the ZXZ entry point. + lmax = 5 + cfg = create_gauss_config(lmax, lmax + 2; nlon=2*lmax + 2) + rng = MersenneTwister(8801) + Qlm = randn(rng, ComplexF64, cfg.nlm) + Qlm[1:lmax+1] .= real.(Qlm[1:lmax+1]) + + X = similar(Qlm) + SH_Xrotate90(cfg, Qlm, X) + + # Ground truth: a pure Rx(+π/2) via the ZXZ entry point, whose own + # converter documents Rx(β) = Rz(-π/2)·Ry(β)·Rz(π/2). + fwd = SHTnsKit.SHTRotation(cfg.lmax, cfg.mmax) + SHTnsKit.shtns_rotation_set_angles_ZXZ(fwd, 0.0, π/2, 0.0) + Xref = similar(Qlm) + SHTnsKit.shtns_rotation_apply_real(fwd, Qlm, Xref) + @test isapprox(X, Xref; rtol=1e-10, atol=1e-12) + + # ...and it is NOT the inverse rotation. + inv_rot = SHTnsKit.SHTRotation(cfg.lmax, cfg.mmax) + SHTnsKit.shtns_rotation_set_angles_ZXZ(inv_rot, 0.0, -π/2, 0.0) + Xinv = similar(Qlm) + SHTnsKit.shtns_rotation_apply_real(inv_rot, Qlm, Xinv) + @test !isapprox(X, Xinv; rtol=1e-3, atol=1e-6) + + # Three forward turns equal one backward turn. + a = similar(Qlm); b = similar(Qlm); c = similar(Qlm) + SH_Xrotate90(cfg, Qlm, a); SH_Xrotate90(cfg, a, b); SH_Xrotate90(cfg, b, c) + @test isapprox(c, Xinv; rtol=1e-9, atol=1e-11) + end + + @testset "angle-axis half-turns pick the right axis" begin + # At β≈π the observable Euler combination is α-γ read off the NEGATED + # first column; reusing the β≈0 formula there returns a rotation about + # the wrong axis (a 180° turn about x̂ came back as Ry(π)). Only exact + # half-turns are affected, so pin them. + cfg = create_gauss_config(5, 7; nlon=12) + rng = MersenneTwister(6041) + Q = randn(rng, ComplexF64, cfg.nlm) + Q[1:cfg.lmax+1] .= real.(Q[1:cfg.lmax+1]) + apply(r) = (R = similar(Q); SHTnsKit.shtns_rotation_apply_real(r, Q, R); R) + byaxis(θ, ax...) = (r = SHTnsKit.SHTRotation(cfg.lmax, cfg.mmax); + SHTnsKit.shtns_rotation_set_angle_axis(r, θ, ax...); apply(r)) + byzxz(β) = (r = SHTnsKit.SHTRotation(cfg.lmax, cfg.mmax); + SHTnsKit.shtns_rotation_set_angles_ZXZ(r, 0.0, β, 0.0); apply(r)) + byzyz(β) = (r = SHTnsKit.SHTRotation(cfg.lmax, cfg.mmax); + SHTnsKit.shtns_rotation_set_angles_ZYZ(r, 0.0, β, 0.0); apply(r)) + + @test isapprox(byaxis(π, 1.0, 0.0, 0.0), byzxz(π); rtol=1e-10, atol=1e-12) # x̂ half-turn + @test isapprox(byaxis(π, 0.0, 1.0, 0.0), byzyz(π); rtol=1e-10, atol=1e-12) # ŷ half-turn + @test !isapprox(byaxis(π, 1.0, 0.0, 0.0), byzyz(π); rtol=1e-3, atol=1e-6) # and they differ + @test isapprox(byaxis(π/2, 1.0, 0.0, 0.0), byzxz(π/2); rtol=1e-10, atol=1e-12) + @test isapprox(byaxis(0.0, 1.0, 0.0, 0.0), Q; rtol=1e-10, atol=1e-12) + + # A half-turn about any axis must square to the identity. + for ax in ((1.0,0.0,0.0), (0.0,1.0,0.0), (0.0,0.0,1.0), (1.0,1.0,0.0), (1.0,1.0,1.0)) + r = SHTnsKit.SHTRotation(cfg.lmax, cfg.mmax) + SHTnsKit.shtns_rotation_set_angle_axis(r, π, ax...) + once = similar(Q); SHTnsKit.shtns_rotation_apply_real(r, Q, once) + twice = similar(Q); SHTnsKit.shtns_rotation_apply_real(r, once, twice) + @test isapprox(twice, Q; rtol=1e-9, atol=1e-11) + end + end + + @testset "order-mixing rotations reject mmax < lmax" begin + # A Wigner-d rotation through a general β couples every |m'| ≤ l. With + # mmax < lmax the |m'| > mmax components were silently dropped — 14.8 % + # of the field energy at lmax=8/mmax=5, 24.0 % at mmax=3. + full = create_gauss_config(8, 10; mmax=8, nlon=18) + cut = create_gauss_config(8, 10; mmax=5, nlon=18) + rng = MersenneTwister(6042) + qf = randn(rng, ComplexF64, full.nlm); qf[1:full.lmax+1] .= real.(qf[1:full.lmax+1]) + qc = randn(rng, ComplexF64, cut.nlm); qc[1:cut.lmax+1] .= real.(qc[1:cut.lmax+1]) + + # mmax == lmax is unaffected and conserves energy. + Rf = similar(qf); SH_Yrotate(full, qf, 0.7, Rf) + w = Float64[full.mi[k] == 0 ? 1.0 : 2.0 for k in 1:full.nlm] + @test sum(w .* abs2.(Rf)) ≈ sum(w .* abs2.(qf)) rtol=1e-10 + + # mmax < lmax with an order-mixing β must raise, not truncate. + @test_throws ArgumentError SH_Yrotate(cut, qc, 0.7, similar(qc)) + @test_throws ArgumentError SH_Xrotate90(cut, qc, similar(qc)) + + # ...but the two non-mixing angles still work at reduced mmax: + # β ≡ 0 is diagonal (a pure Z-rotation) and β ≡ π is anti-diagonal. + for β in (0.0, π) + r = SHTnsKit.SHTRotation(cut.lmax, cut.mmax) + SHTnsKit.shtns_rotation_set_angles_ZYZ(r, 0.4, β, 0.0) + R = similar(qc) + @test (SHTnsKit.shtns_rotation_apply_real(r, qc, R); true) + wc = Float64[cut.mi[k] == 0 ? 1.0 : 2.0 for k in 1:cut.nlm] + @test sum(wc .* abs2.(R)) ≈ sum(wc .* abs2.(qc)) rtol=1e-10 + end + end + @testset "Wigner-d matrix orthogonality" begin # Test orthogonality: d^T d = I for real orthogonal case for l in 0:4 diff --git a/test/serial/test_turbo.jl b/test/serial/test_turbo.jl index 4888db71..cd0afbfb 100644 --- a/test/serial/test_turbo.jl +++ b/test/serial/test_turbo.jl @@ -113,5 +113,54 @@ else @test haskey(b, :synthesis_turbo) @test b.analysis_turbo > 0 end + + @testset "turbo honours cfg.mres" begin + # Regression: both turbo loops iterated a bare `0:mmax`, so for an + # mres>1 config `analysis_turbo` populated coefficient columns the + # transform has no storage for and `synthesis_turbo` summed them + # back in. The disagreement with `analysis`/`synthesis` was O(1), + # not roundoff. + for mres in (2, 3) + cfgm = create_gauss_config(6, 8; mres=mres, nlon=15) + rngm = MersenneTwister(9100 + mres) + fm = randn(rngm, cfgm.nlat, cfgm.nlon) + + a_ref = analysis(cfgm, fm) + a_turbo = SHTnsKit.analysis_turbo(cfgm, fm) + @test a_turbo ≈ a_ref atol=1e-12 + + # Orders absent from the layout must stay exactly zero. + for m in 0:cfgm.mmax + m % mres == 0 && continue + @test all(iszero, @view a_turbo[:, m + 1]) + end + + # Feed a spectrum that DOES carry junk on the skipped columns: + # synthesis must ignore it, exactly as the core does. + dirty = copy(a_ref) + for m in 0:cfgm.mmax + m % mres == 0 && continue + for l in m:cfgm.lmax + dirty[l + 1, m + 1] = 3.0 - 2.0im + end + end + @test SHTnsKit.synthesis_turbo(cfgm, dirty) ≈ synthesis(cfgm, dirty) atol=1e-12 + @test SHTnsKit.synthesis_turbo(cfgm, dirty) ≈ synthesis(cfgm, a_ref) atol=1e-12 + end + end + + @testset "turbo runs inside an outer threaded region" begin + # `@threads :static` cannot be nested; the turbo loops must fall back + # to the caller's thread the way the core orchestrators do. + cfgt = create_gauss_config(5, 7; nlon=11) + ft = randn(MersenneTwister(31), cfgt.nlat, cfgt.nlon) + ref = analysis(cfgt, ft) + out = Vector{Any}(undef, 2) + Threads.@threads for k in 1:2 + out[k] = SHTnsKit.analysis_turbo(cfgt, ft) + end + @test out[1] ≈ ref atol=1e-12 + @test out[2] ≈ ref atol=1e-12 + end end end diff --git a/test/serial/test_vector_transforms.jl b/test/serial/test_vector_transforms.jl index 989e0107..04686c45 100644 --- a/test/serial/test_vector_transforms.jl +++ b/test/serial/test_vector_transforms.jl @@ -8,6 +8,17 @@ using SHTnsKit @isdefined(VERBOSE) || (const VERBOSE = get(ENV, "SHTNSKIT_TEST_VERBOSE", "0") == "1") +# Allocation budgets below are calibrated for a SINGLE-THREADED run, where they +# are exact. The shared m-loop orchestrators start `@threads` tasks whenever +# threads are available, and each threaded region costs a few hundred bytes to +# spawn — a constant independent of problem size, and not the kind of regression +# these budgets exist to catch (a per-row temporary or a dense zero spectrum +# costs tens of KB and grows with lmax). Allow a per-thread slack rather than +# letting the whole suite go red on any multi-core machine; note this makes the +# budgets coarse at `nthreads > 1`, so the tight check is the 1-thread run. +@isdefined(_thread_alloc_slack) || + (_thread_alloc_slack() = Threads.nthreads() > 1 ? 4_096 * Threads.nthreads() : 0) + @testset "Vector Transforms (Spheroidal-Toroidal)" begin @testset "Sphtor roundtrip" begin lmax = 8 @@ -87,8 +98,8 @@ using SHTnsKit synthesis_tor(cfg, Tlm) GC.gc() - @test @allocated(synthesis_sph(cfg, Slm)) <= 15_000 - @test @allocated(synthesis_tor(cfg, Tlm)) <= 15_000 + @test @allocated(synthesis_sph(cfg, Slm)) <= 15_000 + _thread_alloc_slack() + @test @allocated(synthesis_tor(cfg, Tlm)) <= 15_000 + _thread_alloc_slack() end @testset "Gradient transform" begin @@ -453,4 +464,20 @@ using SHTnsKit @test isapprox(Gp_l, Gp_ref; rtol=1e-10, atol=1e-12) end + + @testset "mode-limited sphtor validates its order and truncation" begin + # Out-of-range arguments used to walk off the end of the norm-scale table + # and return ±Inf coefficients; the scalar twin already raised. + cfg = create_gauss_config(3, 5; nlon=8) + v = randn(MersenneTwister(3311), ComplexF64, cfg.nlat) + sl = zeros(ComplexF64, cfg.lmax + 1) + @test_throws ArgumentError SHTnsKit.analysis_sphtor_ml(cfg, 9, v, v, 12) + @test_throws ArgumentError SHTnsKit.analysis_sphtor_ml(cfg, -1, v, v, 3) + @test_throws ArgumentError SHTnsKit.analysis_sphtor_ml(cfg, 1, v, v, 99) + @test_throws ArgumentError SHTnsKit.synthesis_sphtor_ml(cfg, 9, sl, sl, 12) + # a valid call still returns finite coefficients + S, T = SHTnsKit.analysis_sphtor_ml(cfg, 1, v, v, 3) + @test all(isfinite, S) && all(isfinite, T) + end + end