Skip to content

Commit f8ce829

Browse files
marius311marius311
committed
add Healpix field types and improve projections
1 parent 9a27327 commit f8ce829

File tree

3 files changed

+184
-49
lines changed

3 files changed

+184
-49
lines changed

‎src/flat_fields.jl‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -175,6 +175,7 @@ require_unbatched(f::FlatField) = (f.Nbatch==1) || error("This function not impl
175175
getproperty(f::FlatField, ::Val{:Nbatch}) = size(getfield(f,:arr), 4)
176176
getproperty(f::FlatField, ::Val{:Npol}) = size(getfield(f,:arr), 3)
177177
getproperty(f::FlatField, ::Val{:T}) = eltype(f)
178+
getproperty(f::FlatField, ::Val{:proj}) = getfield(f, :metadata)
178179
# sub-components
179180
for (F, props) in _sub_components
180181
for (k,I) in props
@@ -188,7 +189,6 @@ for (F, props) in _sub_components
188189
end
189190

190191

191-
192192
### indices
193193
function getindex(f::FlatS0, k::Symbol; full_plane=false)
194194
maybe_unfold = full_plane ? x->unfold(x,fieldinfo(f).Ny) : identity
@@ -321,7 +321,7 @@ Map(f::FlatField{B}) where {B<:BasisProd} = Map(B)(f)
321321

322322

323323
### pretty printing
324-
typealias_def(::Type{F}) where {B,M,T,A,F<:FlatField{B,M,T,A}} = "Flat$(typealias(B)){$(typealias(A)),$(typealias(M))}"
324+
typealias_def(::Type{F}) where {B,M<:FlatProj,T,A,F<:FlatField{B,M,T,A}} = "Flat$(typealias(B)){$(typealias(A)),$(typealias(M))}"
325325
function Base.summary(io::IO, f::FlatField)
326326
@unpack Nx,Ny,Nbatch,θpix = f
327327
print(io, "$(length(f))-element $Ny×$Nx$(Nbatch==1 ? "" : "(×$Nbatch)")-pixel $(θpix)′-resolution ")

‎src/healpix.jl‎

Lines changed: 175 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -1,68 +1,197 @@
11

2+
# Some initial support for Healpix fields.
3+
# The main functionality of broadcasting, indexing, and projection for
4+
# a few field types is implemented, but not much beyond that.
5+
6+
27
@init global hp = lazy_pyimport("healpy")
38

4-
function θϕ_to_xy((θ,ϕ))
5-
r = 2cos(θ/2)
6-
x = r*cos(ϕ)
7-
y = -r*sin(ϕ)
8-
x, y
9+
struct ProjHealpix <: FieldMetadata
10+
Nside :: Int
11+
end
12+
typealias_def(::Type{<:ProjHealpix}) = "ProjHealpix"
13+
14+
15+
# spin-0
16+
const HealpixMap{ M<:ProjHealpix, T, A<:AbstractArray{T} } = BaseField{Map, M, T, A}
17+
const HealpixFourier{ M<:ProjHealpix, T, A<:AbstractArray{T} } = BaseField{Fourier, M, T, A}
18+
# spin-2
19+
const HealpixQUMap{ M<:ProjHealpix, T, A<:AbstractArray{T} } = BaseField{QUMap, M, T, A}
20+
const HealpixQUFourier{ M<:ProjHealpix, T, A<:AbstractArray{T} } = BaseField{QUFourier, M, T, A}
21+
const HealpixEBMap{ M<:ProjHealpix, T, A<:AbstractArray{T} } = BaseField{EBMap, M, T, A}
22+
const HealpixEBFourier{ M<:ProjHealpix, T, A<:AbstractArray{T} } = BaseField{EBFourier, M, T, A}
23+
# spin-(0,2)
24+
const HealpixIQUMap{ M<:ProjHealpix, T, A<:AbstractArray{T} } = BaseField{IQUMap, M, T, A}
25+
const HealpixIQUFourier{ M<:ProjHealpix, T, A<:AbstractArray{T} } = BaseField{IQUFourier, M, T, A}
26+
const HealpixIEBMap{ M<:ProjHealpix, T, A<:AbstractArray{T} } = BaseField{IEBMap, M, T, A}
27+
const HealpixIEBFourier{ M<:ProjHealpix, T, A<:AbstractArray{T} } = BaseField{IEBFourier, M, T, A}
28+
29+
const HealpixField{B, M<:ProjHealpix, T, A<:AbstractArray{T}} = BaseField{B, M, T, A}
30+
31+
## constructing from arrays
32+
# spin-0
33+
function HealpixMap(I::A) where {T, A<:AbstractArray{T}}
34+
HealpixMap(I, ProjHealpix(hp.npix2nside(length(I))))
35+
end
36+
# spin-2
37+
function HealpixField{B}(X::A, Y::A) where {T, A<:AbstractArray{T}, B<:Basis2Prod{<:Union{𝐐𝐔,𝐄𝐁},Map}}
38+
HealpixField{B}(cat(X, Y, dims=Val(2)), ProjHealpix(hp.npix2nside(length(X))))
39+
end
40+
# spin-(0,2)
41+
function HealpixField{B}(I::A, X::A, Y::A) where {T, A<:AbstractArray{T}, B<:Basis3Prod{𝐈,<:Union{𝐐𝐔,𝐄𝐁},Map}}
42+
HealpixField{B}(cat(I, X, Y, dims=Val(2)), ProjHealpix(hp.npix2nside(length(I))))
943
end
1044

11-
function xy_to_θϕ((x,y))
45+
### pretty printing
46+
typealias_def(::Type{F}) where {B,M<:ProjHealpix,T,A,F<:HealpixField{B,M,T,A}} = "Healpix$(typealias(B)){$(typealias(A)),$(typealias(M))}"
47+
function Base.summary(io::IO, f::HealpixField)
48+
@unpack Nside = f
49+
print(io, "$(length(f))-element Nside=$Nside ")
50+
Base.showarg(io, f, true)
51+
end
52+
53+
# useful for enumerating some cases below
54+
_healpix_sub_components = [
55+
(:HealpixMap, ("Ix"=>:, "I"=>:)),
56+
(:HealpixFourier, ("Il"=>:, "I"=>:)),
57+
(:HealpixQUMap, ("Qx"=>1, "Ux"=>2, "Q" =>1, "U"=>2, "P"=>:)),
58+
(:HealpixQUFourier, ("Ql"=>1, "Ul"=>2, "Q" =>1, "U"=>2, "P"=>:)),
59+
(:HealpixEBMap, ("Ex"=>1, "Bx"=>2, "E" =>1, "B"=>2, "P"=>:)),
60+
(:HealpixEBFourier, ("El"=>1, "Bl"=>2, "E" =>1, "B"=>2, "P"=>:)),
61+
(:HealpixIQUMap, ("Ix"=>1, "Qx"=>2, "Ux"=>3, "I"=>1, "Q"=>2, "U"=>3, "P"=>2:3, "IP"=>:)),
62+
(:HealpixIQUFourier, ("Il"=>1, "Ql"=>2, "Ul"=>3, "I"=>1, "Q"=>2, "U"=>3, "P"=>2:3, "IP"=>:)),
63+
(:HealpixIEBMap, ("Ix"=>1, "Ex"=>2, "Bx"=>3, "I"=>1, "E"=>2, "B"=>3, "P"=>2:3, "IP"=>:)),
64+
(:HealpixIEBFourier, ("Il"=>1, "El"=>2, "Bl"=>3, "I"=>1, "E"=>2, "B"=>3, "P"=>2:3, "IP"=>:)),
65+
]
66+
# sub-components
67+
for (F, props) in _healpix_sub_components
68+
for (k,I) in props
69+
body = if k[end] in "xl"
70+
I==(:) ? :(getfield(f,:arr)) : :(view(getfield(f,:arr), :, $I))
71+
else
72+
I==(:) ? :f : :($(HealpixField{k=="P" ? Basis2Prod{basis(@eval($F)).parameters[end-1:end]...} : basis(@eval($F)).parameters[end]})(view(getfield(f,:arr), :, $I), f.metadata))
73+
end
74+
@eval getproperty(f::$F, ::Val{$(QuoteNode(Symbol(k)))}) = $body
75+
end
76+
end
77+
getproperty(f::HealpixField, ::Val{:proj}) = getfield(f,:metadata)
78+
79+
80+
81+
### broadcasting
82+
83+
function promote_metadata_strict(metadata₁::ProjHealpix, metadata₂::ProjHealpix)
84+
if metadata₁ === metadata₂
85+
metadata₁
86+
else
87+
error("Can't broadcast Healpix maps with two different Nsides ($metadata₁.Nside, $metadata₂.Nside).")
88+
end
89+
end
90+
91+
92+
### Projection
93+
94+
## Coordinate transformations
95+
96+
# choices of negative signs are such that the default projection is
97+
# centered on the center of a healpy.mollview plot and oriented the
98+
# same way.
99+
100+
function xy_to_θϕ(::ProjLambert, x, y)
12101
r = sqrt(x^2+y^2)
13102
θ = 2*acos(r/2)
14-
ϕ = -atan(y,x)
103+
ϕ = atan(-x,-y)
15104
θ, ϕ
16105
end
17106

18-
function healpix_to_flat(healpix_map::Vector{T}, proj::ProjLambert{T}; rots=[((0,90,0),)]) where {T}
19-
20-
Nside_sphere = hp.npix2nside(length(healpix_map))
21-
@unpack Δx, Ny, Nx = proj
22-
23-
# compute projection coordinate mapping
24-
ys = @. Δx * ((-Nx÷2:Nx÷2-1) + 0.5) # x/y switch here intentional
25-
xs = @. Δx * ((-Ny÷2:Ny÷2-1) + 0.5)
26-
xys = tuple.(xs,ys')[:]
27-
θϕs = xy_to_θϕ.(xys)
28-
(θs, ϕs) = first.(θϕs), last.(θϕs)
29-
30-
# the default rots=[(0,90,0)] makes it so you get a view of the
31-
# equator, to match Helapy's azeqview convention
32-
R = prod([hp.Rotator(rot..., eulertype="ZYX") for rot in rots])
33-
(θs′, ϕs′) = eachrow(R.get_inverse()(θs, ϕs))
34-
35-
# interpolate map
36-
FlatMap(reshape(hp.get_interp_val(healpix_map, θs′, ϕs′), Ny, Nx), proj)
37-
107+
function θϕ_to_xy(::ProjLambert, θ, ϕ)
108+
r = 2cos(θ/2)
109+
x = -r*sin(ϕ)
110+
y = -r*cos(ϕ)
111+
x, y
38112
end
39113

40-
function healpix_pixel_centers_to_flat(f::FlatField, Nside_sphere; rots=[((0,90,0),)], healpix_pixels=0:(12*Nside_sphere^2-1))
41-
42-
@unpack Δx, Ny, Nx = f
114+
# adding more flat projections in the future should be as easy as
115+
# defining pairs of functions like the above two for other cases
43116

44-
(θs, ϕs) = hp.pix2ang(Nside_sphere, healpix_pixels)
45-
46-
# the default rots=[(0,90,0)] makes it so you get a view of the
47-
# equator, to match Helapy's azeqview convention
48-
R = prod([hp.Rotator(rot..., eulertype="ZYX") for rot in rots])
49-
(θs′, ϕs′) = eachrow(R(θs, ϕs))
50117

51-
# compute projection coordinate mapping
52-
# (using Ny for xs and vice-versa intentional)
53-
xys = @. θϕ_to_xy(tuple(θs′, ϕs′))
54-
xs = @. first(xys) / Δx + Ny÷2 + 0.5 # x/y switch here intentional
55-
ys = @. last(xys) / Δx + Nx÷2 + 0.5
118+
function xy_to_θϕ(flat_proj::FlatProj; rotator)
119+
@unpack Δx, Ny, Nx = flat_proj
120+
xs = @. Δx * ((-Nx÷2:Nx÷2-1) + 0.5)
121+
ys = @. Δx * ((-Ny÷2:Ny÷2-1) + 0.5)
122+
θϕs = xy_to_θϕ.(Ref(flat_proj), xs, ys')
123+
(eachrow(rotator.get_inverse()(first.(θϕs)[:], last.(θϕs)[:]))...,)
124+
end
56125

126+
function θϕ_to_xy((hpx_proj, flat_proj)::Pair{<:ProjHealpix, <:FlatProj}; rotator)
127+
@unpack Δx, Ny, Nx = flat_proj
128+
@unpack Nside = hpx_proj
129+
(θs, ϕs) = hp.pix2ang(Nside, 0:(12*Nside^2-1))
130+
(θs′, ϕs′) = eachrow(rotator(θs, ϕs))
131+
xys = θϕ_to_xy.(Ref(flat_proj), θs′, ϕs′)
132+
xs = @. first(xys) / Δx + Nx÷2 + 0.5
133+
ys = @. last(xys) / Δx + Ny÷2 + 0.5
57134
(xs, ys)
135+
end
136+
137+
138+
## Healpix => Flat
139+
140+
"""
141+
project(healpix_map::HealpixField => flat_proj::FlatProj; [rotator])
58142
143+
Project `healpix_map` to a flat projection specified by `flat_proj`.
144+
E.g.:
145+
146+
healpix_map = HealpixMap(rand(12*2048^2))
147+
flat_proj = ProjLambert(Ny=128, Nx=128, θpix=3, T=Float32)
148+
f = project(healpix_map => flat_proj; rotator=pyimport("healpy").Rotator((0,28,23)))
149+
150+
The use of `=>` is to help remember in which order the arguments are
151+
specified. Optional keyword argument `rotator` is a `healpy.Rotator`
152+
object specifying a rotation which rotates the north pole to the
153+
center of the desired field. If `healpix_map` is a `HealpixQU`, Q/U
154+
polarization angles are rotated to be aligned with the local
155+
coordinates (sometimes called "polarization flattening").
156+
"""
157+
function project((hpx_map, flat_proj)::Pair{<:HealpixMap,<:ProjLambert{T}}; rotator=hp.Rotator((0,90,0))) where {T}
158+
_project(xy_to_θϕ(flat_proj; rotator), hpx_map => flat_proj)
159+
end
160+
161+
function project((hpx_map, flat_proj)::Pair{<:HealpixQUMap,<:ProjLambert{T}}; rotator=hp.Rotator((0,90,0))) where {T}
162+
@unpack Ny, Nx = flat_proj
163+
(θs, ϕs) = xy_to_θϕ(flat_proj; rotator)
164+
Q = _project((θs, ϕs), hpx_map.Q => flat_proj)
165+
U = _project((θs, ϕs), hpx_map.U => flat_proj)
166+
ψpol = permutedims(reshape((hp.Rotator((0,-90,0)) * rotator).angle_ref(θs, ϕs), Nx, Ny))
167+
QU_flat = @. (Q.arr + im * U.arr) * exp(im * 2 * ψpol)
168+
FlatQUMap(real.(QU_flat), imag.(QU_flat), flat_proj)
169+
end
170+
171+
function _project((θs,ϕs)::Tuple, (hpx_map, flat_proj)::Pair{<:HealpixMap,<:ProjLambert{T}}) where {T}
172+
@unpack Ny, Nx = flat_proj
173+
FlatMap(T.(permutedims(reshape(hp.get_interp_val(collect(hpx_map), θs, ϕs), Nx, Ny))), flat_proj)
59174
end
60175

61176

62-
function flat_to_healpix(f::FlatS0, Nside_sphere; kwargs...)
63-
# get pixel centers of Healpix pixels in 2D map
64-
xs,ys = healpix_pixel_centers_to_flat(f, Nside_sphere; kwargs...)
65-
# interpolate map
66-
@ondemand(Images.bilinear_interpolation).(Ref(f[:Ix]), xs, ys)
177+
## Flat => Healpix
178+
179+
"""
180+
project(flat_map::FlatField => healpix_proj::ProjHealpix; [rotator])
181+
182+
Reproject a `flat_map` back onto the sphere in a Healpix projection
183+
specified by `healpix_proj`. E.g.
184+
185+
flat_map = FlatMap(rand(128,128))
186+
f = project(flat_map => ProjHealpix(2048); rotator=pyimport("healpy").Rotator((0,28,23)))
187+
188+
The use of `=>` is to help remember in which order the arguments are
189+
specified. Optional keyword argument `rotator` is a `healpy.Rotator`
190+
object specifying a rotation which rotates the north pole to the
191+
center of the desired field.
192+
"""
193+
function project((flat_map, hpx_proj)::Pair{<:FlatS0, <:ProjHealpix}; rotator=hp.Rotator((0,90,0)))
194+
(xs, ys) = θϕ_to_xy(hpx_proj => flat_map.proj; rotator)
195+
HealpixMap(@ondemand(Images.bilinear_interpolation).(Ref(cpu(flat_map[:Ix])), ys, xs), hpx_proj)
67196
end
68197

‎src/plotting.jl‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -247,7 +247,13 @@ function animate(fields::AbstractVecOrMat{<:AbstractVecOrMat{<:FlatField}}; fps=
247247
end
248248

249249

250-
### Plotting Loess interpolated objects
250+
251+
### plotting HealpixFields
252+
253+
plot(f::HealpixMap; kwargs...) = hp.mollview(collect(f.arr); cmap="RdBu_r", kwargs...)
254+
255+
256+
251257

252258

253259
### convenience

0 commit comments

Comments
 (0)