Skip to content

Commit 7431a93

Browse files
xzacklixzackli
committed
new healpix workspace
1 parent 1902953 commit 7431a93

File tree

12 files changed

+799
-651
lines changed

12 files changed

+799
-651
lines changed

‎src/XGPaint.jl‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ abstract type AbstractInterpolatorProfile{T} <: AbstractProfile{T} end
3333

3434

3535
include("./util.jl")
36+
include("./healpixworkspace.jl")
3637
include("./model.jl")
3738
include("./profiles.jl")
3839
include("./profiles_y.jl")

‎src/healpixworkspace.jl‎

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
2+
# Ring workspace for fast disc queries
3+
struct RingWorkspace{T<:AbstractFloat}
4+
res::Healpix.Resolution
5+
ring_thetas::Vector{T} # Colatitude of each ring
6+
ring_shifts::Vector{T} # Shift offset for each ring (0.0 or 0.5)
7+
ring_num_pixels::Vector{Int} # Number of pixels in each ring
8+
ring_first_pixels::Vector{Int} # First pixel index for each ring
9+
end
10+
11+
function RingWorkspace{T}(res::Healpix.Resolution) where T<:AbstractFloat
12+
num_rings = res.nsideTimesFour - 1
13+
14+
ring_thetas = Vector{T}(undef, num_rings)
15+
ring_shifts = Vector{T}(undef, num_rings)
16+
ring_num_pixels = Vector{Int}(undef, num_rings)
17+
ring_first_pixels = Vector{Int}(undef, num_rings)
18+
19+
for ring_idx in 1:num_rings
20+
ringinfo = Healpix.getringinfo(res, ring_idx)
21+
ring_thetas[ring_idx] = T(ringinfo.colatitude_rad)
22+
ring_shifts[ring_idx] = ringinfo.shifted ? T(0.5) : T(0.0)
23+
ring_num_pixels[ring_idx] = ringinfo.numOfPixels
24+
ring_first_pixels[ring_idx] = ringinfo.firstPixIdx
25+
end
26+
27+
RingWorkspace{T}(res, ring_thetas, ring_shifts, ring_num_pixels, ring_first_pixels)
28+
end
29+
30+
RingWorkspace(res::Healpix.Resolution) = RingWorkspace{Float64}(res)
31+
32+
function get_relevant_rings(res::Healpix.Resolution, θ_center, radius)
33+
z_top = cos(max(0, θ_center - radius))
34+
z_bottom = cos(min(π, θ_center + radius))
35+
36+
ring_top = Healpix.ringAbove(res, z_top)
37+
ring_bottom = Healpix.ringAbove(res, z_bottom) + 1
38+
39+
return max(1, ring_top), min(res.nsideTimesFour - 1, ring_bottom)
40+
end
41+
42+
"""
43+
get_ring_disc_ranges(workspace, ring_idx, θ_center, ϕ_center, radius)
44+
45+
Returns (range1, range2) of pixel indices on a ring that lie within the disc.
46+
47+
The spherical distance between points (θ₁,φ₁) and (θ₂,φ₂) is given by:
48+
cos(d) = cos(θ₁)cos(θ₂) + sin(θ₁)sin(θ₂)cos(φ₁ - φ₂)
49+
50+
For a point to be included in the disc: d ≤ radius.
51+
52+
Returns two ranges to handle φ wraparound at the 0/2π boundary:
53+
- range1: Main contiguous range of pixels
54+
- range2: Additional range when disc crosses φ=0/2π seam (empty if no crossing)
55+
"""
56+
function get_ring_disc_ranges(workspace::RingWorkspace, ring_idx::Int,
57+
center_theta::T, center_phi::T, radius::T) where T
58+
nr = workspace.ring_num_pixels[ring_idx]
59+
ring_theta = workspace.ring_thetas[ring_idx]
60+
shift = workspace.ring_shifts[ring_idx]
61+
62+
# Polar cap regions require special handling
63+
is_north_cap = ring_idx ≤ workspace.res.nside
64+
is_south_cap = ring_idx ≥ 3 * workspace.res.nside
65+
66+
if is_north_cap || is_south_cap
67+
# Exact pole centers: simple distance check
68+
if abs(center_theta) < T(1e-10)
69+
return ring_theta ≤ radius ? (1:nr, 1:0) : (1:0, 1:0)
70+
end
71+
if abs(center_theta - T(π)) < T(1e-10)
72+
return (T(π) - ring_theta) ≤ radius ? (1:nr, 1:0) : (1:0, 1:0)
73+
end
74+
75+
# Solve spherical triangle for phi half-width
76+
cos_center, sin_center = cos(center_theta), sin(center_theta)
77+
cos_ring, sin_ring = cos(ring_theta), sin(ring_theta)
78+
cos_radius = cos(radius)
79+
80+
denominator = sin_ring * sin_center
81+
if abs(denominator) < T(1e-12)
82+
condition = abs(cos_radius - cos_ring * cos_center) < T(1e-12)
83+
return condition ? (1:nr, 1:0) : (1:0, 1:0)
84+
end
85+
86+
cos_delta_phi = (cos_radius - cos_ring * cos_center) / denominator
87+
if cos_delta_phi > T(1); return (1:0, 1:0); end
88+
if cos_delta_phi < T(-1); return (1:nr, 1:0); end
89+
90+
delta_phi = acos(cos_delta_phi)
91+
phi_min, phi_max = center_phi - delta_phi, center_phi + delta_phi
92+
93+
# Convert phi to pixel indices
94+
ip_lo = floor(Int, nr / T(2π) * phi_min - shift) + 1
95+
ip_hi = floor(Int, nr / T(2π) * phi_max - shift)
96+
if ip_lo > ip_hi; return (1:0, 1:0); end
97+
98+
# Handle phi wraparound at 0/2π boundary
99+
if ip_hi >= nr; ip_lo -= nr; ip_hi -= nr; end
100+
return ip_lo < 0 ? (1:(ip_hi + 1), (ip_lo + nr + 1):nr) : ((ip_lo + 1):(ip_hi + 1), 1:0)
101+
else
102+
# Equatorial region: standard HEALPix algorithm
103+
z, z0 = cos(ring_theta), cos(center_theta)
104+
xa = T(1) / sqrt((T(1) - z0) * (T(1) + z0))
105+
x = (cos(radius) - z * z0) * xa
106+
ysq = T(1) - z^2 - x^2
107+
108+
dphi = ysq < T(0) ? T(0) : atan(sqrt(ysq), x)
109+
if dphi ≤ T(0); return (1:0, 1:0); end
110+
111+
ip_lo = floor(Int, nr / T(2π) * (center_phi - dphi) - shift) + 1
112+
ip_hi = floor(Int, nr / T(2π) * (center_phi + dphi) - shift)
113+
114+
if ip_lo > ip_hi; return (1:0, 1:0); end
115+
if ip_hi >= nr; ip_lo -= nr; ip_hi -= nr; end
116+
117+
return ip_lo < 0 ? (1:(ip_hi + 1), (ip_lo + nr + 1):nr) : ((ip_lo + 1):(ip_hi + 1), 1:0)
118+
end
119+
end

‎src/profiles.jl‎

Lines changed: 62 additions & 78 deletions
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,6 @@ function profile_grid(model::AbstractGNFW{T}, logθs, redshifts, logMs) where T
6565
end
6666

6767

68-
6968
"""
7069
Computes a real-space beam interpolator and a maximum
7170
"""
@@ -86,72 +85,17 @@ function realspacegaussbeam(::Type{T}, θ_FWHM::Ti; rtol=1e-24, N_θ::Int=2000)
8685
end
8786

8887

89-
# heuristic for sizehint
90-
_approx_discbuffer_size(nside, θmax) = ceil(Int, 1.1 * π * θmax^2 / (nside2pixarea(nside)))
91-
92-
struct HealpixSerialProfileWorkspace{T} <: AbstractProfileWorkspace{T}
93-
nside::Int
94-
ringinfo::RingInfo
95-
disc_buffer::Vector{Int}
96-
θmax::T
97-
posmap::Vector{Tuple{T, T, T}}
98-
99-
HealpixSerialProfileWorkspace{T}(nside::Int, θmax) where T = new{T}(
100-
nside,
101-
RingInfo(0, 0, 0, 0.0, true),
102-
sizehint!(Int[], _approx_discbuffer_size(nside, θmax)),
103-
T(θmax),
104-
vectorhealpixmap(T, nside)
105-
)
106-
107-
HealpixSerialProfileWorkspace{T}(nside, ringinfo, disc_buffer, θmax, posmap) where T = new{T}(
108-
nside, ringinfo, disc_buffer, T(θmax), posmap)
109-
end
110-
111-
# default outer constructors for Float64
112-
HealpixSerialProfileWorkspace(nside::Int, θmax) = HealpixSerialProfileWorkspace{Float64}(nside, θmax)
113-
114-
function Base.show(io::IO, ::HealpixSerialProfileWorkspace{T}) where T
115-
expr = "HealpixSerialProfileWorkspace{$(T)}"
116-
print(io, expr)
117-
end
118-
119-
struct HealpixProfileWorkspace{T} <: AbstractProfileWorkspace{T}
120-
nside::Int
121-
ringinfo::RingInfo
122-
disc_buffer::Vector{Vector{Int}}
123-
θmax::T
124-
posmap::Vector{Tuple{T, T, T}}
125-
126-
HealpixProfileWorkspace{T}(nside::Int, θmax, nthreads=Threads.nthreads()) where T = new{T}(
127-
nside,
128-
RingInfo(0, 0, 0, 0.0, true),
129-
Vector{Int}[sizehint!(Int[], _approx_discbuffer_size(nside, θmax)) for i in 1:nthreads],
130-
T(θmax),
131-
vectorhealpixmap(T, nside)
132-
)
133-
end
134-
135-
HealpixProfileWorkspace(nside::Int, θmax) = HealpixProfileWorkspace{Float64}(nside, θmax)
136-
137-
# unlike other workspaces, Healpix workspace needs a mutable buffer per thread
138-
# so asking for a workspace with a thread id will return the appropriate mutable workspace
139-
function wrapserialworkspace(w::HealpixProfileWorkspace{T}, tid) where T
140-
return HealpixSerialProfileWorkspace{T}(w.nside, w.ringinfo, w.disc_buffer[tid], w.θmax, w.posmap)
141-
end
88+
# function realspacebeampaint!(hp_map, w::HealpixSerialProfileWorkspace, realprofile, flux, θ₀, ϕ₀)
89+
# x₀, y₀, z₀ = ang2vec(θ₀, ϕ₀)
90+
# XGPaint.queryDiscRing!(w.disc_buffer, w.ringinfo, hp_map.resolution, θ₀, ϕ₀, w.θmax)
14291

143-
144-
function realspacebeampaint!(hp_map, w::HealpixSerialProfileWorkspace, realprofile, flux, θ₀, ϕ₀)
145-
x₀, y₀, z₀ = ang2vec(θ₀, ϕ₀)
146-
XGPaint.queryDiscRing!(w.disc_buffer, w.ringinfo, hp_map.resolution, θ₀, ϕ₀, w.θmax)
147-
148-
for ir in w.disc_buffer
149-
x₁, y₁, z₁ = w.posmap.pixels[ir]
150-
d² = (x₁ - x₀)^2 + (y₁ - y₀)^2 + (z₁ - z₀)^2
151-
θ = acos(1 - d² / 2)
152-
hp_map.pixels[ir] += flux * realprofile(θ)
153-
end
154-
end
92+
# for ir in w.disc_buffer
93+
# x₁, y₁, z₁ = w.posmap.pixels[ir]
94+
# d² = (x₁ - x₀)^2 + (y₁ - y₀)^2 + (z₁ - z₀)^2
95+
# θ = acos(1 - d² / 2)
96+
# hp_map.pixels[ir] += flux * realprofile(θ)
97+
# end
98+
# end
15599

156100

157101
"""Apply a beam to a profile grid"""
@@ -385,26 +329,66 @@ function profile_paint!(m::Enmap{T, 2, Matrix{T}, Gnomonic{T}},
385329
end
386330

387331

388-
function profile_paint_generic!(m::HealpixMap{T, RingOrder}, w::HealpixSerialProfileWorkspace,
332+
# function profile_paint_generic!(m::HealpixMap{T, RingOrder}, w::HealpixSerialProfileWorkspace,
333+
# model, Mh, z, α₀, δ₀, θmax, normalization=1) where T
334+
# ϕ₀ = α₀
335+
# θ₀ = T(π)/2 - δ₀
336+
# x₀, y₀, z₀ = ang2vec(θ₀, ϕ₀)
337+
# θmin = compute_θmin(model)
338+
# XGPaint.queryDiscRing!(w.disc_buffer, w.ringinfo, m.resolution, θ₀, ϕ₀, θmax)
339+
# for ir in w.disc_buffer
340+
# x₁, y₁, z₁ = w.posmap[ir]
341+
# d² = (x₁ - x₀)^2 + (y₁ - y₀)^2 + (z₁ - z₀)^2
342+
# θ = acos(clamp(1 - d² / 2, -one(T), one(T)))
343+
# θ = max(θmin, θ) # clamp to minimum θ
344+
# m.pixels[ir] += ifelse(θ < θmax,
345+
# normalization * model(θ, Mh, z),
346+
# zero(T))
347+
# end
348+
# end
349+
350+
function profile_paint_generic!(m::HealpixMap{T, RingOrder}, workspace::RingWorkspace{T},
389351
model, Mh, z, α₀, δ₀, θmax, normalization=1) where T
390-
ϕ₀ = α₀
352+
ϕ₀ = α₀
391353
θ₀ = T(π)/2 - δ₀
392354
x₀, y₀, z₀ = ang2vec(θ₀, ϕ₀)
393355
θmin = compute_θmin(model)
394-
XGPaint.queryDiscRing!(w.disc_buffer, w.ringinfo, m.resolution, θ₀, ϕ₀, θmax)
395-
for ir in w.disc_buffer
396-
x₁, y₁, z₁ = w.posmap[ir]
397-
d² = (x₁ - x₀)^2 + (y₁ - y₀)^2 + (z₁ - z₀)^2
398-
θ = acos(clamp(1 - d² / 2, -one(T), one(T)))
399-
θ = max(θmin, θ) # clamp to minimum θ
400-
m.pixels[ir] += ifelse(θ < θmax,
401-
normalization * model(θ, Mh, z),
402-
zero(T))
356+
357+
# Get relevant rings for this disc
358+
ring_start, ring_end = get_relevant_rings(workspace.res, θ₀, θmax)
359+
360+
for ring_idx in ring_start:ring_end
361+
# Get pixel ranges on this ring that intersect the disc
362+
range1, range2 = get_ring_disc_ranges(workspace, ring_idx, θ₀, ϕ₀, θmax)
363+
364+
# Get precomputed ring info
365+
first_pixel = workspace.ring_first_pixels[ring_idx]
366+
367+
# Process both ranges (range2 may be empty for no phi wraparound)
368+
for pixel_range in (range1, range2)
369+
for pix_idx in pixel_range
370+
# Convert ring pixel index to global healpix pixel index
371+
global_pix = first_pixel + pix_idx - 1
372+
373+
# Get position of this pixel
374+
x₁, y₁, z₁ = pix2vecRing(workspace.res, global_pix)
375+
376+
# Compute angular distance
377+
d² = (x₁ - x₀)^2 + (y₁ - y₀)^2 + (z₁ - z₀)^2
378+
θ = acos(clamp(1 - d² / 2, -one(T), one(T)))
379+
θ = max(θmin, θ) # clamp to minimum θ
380+
381+
# Add contribution to map
382+
m.pixels[global_pix] += ifelse(θ < θmax,
383+
normalization * model(θ, Mh, z),
384+
zero(T))
385+
end
386+
end
403387
end
404388
end
405389

406390
# fall back to generic profile painter if no specialized painter is defined for the model
407-
function profile_paint!(m::HealpixMap{T, RingOrder}, w::HealpixSerialProfileWorkspace, model,
391+
function profile_paint!(m::HealpixMap{T, RingOrder}, w::RingWorkspace{T}, model,
408392
Mh, z, α₀, δ₀, θmax, normalization=1) where T
409393
profile_paint_generic!(m, w, model, Mh, z, α₀, δ₀, θmax, normalization)
410394
end

‎src/profiles_rsz.jl‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -133,8 +133,8 @@ end
133133

134134
# like the usual paint, but use the sign of the null as the sign of the perturbation
135135
function profile_paint!(m::HealpixMap{T, RingOrder},
136-
workspace::HealpixSerialProfileWorkspace, model::RSZPerturbativeProfile,
136+
workspace::RingWorkspace{T}, model::RSZPerturbativeProfile,
137137
Mh, z, α₀, δ₀, θmax) where T
138138
X_0 = calc_null(model, Mh, z)
139-
profile_paint_generic!(m, workspace, model, Mh, z, α₀, δ₀, θmax, sign(X - X_0))
139+
profile_paint_generic!(m, workspace, model, Mh, z, α₀, δ₀, θmax, sign(model.X - X_0))
140140
end

‎src/profiles_szp.jl‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -99,8 +99,8 @@ function profile_paint!(m::Enmap{T, 2, Matrix{T}, CarClenshawCurtis{T}},
9999
end
100100

101101
function profile_paint!(m::HealpixMap{T, RingOrder},
102-
workspace::HealpixSerialProfileWorkspace, model::SZPackRSZProfile,
102+
workspace::RingWorkspace{T}, model::SZPackRSZProfile,
103103
Mh, z, α₀, δ₀, θmax) where T
104-
rsz_factor_I_over_y = compute_rsz_factor_I_over_y(model_szp, Mh, z)
105-
profile_paint_generic!(m, workspace, model, Mh, z, α₀, δ₀, θmax, rsz_factor_I_over_y)
104+
rsz_factor_I_over_y = compute_rsz_factor_I_over_y(model, Mh, z)
105+
profile_paint_generic!(m, workspace, model.y_model_interp, Mh, z, α₀, δ₀, θmax, rsz_factor_I_over_y)
106106
end

0 commit comments

Comments
 (0)