Skip to content

Commit b8d9fa1

Browse files
marius311marius311
committed
improve plotting
1 parent 00ed4a5 commit b8d9fa1

File tree

1 file changed

+108
-66
lines changed

1 file changed

+108
-66
lines changed

‎src/plotting.jl‎

Lines changed: 108 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -22,98 +22,136 @@ end
2222

2323
plotsize₀ = 4
2424

25-
pretty_name(s::Symbol) = pretty_name(Val.(Symbol.(split(string(s),"")))...)
26-
pretty_name(::Val{s},::Val{:x}) where {s} = "$s Map"
27-
pretty_name(::Val{s},::Val{:l}) where {s} = "$s Fourier"
28-
pretty_name(::Val{:T},::Val{:x}) where {s} = "Map"
29-
pretty_name(::Val{:T},::Val{:l}) where {s} = "Fourier"
30-
31-
# generic plotting some components of a FlatField
32-
function _plot(f::FlatField{P}, ax, k, title, vlim; units=:deg, ticklabels=true, axeslabels=false, kwargs...) where {N,θ,P<:Flat{N,θ}}
33-
if string(k)[2] == 'x'
34-
x = θ*N/Dict(:deg=>60,:arcmin=>1)[units]/2
35-
elseif string(k)[2] == 'l'
36-
x = fieldinfo(f).nyq
37-
else
38-
throw(ArgumentError("Invalid `which`: $k"))
39-
end
40-
extent = [-x,x,-x,x]
41-
(title == nothing) && (title="$(pretty_name(k)) ($(N)x$(N) @ $(θ)')")
42-
(vlim == nothing) && (vlim=:sym)
43-
_plot(Array(f[k]); ax=ax, extent=extent, title=title, vlim=vlim, kwargs...)
44-
if ticklabels
45-
if string(k)[2] == 'x'
46-
@pydef mutable struct MyFmt <: pyimport(:matplotlib).ticker.ScalarFormatter
47-
__call__(self,v,p=nothing) = py"super"(MyFmt,self).__call__(v,p)*Dict(:deg=>"°",:arcmin=>"′")[units]
48-
end
49-
ax.xaxis.set_major_formatter(MyFmt())
50-
ax.yaxis.set_major_formatter(MyFmt())
51-
if axeslabels
52-
ax.set_xlabel("RA")
53-
ax.set_ylabel("Dec")
54-
end
55-
else
56-
ax.set_xlabel(raw"$\ell_x$")
57-
ax.set_ylabel(raw"$\ell_y$")
58-
ax.tick_params(axis="x", rotation=45)
59-
end
60-
ax.tick_params(labeltop=false, labelbottom=true)
61-
else
62-
ax.tick_params(labeltop=false, labelleft=false)
63-
end
64-
end
25+
pretty_name(s) = pretty_name(Val.(Symbol.(split(string(s),"")))...)
26+
pretty_name(::Val{s}, b::Val) where {s} = "$s "*pretty_name(b)
27+
pretty_name(::Val{:x}) = "Map"
28+
pretty_name(::Val{:l}) = "Fourier"
6529

66-
# plotting a map
67-
function _plot(m::AbstractMatrix{<:Real}; ax=gca(), title=nothing, vlim=:sym, cmap="RdBu_r", vscale=:linear, cbar=true, kwargs...)
30+
function _plot(f, ax, k, title, vlim, vscale, cmap; cbar=true, units=:deg, ticklabels=true, axeslabels=false, kwargs...)
6831

69-
# some logic to automatically get upper/lower limits
70-
if vlim==:sym
71-
vmax = quantile(abs.(m[@. !isnan(m)][:]),0.999)
32+
@unpack Nside, θpix = fieldinfo(f)
33+
ismap = endswith(string(k), "x")
34+
35+
# default values
36+
if title == nothing
37+
if f isa FlatS0
38+
title = pretty_name(string(k)[2])
39+
else
40+
title = pretty_name(k)
41+
end
42+
title *= " ($(Nside)x$(Nside) @ $(θpix)')"
43+
end
44+
if vlim == nothing
45+
vlim = ismap ? :sym : :asym
46+
end
47+
if vscale == nothing
48+
vscale = ismap ? :linear : :log
49+
end
50+
if cmap == nothing
51+
if ismap
52+
cmap = get_cmap("RdBu_r")
53+
else
54+
cmap = get_cmap("viridis")
55+
cmap.set_bad("lightgray")
56+
end
57+
end
58+
59+
# build array
60+
if ismap
61+
arr = Array(f[k])
62+
else
63+
arr = abs.(ifftshift(unfold(Array(f[k]))))
64+
end
65+
if vscale == :log
66+
arr[arr .== 0] .= NaN
67+
end
68+
69+
# auto vlim's
70+
if vlim==:sym
71+
vmax = quantile(abs.(arr[@. !isnan(arr)][:]),0.999)
7272
vmin = -vmax
7373
elseif vlim==:asym
74-
vmin, vmax = (quantile(m[@. !isnan(m)][:],q) for q=(0.001,0.999))
74+
vmin, vmax = (quantile(arr[@. !isnan(arr)][:],q) for q=(0.001,0.999))
7575
elseif isa(vlim,Tuple)
7676
vmin, vmax = vlim
7777
else
7878
vmax = vlim
7979
vmin = -vmax
8080
end
81-
82-
m = Float64.(m)
83-
m[isinf.(m)] .= NaN
84-
85-
cax = ax.matshow(clamp.(m,vmin,vmax); vmin=vmin, vmax=vmax, cmap=cmap, rasterized=true, kwargs...)
86-
cbar && gcf().colorbar(cax,ax=ax)
87-
title!=nothing && ax.set_title(title, y=1)
88-
ax
89-
end
90-
91-
# plotting fourier coefficients
92-
function _plot(m::AbstractMatrix{<:Complex}; vscale=:log, kwargs...)
93-
dat = ifftshift(unfold(m))
94-
if vscale==:log
95-
dat .= log10.(abs.(dat))
81+
82+
# make the plot
83+
if ismap
84+
extent = [-1,1,-1,1] .* θpix*Nside/Dict(:deg=>60,:arcmin=>1)[units]/2
85+
else
86+
extent = [-1,1,-1,1] .* fieldinfo(f).nyq
87+
end
88+
norm = vscale == :log ? matplotlib.colors.LogNorm() : nothing
89+
cax = ax.matshow(
90+
arr;
91+
vmin=vmin, vmax=vmax, extent=extent,
92+
cmap=cmap, rasterized=true, norm=norm,
93+
kwargs...
94+
)
95+
96+
# annonate
97+
if cbar
98+
colorbar(cax,ax=ax)
99+
end
100+
ax.set_title(title, y=1)
101+
if ticklabels
102+
if ismap
103+
@pydef mutable struct MyFmt <: pyimport(:matplotlib).ticker.ScalarFormatter
104+
__call__(self,v,p=nothing) = py"super"(MyFmt,self).__call__(v,p)*Dict(:deg=>"°",:arcmin=>"′")[units]
105+
end
106+
ax.xaxis.set_major_formatter(MyFmt())
107+
ax.yaxis.set_major_formatter(MyFmt())
108+
if axeslabels
109+
ax.set_xlabel("RA")
110+
ax.set_ylabel("Dec")
111+
end
112+
else
113+
ax.set_xlabel(raw"$\ell_x$")
114+
ax.set_ylabel(raw"$\ell_y$")
115+
ax.tick_params(axis="x", rotation=45)
116+
end
117+
ax.tick_params(labeltop=false, labelbottom=true)
118+
else
119+
ax.tick_params(labeltop=false, labelleft=false)
96120
end
97-
_plot(real.(dat); vlim=(nothing,nothing), cmap=nothing, kwargs...)
121+
98122
end
99123

100124

101-
102125
@doc doc"""
103126
plot(f::Field; kwargs...)
104127
plot(fs::VecOrMat{\<:Field}; kwarg...)
105128
106129
Plotting fields.
107130
"""
108131
plot(f::Field; kwargs...) = plot([f]; kwargs...)
109-
function plot(fs::AbstractVecOrMat{F}; plotsize=plotsize₀, which=default_which(fs), title=nothing, vlim=nothing, return_all=false, kwargs...) where {F<:Field}
132+
plot(D::DiagOp; kwargs...) =
133+
plot([diag(D)]; which=permutedims([x for x in propertynames(diag(D)) if string(x)[end] in "xl"]), kwargs...)
134+
135+
function plot(
136+
fs::AbstractVecOrMat{F};
137+
plotsize = plotsize₀,
138+
which = default_which(fs),
139+
title = nothing,
140+
vlim = nothing,
141+
vscale = nothing,
142+
cmap = nothing,
143+
return_all = false,
144+
kwargs...) where {F<:Field}
145+
110146
(m,n) = size(tuple.(fs, which)[:,:])
111147
fig,axs = subplots(m, n; figsize=plotsize.*[1.4*n,m], squeeze=false)
112148
axs = getindex.(Ref(axs), 1:m, (1:n)') # see https://github.com/JuliaPy/PyCall.jl/pull/487#issuecomment-456998345
113-
_plot.(fs,axs,which,title,vlim; kwargs...)
149+
_plot.(fs,axs,which,title,vlim,vscale,cmap; kwargs...)
114150
tight_layout(w_pad=-10)
115151
return_all ? (fig,axs,which) : isjuno ? fig : nothing
152+
116153
end
154+
117155
default_which(::AbstractVecOrMat{<:FlatS0}) = [:Ix]
118156
default_which(::AbstractVecOrMat{<:FlatS2}) = [:Ex :Bx]
119157
default_which(::AbstractVecOrMat{<:FlatS02}) = [:Ix :Ex :Bx]
@@ -126,6 +164,8 @@ function default_which(fs::AbstractVecOrMat{<:Field})
126164
end
127165

128166

167+
### animations of FlatFields
168+
129169
@doc doc"""
130170
animate(fields::Vector{\<:Vector{\<:Field}}; interval=50, motionblur=false, kwargs...)
131171
@@ -168,7 +208,9 @@ for plot in (:plot, :loglog, :semilogx, :semilogy)
168208

169209
@eval function ($plot)(f::Function, m::Loess.LoessModel, args...; kwargs...)
170210
l, = ($plot)(m.xs, f.(m.ys), ".", args...; kwargs...)
171-
xs′ = range(first(m.xs),last(m.xs),length=10*length(m.xs))
211+
xs′ = vcat(map(1:length(m.xs)-1) do i
212+
collect(range(m.xs[i],m.xs[i+1],length=10))[1:end-1]
213+
end..., [last(m.xs)])
172214
($plot)(xs′, f.(m.(xs′)), args...; c=l.get_color(), kwargs...)
173215
end
174216

0 commit comments

Comments
 (0)