# MIT License # # Copyright (c) 2017–2022: Justin Willmert # # See https://justinwillmert.com/articles/2022/solving-for-spherical-harmonic-analysis-quadrature-weights/ # and https://github.com/jmert/CMB.jl using AssociatedLegendrePolynomials: λlm! using FFTW using FastGaussQuadrature: gausslegendre struct RingInfo offsets::Tuple{Int,Int} stride::Int nϕ::Int cosθ::Float64 ϕ₀::Float64 ΔΩ::Float64 end struct MapRings rings::Vector{RingInfo} end """ maprings = ecp_maprings(nθ, nϕ) Constructs a `MapRings` structure describing an equidistant cylindrical projection (ECP) grid of dimensions `nθ × nϕ` pixels in column-major ordering. """ function ecp_maprings(nθ, nϕ) nθh = (nθ + 1) ÷ 2 # number of rings in N. hemisphere + equator (if nθ is odd) ϕ₀ = 2π / 2nϕ # constant for all rings ΔθΔϕ = (2π / nϕ) * (π / nθ) # constant portion of ΔΩ = sinθ * ΔθΔϕ stride = nθ # for column-major ordering rings = Vector{RingInfo}(undef, nθh) for ii in 1:nθh o₁ = ii # northern-ring offset o₂ = nθ - ii + 1 # southern-ring offset offsets = (o₁, o₁ == o₂ ? 0 : o₂) # for equator of odd nθ, disable southern ring θ_π = (ii - 0.5) / nθ # θ/π for the ring sθ, cθ = sincospi(θ_π) ΔΩ = sθ * ΔθΔϕ rings[ii] = RingInfo(offsets, stride, nϕ, cθ, ϕ₀, ΔΩ) end return MapRings(rings) end """ maprings = gl_maprings(nθ, nϕ) Constructs a `MapRings` structure describing an Gauss-Legendre (GL) grid of dimensions `nθ × nϕ` pixels in column-major ordering. """ function gl_maprings(nθ, nϕ) nθh = (nθ + 1) ÷ 2 # number of rings in N. hemisphere + equator (if nθ is odd) ϕ₀, Δϕ = 2π / 2nϕ, 2π / nϕ # constant for all rings cθ, sθΔθ = gausslegendre(nθ) # z = cos(θ) and corresponding weights stride = nθ # for column-major ordering rings = Vector{RingInfo}(undef, nθh) for ii in 1:nθh o₁ = ii # northern-ring offset o₂ = nθ - ii + 1 # southern-ring offset offsets = (o₁, o₁ == o₂ ? 0 : o₂) # for equator of odd nθ, disable southern ring # N.B. z ∈ [-1, 1] and θ ∈ [0, π] run in opposite directions, and cos(π-θ) == -z rings[ii] = RingInfo(offsets, stride, nϕ, -cθ[ii], ϕ₀, sθΔθ[ii] * Δϕ) end return MapRings(rings) end """ (; npix, nϕmax) = _mapring_counts(maprings::MapRings) Counts total number of pixels (`npix`) described by `maprings` and the longest observed ring (`nϕmax`). """ function _mapring_counts(maprings::MapRings) nϕmax, npix = 0, 0 for ring in maprings.rings mult = mapreduce(!=(0), +, ring.offsets) # single or paired rings? npix += mult * ring.nϕ # add 1 or 2 times Nϕ to total count nϕmax = max(nϕmax, ring.nϕ) # keep track of longest seen ring end return (; npix, nϕmax) end """ (i′, isconj, isnyq) alias_index(len::Int, i::Int) Calculates aliased index `i′` given original index `i` and Fourier axis length 'len'. Also indicates whether conjugation (`isconj`) or twice-real symmetries (`isnyq`) are needed. """ @inline function alias_index(len::Int, i::Int) isconj, isnyq = false, false nyq = max(1, len ÷ 2) i < nyq && return (i, isconj, isnyq) i = mod(i, len) if i > nyq i = len - i isconj = true elseif i == 0 || (iseven(len) && i == nyq) isnyq = true end return (i, isconj, isnyq) end """ coeffs′ = alias_coeffs(coeffs::NTuple, isconj::Bool, isnyq::Bool) Applies conjugation (`isconj`) or twice-real (`isnyq`) transformation, if required, in mode aliasing to the harmonic coefficients `coeffs`. """ @inline function alias_coeffs(coeffs::NTuple, isconj::Bool, isnyq::Bool) return ifelse(isnyq, complex.(2 .* real.(coeffs)), ifelse(isconj, conj.(coeffs), coeffs)) end function synthesize(maprings::MapRings, alms::Matrix{C}) where {C<:Complex} R = real(C) lmax, mmax = size(alms) .- 1 (; npix, nϕmax) = _mapring_counts(maprings) nϕr = nϕmax ÷ 2 + 1 # (maximum) real-symmetric FFT's Nyquist length (index) mapbuf = Vector{R}(undef, npix) Λ = Matrix{R}(undef, lmax + 1, mmax + 1) λ₁ = zeros(C, nϕr) # phase factors for northern ring λ₂ = zeros(C, nϕr) # phase factors for southern ring for ring in maprings.rings nϕ, ϕ₀, cθ = ring.nϕ, ring.ϕ₀, ring.cosθ λlm!(Λ, lmax, mmax, cθ) for m in 0:mmax acc₁, acc₂ = zero(C), zero(C) # Σ_{ℓ = m}^{lmax} for ℓ in m:lmax term = alms[ℓ+1,m+1] * Λ[ℓ+1,m+1] acc₁ += term acc₂ += isodd(ℓ + m) ? -term : term end acc₁, acc₂ = (acc₁, acc₂) .* cis(m * ϕ₀) # ring offset rotation i, isconj, isnyq = alias_index(nϕ, m) acc₁, acc₂ = alias_coeffs((acc₁, acc₂), isconj, isnyq) λ₁[i+1] += acc₁ λ₂[i+1] += acc₂ end stride = ring.stride o₁, o₂ = ring.offsets nϕh = nϕ ÷ 2 + 1 if o₁ != 0 idx = range(o₁, step = stride, length = nϕ) @views mapbuf[idx] = brfft(λ₁[1:nϕh], nϕ) end if o₂ != 0 idx = range(o₂, step = stride, length = nϕ) @views mapbuf[idx] = brfft(λ₂[1:nϕh], nϕ) end λ₁ .= zero(C) λ₂ .= zero(C) end return mapbuf end # function synthesize function analyze(maprings::MapRings, mapbuf::Vector{R}, lmax::Integer, mmax::Integer = lmax) where {R<:Real} C = complex(R) (; npix, nϕmax) = _mapring_counts(maprings) nϕr = nϕmax ÷ 2 + 1 # (maximum) real-symmetric FFT's Nyquist length (index) alms = fill(zero(C), lmax + 1, mmax + 1) Λ = zeros(R, lmax + 1, mmax + 1) λ₁ = zeros(C, nϕr) # phase factors for northern ring λ₂ = zeros(C, nϕr) # phase factors for southern ring for ring in maprings.rings nϕ, ϕ₀, cθ, ΔΩ = ring.nϕ, ring.ϕ₀, ring.cosθ, ring.ΔΩ λlm!(Λ, lmax, mmax, cθ) # λ_ℓ^m(cos θ) factors stride = ring.stride o₁, o₂ = ring.offsets nϕh = nϕ ÷ 2 + 1 if o₁ != 0 idx = range(o₁, step = stride, length = nϕ) @views λ₁[1:nϕh] = rfft(mapbuf[idx]) end if o₂ != 0 idx = range(o₂, step = stride, length = nϕ) @views λ₂[1:nϕh] = rfft(mapbuf[idx]) end for m in 0:mmax i, isconj, isnyq = alias_index(nϕ, m) a₁, a₂ = alias_coeffs((λ₁[i+1], λ₂[i+1]), isconj, isnyq) a₁, a₂ = (a₁, a₂) .* (ΔΩ * cis(m * -ϕ₀)) # Σ_{ℓ = m}^{lmax} for ℓ in m:lmax c = isodd(ℓ+m) ? a₁ - a₂ : a₁ + a₂ alms[ℓ+1,m+1] += c * Λ[ℓ+1,m+1] end end λ₁ .= zero(C) λ₂ .= zero(C) end return alms end # function analyze function alm2cl(alms::Matrix{<:Complex{R}}) where R lmax, mmax = size(alms) .- 1 Cl = zeros(R, lmax + 1) for ll in 0:lmax Cl[ll+1] = abs2(alms[ll+1,1]) for mm in 1:min(ll, mmax) Cl[ll+1] += 2 * abs2(alms[ll+1,mm+1]) end Cl[ll+1] /= 2ll + 1 end return Cl end