Skip to content

Commit cc0da24

Browse files
haakon-ehaakon-e
committed
migrate plot_timeheight to Makie
1 parent d7c087d commit cc0da24

File tree

1 file changed

+85
-109
lines changed

1 file changed

+85
-109
lines changed

‎test/plotting_utils.jl‎

Lines changed: 85 additions & 109 deletions
Original file line numberDiff line numberDiff line change
@@ -352,75 +352,82 @@ function plot_timeheight_p3(nc_data_file, precip; output = "output")
352352
Plots.png(p, joinpath(path, "timeheight.png"))
353353
end
354354

355+
function make_tz_hm(i, j, t, z, var; imax = 3, title = "", kw...)
356+
fig = current_figure()
357+
ax = Axis(fig[i, j];
358+
xlabel = i == imax ? "time [s]" : "",
359+
ylabel = j == 1 ? "z [m]" : "",
360+
yticklabelsvisible = j == 1,
361+
xticklabelsvisible = i == 3,
362+
title,
363+
)
364+
hm = CairoMakie.heatmap!(ax, t, z, var; colormap = :BuPu, kw...)
365+
Makie.Colorbar(fig[i, j], hm;
366+
halign = 1, valign = 0,
367+
height = Makie.Relative(0.3), width = 10, tellwidth = false,
368+
ticks = WilkinsonTicks(3; k_min = 3, k_max = 6),
369+
)
370+
end
371+
355372
function plot_timeheight(nc_data_file; output = "output", mixed_phase = true, pysdm = false)
356373
path = joinpath(@__DIR__, output)
357374
mkpath(path)
358375

359376
ds = NC.NCDataset(joinpath(@__DIR__, nc_data_file))
360377
if pysdm
361-
t_plt = Array(ds["time"])
362-
z_plt = Array(ds["height"])
363-
q_liq_plt = transpose(Array(ds["cloud water mixing ratio"]))
364-
q_rai_plt = transpose(Array(ds["rain water mixing ratio"]))
365-
q_ice_plt = transpose(Array(ds["rain water mixing ratio"])) * FT(0)
366-
q_sno_plt = transpose(Array(ds["rain water mixing ratio"])) * FT(0)
367-
q_vap = transpose(Array(ds["water_vapour_mixing_ratio"])) * 1e3
368-
q_tot_plt = q_vap + q_liq_plt
369-
N_aer_plt = transpose(Array(ds["na"]))
370-
N_liq_plt = transpose(Array(ds["nc"]))
371-
N_rai_plt = transpose(Array(ds["nr"]))
378+
t = Array(ds["time"])
379+
z = Array(ds["height"])
380+
q_liq = Array(ds["cloud water mixing ratio"])
381+
q_rai = Array(ds["rain water mixing ratio"])
382+
q_ice = Array(ds["rain water mixing ratio"]) * FT(0)
383+
q_sno = Array(ds["rain water mixing ratio"]) * FT(0)
384+
q_vap = Array(ds["water_vapour_mixing_ratio"]) * 1e3
385+
q_tot = q_vap + q_liq
386+
N_aer = Array(ds["na"])
387+
N_liq = Array(ds["nc"])
388+
N_rai = Array(ds["nr"])
372389
else
373-
t_plt = Array(ds.group["profiles"]["t"])
374-
z_plt = Array(ds.group["profiles"]["zc"])
375-
q_tot_plt = Array(ds.group["profiles"]["q_tot"])
376-
q_liq_plt = Array(ds.group["profiles"]["q_liq"])
377-
q_ice_plt = Array(ds.group["profiles"]["q_ice"])
378-
q_rai_plt = Array(ds.group["profiles"]["q_rai"])
379-
q_sno_plt = Array(ds.group["profiles"]["q_sno"])
380-
N_aer_plt = Array(ds.group["profiles"]["N_aer"])
381-
N_liq_plt = Array(ds.group["profiles"]["N_liq"])
382-
N_rai_plt = Array(ds.group["profiles"]["N_rai"])
390+
t = Array(ds.group["profiles"]["t"])
391+
z = Array(ds.group["profiles"]["zc"])
392+
q_tot = permutedims(Array(ds.group["profiles"]["q_tot"]))
393+
q_liq = permutedims(Array(ds.group["profiles"]["q_liq"]))
394+
q_ice = permutedims(Array(ds.group["profiles"]["q_ice"]))
395+
q_rai = permutedims(Array(ds.group["profiles"]["q_rai"]))
396+
q_sno = permutedims(Array(ds.group["profiles"]["q_sno"]))
397+
N_aer = permutedims(Array(ds.group["profiles"]["N_aer"]))
398+
N_liq = permutedims(Array(ds.group["profiles"]["N_liq"]))
399+
N_rai = permutedims(Array(ds.group["profiles"]["N_rai"]))
383400
end
384-
#! format: off
385-
p1 = Plots.heatmap(t_plt, z_plt, q_tot_plt .* 1e3, title = "q_tot [g/kg]", xlabel = "time [s]", ylabel = "z [m]", color = :BuPu, clims=(8, 15))
386-
p2 = Plots.heatmap(t_plt, z_plt, q_liq_plt .* 1e3, title = "q_liq [g/kg]", xlabel = "time [s]", ylabel = "z [m]", color = :BuPu, clims=(0, 1))
387-
p3 = Plots.heatmap(t_plt, z_plt, q_ice_plt .* 1e3, title = "q_ice [g/kg]", xlabel = "time [s]", ylabel = "z [m]", color = :BuPu, clims=(0, 0.25))
388-
p4 = Plots.heatmap(t_plt, z_plt, q_rai_plt .* 1e3, title = "q_rai [g/kg]", xlabel = "time [s]", ylabel = "z [m]", color = :BuPu, clims=(0, 0.25))
389-
p5 = Plots.heatmap(t_plt, z_plt, q_sno_plt .* 1e3, title = "q_sno [g/kg]", xlabel = "time [s]", ylabel = "z [m]", color = :BuPu)
390-
p6 = Plots.heatmap(t_plt, z_plt, N_aer_plt .* 1e-6, title = "N_aer [1/cm^3]", xlabel = "time [s]", ylabel = "z [m]", color = :BuPu, clims=(0, 100))
391-
p7 = Plots.heatmap(t_plt, z_plt, N_liq_plt .* 1e-6, title = "N_liq [1/cm^3]", xlabel = "time [s]", ylabel = "z [m]", color = :BuPu, clims=(0, 50))
392-
p8 = Plots.heatmap(t_plt, z_plt, N_rai_plt .* 1e-6, title = "N_rai [1/cm^3]", xlabel = "time [s]", ylabel = "z [m]", color = :BuPu, clims=(0, 1))
393-
#! format: on
401+
402+
kg_to_g = 1e3
403+
m⁻³_to_cm⁻³ = 1e-6
404+
394405
if mixed_phase
395-
p = Plots.plot(
396-
p1,
397-
p2,
398-
p3,
399-
p4,
400-
p5,
401-
p6,
402-
p7,
403-
p8,
404-
size = (1200.0, 1200.0),
405-
bottom_margin = 30.0 * Plots.PlotMeasures.px,
406-
left_margin = 30.0 * Plots.PlotMeasures.px,
407-
layout = (3, 3),
408-
)
406+
fig = Figure(size = (1200, 1200))
407+
make_tz_hm(1, 1, t, z, q_tot .* kg_to_g; colorrange = (8, 15), title = "q_tot [g/kg]")
408+
make_tz_hm(1, 2, t, z, q_liq .* kg_to_g; colorrange = (0, 1), title = "q_liq [g/kg]")
409+
make_tz_hm(1, 3, t, z, q_ice .* kg_to_g; colorrange = (0, 0.25), title = "q_ice [g/kg]")
410+
make_tz_hm(2, 1, t, z, q_rai .* kg_to_g; colorrange = (0, 0.25), title = "q_rai [g/kg]")
411+
make_tz_hm(2, 2, t, z, q_sno .* kg_to_g; title = "q_sno [g/kg]")
412+
make_tz_hm(2, 3, t, z, N_aer .* m⁻³_to_cm⁻³; colorrange = (0, 100), title = "N_aer [1/cm³]")
413+
make_tz_hm(3, 1, t, z, N_liq .* m⁻³_to_cm⁻³; colorrange = (0, 50), title = "N_liq [1/cm³]")
414+
make_tz_hm(3, 2, t, z, N_rai .* m⁻³_to_cm⁻³; colorrange = (0, 1), title = "N_rai [1/cm³]")
415+
409416
else
410-
p = Plots.plot(
411-
p1,
412-
p2,
413-
p4,
414-
p6,
415-
p7,
416-
p8,
417-
size = (1200.0, 600.0),
418-
bottom_margin = 30.0 * Plots.PlotMeasures.px,
419-
left_margin = 30.0 * Plots.PlotMeasures.px,
420-
layout = (2, 3),
421-
)
417+
fig = Figure(size = (1200, 600))
418+
make_tz_hm(1, 1, t, z, q_tot .* kg_to_g; colorrange = (8, 15), title = "q_tot [g/kg]")
419+
make_tz_hm(1, 2, t, z, q_liq .* kg_to_g; colorrange = (0, 1), title = "q_liq [g/kg]")
420+
make_tz_hm(1, 3, t, z, q_rai .* kg_to_g; colorrange = (0, 0.25), title = "q_rai [g/kg]")
421+
make_tz_hm(2, 1, t, z, N_aer .* m⁻³_to_cm⁻³; colorrange = (0, 100), title = "N_aer [1/cm³]")
422+
make_tz_hm(2, 2, t, z, N_liq .* m⁻³_to_cm⁻³; colorrange = (0, 50), title = "N_liq [1/cm³]")
423+
make_tz_hm(2, 3, t, z, N_rai .* m⁻³_to_cm⁻³; colorrange = (0, 1), title = "N_rai [1/cm³]")
422424
end
423-
Plots.png(p, joinpath(path, "timeheight.png"))
425+
fig
426+
427+
axs = filter(x -> x isa Axis, contents(fig[:, :]))
428+
linkaxes!(axs...)
429+
save(joinpath(path, "timeheight.png"), fig)
430+
nothing
424431
end
425432

426433
function plot_timeheight_no_ice_snow(nc_data_file; output = "output", pysdm = false)
@@ -429,58 +436,27 @@ function plot_timeheight_no_ice_snow(nc_data_file; output = "output", pysdm = fa
429436

430437
ds = NC.NCDataset(joinpath(@__DIR__, nc_data_file))
431438
if pysdm
432-
t_plt = Array(ds["time"])
433-
z_plt = Array(ds["height"])
434-
q_liq_plt = transpose(Array(ds["cloud water mixing ratio"])) / 1e3
435-
q_rai_plt = transpose(Array(ds["rain water mixing ratio"])) / 1e3
436-
q_vap = transpose(Array(ds["water_vapour_mixing_ratio"]))
437-
q_tot_plt = q_vap + q_liq_plt
439+
t = Array(ds["time"])
440+
z = Array(ds["height"])
441+
q_liq = Array(ds["cloud water mixing ratio"]) / 1e3
442+
q_rai = Array(ds["rain water mixing ratio"]) / 1e3
443+
q_vap = Array(ds["water_vapour_mixing_ratio"])
444+
q_tot = q_vap + q_liq
438445
else
439-
t_plt = Array(ds.group["profiles"]["t"])
440-
z_plt = Array(ds.group["profiles"]["zc"])
441-
q_tot_plt = Array(ds.group["profiles"]["q_tot"])
442-
q_liq_plt = Array(ds.group["profiles"]["q_liq"])
443-
q_rai_plt = Array(ds.group["profiles"]["q_rai"])
446+
t = Array(ds.group["profiles"]["t"])
447+
z = Array(ds.group["profiles"]["zc"])
448+
q_tot = permutedims(Array(ds.group["profiles"]["q_tot"]))
449+
q_liq = permutedims(Array(ds.group["profiles"]["q_liq"]))
450+
q_rai = permutedims(Array(ds.group["profiles"]["q_rai"]))
444451
end
445452

446-
p1 = Plots.heatmap(
447-
t_plt,
448-
z_plt,
449-
q_tot_plt .* 1e3,
450-
title = "q_tot [g/kg]",
451-
xlabel = "time [s]",
452-
ylabel = "z [m]",
453-
color = :BuPu,
454-
clims = (0, 1),
455-
)
456-
p2 = Plots.heatmap(
457-
t_plt,
458-
z_plt,
459-
q_liq_plt .* 1e3,
460-
title = "q_liq [g/kg]",
461-
xlabel = "time [s]",
462-
ylabel = "z [m]",
463-
color = :BuPu,
464-
clims = (0, 1),
465-
)
466-
p3 = Plots.heatmap(
467-
t_plt,
468-
z_plt,
469-
q_rai_plt .* 1e3,
470-
title = "q_rai [g/kg]",
471-
xlabel = "time [s]",
472-
ylabel = "z [m]",
473-
color = :BuPu,
474-
clims = (0, 0.25),
475-
)
476-
p = Plots.plot(
477-
p1,
478-
p2,
479-
p3,
480-
size = (1200.0, 300.0),
481-
bottom_margin = 30.0 * Plots.PlotMeasures.px,
482-
left_margin = 30.0 * Plots.PlotMeasures.px,
483-
layout = (1, 3),
484-
)
485-
Plots.png(p, joinpath(path, "timeheight.png"))
453+
kg_to_g = 1e3
454+
fig = Figure(size = (1200, 300))
455+
make_tz_hm(1, 1, t, z, q_tot .* kg_to_g; title = "q_tot [g/kg]", imax = 1, colorrange = (0, 1))
456+
make_tz_hm(1, 2, t, z, q_liq .* kg_to_g; title = "q_liq [g/kg]", imax = 1, colorrange = (0, 1))
457+
make_tz_hm(1, 3, t, z, q_rai .* kg_to_g; title = "q_rai [g/kg]", imax = 1, colorrange = (0, 0.25))
458+
axs = filter(x -> x isa Axis, contents(fig[:, :]))
459+
linkaxes!(axs...)
460+
save(joinpath(path, "timeheight.png"), fig)
461+
nothing
486462
end

0 commit comments

Comments
 (0)