|
| 1 | + |
| 2 | +# Ring workspace for fast disc queries; this is designed for CPU for now |
| 3 | +struct HealpixRingProfileWorkspace{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 HealpixRingProfileWorkspace{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 | + HealpixRingProfileWorkspace{T}(res, ring_thetas, ring_shifts, ring_num_pixels, ring_first_pixels) |
| 28 | +end |
| 29 | + |
| 30 | +HealpixRingProfileWorkspace(res::Healpix.Resolution) = HealpixRingProfileWorkspace{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::HealpixRingProfileWorkspace{T}, ring_idx::Int, |
| 57 | + center_theta::T, center_phi::T, radius::T) where T |
| 58 | + # Validate that phi is in the valid range [0, 2π) |
| 59 | + @assert 0 <= center_phi < T(2π) "center_phi must be in the range [0, 2π), got $(center_phi)" |
| 60 | + |
| 61 | + nr = workspace.ring_num_pixels[ring_idx] |
| 62 | + ring_theta = workspace.ring_thetas[ring_idx] |
| 63 | + shift = workspace.ring_shifts[ring_idx] |
| 64 | + |
| 65 | + # Polar cap regions require special handling |
| 66 | + is_north_cap = ring_idx ≤ workspace.res.nside |
| 67 | + is_south_cap = ring_idx ≥ 3 * workspace.res.nside |
| 68 | + |
| 69 | + if is_north_cap || is_south_cap |
| 70 | + # Exact pole centers: simple distance check |
| 71 | + if abs(center_theta) < T(1e-10) |
| 72 | + return ring_theta ≤ radius ? (1:nr, 1:0) : (1:0, 1:0) |
| 73 | + end |
| 74 | + if abs(center_theta - T(π)) < T(1e-10) |
| 75 | + return (T(π) - ring_theta) ≤ radius ? (1:nr, 1:0) : (1:0, 1:0) |
| 76 | + end |
| 77 | + |
| 78 | + # Solve spherical triangle for phi half-width |
| 79 | + cos_center, sin_center = cos(center_theta), sin(center_theta) |
| 80 | + cos_ring, sin_ring = cos(ring_theta), sin(ring_theta) |
| 81 | + cos_radius = cos(radius) |
| 82 | + |
| 83 | + denominator = sin_ring * sin_center |
| 84 | + if abs(denominator) < T(1e-12) |
| 85 | + condition = abs(cos_radius - cos_ring * cos_center) < T(1e-12) |
| 86 | + return condition ? (1:nr, 1:0) : (1:0, 1:0) |
| 87 | + end |
| 88 | + |
| 89 | + cos_delta_phi = (cos_radius - cos_ring * cos_center) / denominator |
| 90 | + if cos_delta_phi > T(1); return (1:0, 1:0); end |
| 91 | + if cos_delta_phi < T(-1); return (1:nr, 1:0); end |
| 92 | + |
| 93 | + delta_phi = acos(cos_delta_phi) |
| 94 | + phi_min, phi_max = center_phi - delta_phi, center_phi + delta_phi |
| 95 | + |
| 96 | + # Convert phi to pixel indices |
| 97 | + ip_lo = floor(Int, nr / T(2π) * phi_min - shift) + 1 |
| 98 | + ip_hi = floor(Int, nr / T(2π) * phi_max - shift) |
| 99 | + if ip_lo > ip_hi; return (1:0, 1:0); end |
| 100 | + |
| 101 | + # Handle phi wraparound at 0/2π boundary |
| 102 | + if ip_hi >= nr; ip_lo -= nr; ip_hi -= nr; end |
| 103 | + return ip_lo < 0 ? (1:(ip_hi + 1), (ip_lo + nr + 1):nr) : ((ip_lo + 1):(ip_hi + 1), 1:0) |
| 104 | + else |
| 105 | + # Equatorial region: standard HEALPix algorithm |
| 106 | + z, z0 = cos(ring_theta), cos(center_theta) |
| 107 | + xa = T(1) / sqrt((T(1) - z0) * (T(1) + z0)) |
| 108 | + x = (cos(radius) - z * z0) * xa |
| 109 | + ysq = T(1) - z^2 - x^2 |
| 110 | + |
| 111 | + dphi = ysq < T(0) ? T(0) : atan(sqrt(ysq), x) |
| 112 | + if dphi ≤ T(0); return (1:0, 1:0); end |
| 113 | + |
| 114 | + ip_lo = floor(Int, nr / T(2π) * (center_phi - dphi) - shift) + 1 |
| 115 | + ip_hi = floor(Int, nr / T(2π) * (center_phi + dphi) - shift) |
| 116 | + |
| 117 | + if ip_lo > ip_hi; return (1:0, 1:0); end |
| 118 | + if ip_hi >= nr; ip_lo -= nr; ip_hi -= nr; end |
| 119 | + |
| 120 | + return ip_lo < 0 ? (1:(ip_hi + 1), (ip_lo + nr + 1):nr) : ((ip_lo + 1):(ip_hi + 1), 1:0) |
| 121 | + end |
| 122 | +end |
0 commit comments