Use PyPlot to work around Plots not supporting 2d heatmap coordinates

Making animations directly with PyPlot.jl and matplotlib is not as neat,
and the output is not consistent with Plots.jl, but it works without
needing the update to Plots.jl in
JuliaPlots/Plots.jl#4298

Manifest.toml

@@ -1,6 +1,6 @@
 # This file is machine-generated - editing it directly is not advised
 
-julia_version = "1.7.0"
+julia_version = "1.7.2"
 manifest_format = "2.0"
 
 [[deps.AbstractFFTs]]
@@ -50,9 +50,9 @@ version = "0.1.1"
 
 [[deps.Cairo_jll]]
 deps = ["Artifacts", "Bzip2_jll", "Fontconfig_jll", "FreeType2_jll", "Glib_jll", "JLLWrappers", "LZO_jll", "Libdl", "Pixman_jll", "Pkg", "Xorg_libXext_jll", "Xorg_libXrender_jll", "Zlib_jll", "libpng_jll"]
-git-tree-sha1 = "f2202b55d816427cd385a9a4f3ffb226bee80f99"
+git-tree-sha1 = "4b859a208b2397a7a623a03449e4636bdb17bcf2"
 uuid = "83423d85-b0ee-5818-9007-b63ccbeb887a"
-version = "1.16.1+0"
+version = "1.16.1+1"
 
 [[deps.ChainRulesCore]]
 deps = ["Compat", "LinearAlgebra", "SparseArrays"]
@@ -324,9 +324,9 @@ version = "0.21.0+0"
 
 [[deps.Glib_jll]]
 deps = ["Artifacts", "Gettext_jll", "JLLWrappers", "Libdl", "Libffi_jll", "Libiconv_jll", "Libmount_jll", "PCRE_jll", "Pkg", "Zlib_jll"]
-git-tree-sha1 = "7bf67e9a481712b3dbe9cb3dac852dc4b1162e02"
+git-tree-sha1 = "a32d672ac2c967f3deb8a81d828afc739c838a06"
 uuid = "7746bdde-850d-59dc-9ae8-88ece973131d"
-version = "2.68.3+0"
+version = "2.68.3+2"
 
 [[deps.Glob]]
 git-tree-sha1 = "4df9f7e06108728ebf00a0a11edee4b29a482bb2"
@@ -358,9 +358,9 @@ version = "0.9.17"
 
 [[deps.HarfBuzz_jll]]
 deps = ["Artifacts", "Cairo_jll", "Fontconfig_jll", "FreeType2_jll", "Glib_jll", "Graphite2_jll", "JLLWrappers", "Libdl", "Libffi_jll", "Pkg"]
-git-tree-sha1 = "8a954fed8ac097d5be04921d595f741115c1b2ad"
+git-tree-sha1 = "129acf094d168394e80ee1dc4bc06ec835e510a3"
 uuid = "2e76f6c2-a576-52d4-95c1-20adfe4de566"
-version = "2.8.1+0"
+version = "2.8.1+1"
 
 [[deps.IJulia]]
 deps = ["Base64", "Conda", "Dates", "InteractiveUtils", "JSON", "Libdl", "Markdown", "MbedTLS", "Pkg", "Printf", "REPL", "Random", "SoftGlobalScope", "Test", "UUIDs", "ZMQ"]
@@ -440,6 +440,12 @@ git-tree-sha1 = "f6250b16881adf048549549fba48b1161acdac8c"
 uuid = "c1c5ebd0-6772-5130-a774-d5fcae4a789d"
 version = "3.100.1+0"
 
+[[deps.LERC_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
+git-tree-sha1 = "bf36f528eec6634efc60d7ec062008f171071434"
+uuid = "88015f11-f218-50d7-93a8-a6af411a945d"
+version = "3.0.0+1"
+
 [[deps.LZO_jll]]
 deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
 git-tree-sha1 = "e5b909bcf985c5e2605737d2ce278ed791b89be6"
@@ -517,10 +523,10 @@ uuid = "4b2f31a3-9ecc-558c-b454-b3730dcb73e9"
 version = "2.35.0+0"
 
 [[deps.Libtiff_jll]]
-deps = ["Artifacts", "JLLWrappers", "JpegTurbo_jll", "Libdl", "Pkg", "Zlib_jll", "Zstd_jll"]
-git-tree-sha1 = "340e257aada13f95f98ee352d316c3bed37c8ab9"
+deps = ["Artifacts", "JLLWrappers", "JpegTurbo_jll", "LERC_jll", "Libdl", "Pkg", "Zlib_jll", "Zstd_jll"]
+git-tree-sha1 = "c9551dd26e31ab17b86cbd00c2ede019c08758eb"
 uuid = "89763e89-9b03-5906-acba-b20f662cd828"
-version = "4.3.0+0"
+version = "4.3.0+1"
 
 [[deps.Libuuid_jll]]
 deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
@@ -647,9 +653,9 @@ uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908"
 
 [[deps.Ogg_jll]]
 deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
-git-tree-sha1 = "7937eda4681660b4d6aeeecc2f7e1c81c8ee4e2f"
+git-tree-sha1 = "887579a3eb005446d514ab7aeac5d1d027658b8f"
 uuid = "e7412a2a-1a6e-54c0-be00-318e2571c051"
-version = "1.3.5+0"
+version = "1.3.5+1"
 
 [[deps.OpenBLAS_jll]]
 deps = ["Artifacts", "CompilerSupportLibraries_jll", "Libdl"]
@@ -761,11 +767,23 @@ version = "0.5.0"
 deps = ["Unicode"]
 uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 
+[[deps.PyCall]]
+deps = ["Conda", "Dates", "Libdl", "LinearAlgebra", "MacroTools", "Serialization", "VersionParsing"]
+git-tree-sha1 = "53b8b07b721b77144a0fbbbc2675222ebf40a02d"
+uuid = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"
+version = "1.94.1"
+
+[[deps.PyPlot]]
+deps = ["Colors", "LaTeXStrings", "PyCall", "Sockets", "Test", "VersionParsing"]
+git-tree-sha1 = "f9d953684d4d21e947cb6d642db18853d43cb027"
+uuid = "d330b81b-6aea-500a-939a-2ce795aea3ee"
+version = "2.11.0"
+
 [[deps.Qt5Base_jll]]
 deps = ["Artifacts", "CompilerSupportLibraries_jll", "Fontconfig_jll", "Glib_jll", "JLLWrappers", "Libdl", "Libglvnd_jll", "OpenSSL_jll", "Pkg", "Xorg_libXext_jll", "Xorg_libxcb_jll", "Xorg_xcb_util_image_jll", "Xorg_xcb_util_keysyms_jll", "Xorg_xcb_util_renderutil_jll", "Xorg_xcb_util_wm_jll", "Zlib_jll", "xkbcommon_jll"]
-git-tree-sha1 = "ad368663a5e20dbb8d6dc2fddeefe4dae0781ae8"
+git-tree-sha1 = "c6c0f690d0cc7caddb74cef7aa847b824a16b256"
 uuid = "ea2cea3b-5b76-57ae-a6ef-0a8af62496e1"
-version = "5.15.3+0"
+version = "5.15.3+1"
 
 [[deps.QuadGK]]
 deps = ["DataStructures", "LinearAlgebra"]
@@ -1189,9 +1207,9 @@ version = "1.0.20+0"
 
 [[deps.libvorbis_jll]]
 deps = ["Artifacts", "JLLWrappers", "Libdl", "Ogg_jll", "Pkg"]
-git-tree-sha1 = "c45f4e40e7aafe9d086379e5578947ec8b95a8fb"
+git-tree-sha1 = "b910cb81ef3fe6e78bf6acee440bda86fd6ae00c"
 uuid = "f27f6e37-5d2b-51aa-960f-b287f2bc3b7a"
-version = "1.3.7+0"
+version = "1.3.7+1"
 
 [[deps.nghttp2_jll]]
 deps = ["Artifacts", "Libdl"]

Project.toml

@@ -20,6 +20,8 @@ OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 PackageCompiler = "9b87118b-4619-50d2-8e1e-99f35a4d4d9d"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 Primes = "27ebfcd6-29c5-5fa9-bf4b-fb8fc14df3ae"
+PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"
+PyPlot = "d330b81b-6aea-500a-939a-2ce795aea3ee"
 Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"
 Roots = "f2b01f46-fcfa-551c-844a-d8ac1e96c665"
 SHA = "ea8e919c-243c-51af-8825-aaa63cd721ce"

src/post_processing.jl

@@ -4,6 +4,10 @@ module post_processing
 
 export analyze_and_plot_data
 
+# Next three lines only used for workaround needed by plot_unnormalised()
+using PyCall
+import PyPlot
+
 # packages
 using Plots
 using IJulia
@@ -344,34 +348,84 @@ function plot_1D_1V_diagnostics(run_names, run_labels, nc_files, nwrite_movie,
  gif(anim, outfile, fps=5)
  end
  if pp.animate_f_unnormalized
- # make a gif animation of f_unnorm(v_parallel_unnorm,z,t)
- anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
- subplots = (@views plot_unnormalised(f[:,:,is,i], this_z, this_vpa,
- n[:,is,i], upar[:,is,i], vth[:,is,i], ev_n, ev_u,
- ev_p, title=run_label)
- for (f, n, upar, vth, ev_n, ev_u, ev_p, this_z,
- this_vpa, run_label) ∈
- zip(ff, density, parallel_flow, thermal_speed,
- evolve_density, evolve_upar, evolve_ppar, z, vpa,
- run_labels))
- plot(subplots..., layout=(1,n_runs), size=(400*n_runs, 400))
+ ## The nice, commented out version will only work when plot_unnormalised can
+ ## use Plots.jl...
+ ## make a gif animation of f_unnorm(v_parallel_unnorm,z,t)
+ #anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
+ # subplots = (@views plot_unnormalised(f[:,:,is,i], this_z, this_vpa,
+ # n[:,is,i], upar[:,is,i], vth[:,is,i], ev_n, ev_u,
+ # ev_p, title=run_label)
+ # for (f, n, upar, vth, ev_n, ev_u, ev_p, this_z,
+ # this_vpa, run_label) ∈
+ # zip(ff, density, parallel_flow, thermal_speed,
+ # evolve_density, evolve_upar, evolve_ppar, z, vpa,
+ # run_labels))
+ # plot(subplots..., layout=(1,n_runs), size=(400*n_runs, 400))
+ #end
+ #outfile = string(prefix, "_f_unnorm_vs_vpa_z", spec_string, ".gif")
+ #gif(anim, outfile, fps=5)
+ ## make a gif animation of log(f_unnorm)(v_parallel_unnorm,z,t)
+ #anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
+ # subplots = (@views plot_unnormalised(f[:,:,is,i], this_z, this_vpa, n[:,is,i],
+ # upar[:,is,i], vth[:,is,i], ev_n, ev_u, ev_p,
+ # plot_log=true, title=run_label)
+ # for (f, n, upar, vth, ev_n, ev_u, ev_p, this_z,
+ # this_vpa, run_label) ∈
+ # zip(ff, density, parallel_flow, thermal_speed,
+ # evolve_density, evolve_upar, evolve_ppar, z,
+ # vpa, run_labels))
+ # plot(subplots..., layout=(1,n_runs), size=(400*n_runs, 400))
+ #end
+ #outfile = string(prefix, "_logf_unnorm_vs_vpa_z", spec_string, ".gif")
+ #gif(anim, outfile, fps=5)
+
+ matplotlib = pyimport("matplotlib")
+ matplotlib.use("agg")
+ matplotlib_animation = pyimport("matplotlib.animation")
+ iframes = collect(itime_min:nwrite_movie:itime_max)
+ nframes = length(iframes)
+ function make_frame(i)
+ PyPlot.clf()
+ iframe = iframes[i+1]
+ # i counts from 0, Python-style
+ for (run_ind, f, n, upar, vth, ev_n, ev_u, ev_p, this_z, this_vpa,
+ run_label) ∈ zip(1:n_runs, ff, density, parallel_flow,
+ thermal_speed, evolve_density, evolve_upar,
+ evolve_ppar, z, vpa, run_labels)
+
+ PyPlot.subplot(1, n_runs, run_ind)
+ @views plot_unnormalised(f[:,:,is,iframe], this_z, this_vpa,
+ n[:,is,iframe], upar[:,is,iframe],
+ vth[:,is,iframe], ev_n, ev_u, ev_p,
+ title=run_label)
+ end
  end
+ fig = PyPlot.figure(1, figsize=(4,4))
+ fig.clf()
+ myanim = matplotlib_animation.FuncAnimation(fig, make_frame, frames=nframes)
  outfile = string(prefix, "_f_unnorm_vs_vpa_z", spec_string, ".gif")
- gif(anim, outfile, fps=5)
- # make a gif animation of log(f_unnorm)(v_parallel_unnorm,z,t)
- anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
- subplots = (@views plot_unnormalised(f[:,:,is,i], this_z, this_vpa, n[:,is,i],
- upar[:,is,i], vth[:,is,i], ev_n, ev_u, ev_p,
- plot_log=true, title=run_label)
- for (f, n, upar, vth, ev_n, ev_u, ev_p, this_z,
- this_vpa, run_label) ∈
- zip(ff, density, parallel_flow, thermal_speed,
- evolve_density, evolve_upar, evolve_ppar, z,
- vpa, run_labels))
- plot(subplots..., layout=(1,n_runs), size=(400*n_runs, 400))
+ myanim.save(outfile, writer=matplotlib_animation.PillowWriter(fps=30))
+
+ function make_frame_log(i)
+ PyPlot.clf()
+ iframe = iframes[i+1]
+ # i counts from 0, Python-style
+ for (run_ind, f, n, upar, vth, ev_n, ev_u, ev_p, this_z, this_vpa,
+ run_label) ∈ zip(1:n_runs, ff, density, parallel_flow,
+ thermal_speed, evolve_density, evolve_upar,
+ evolve_ppar, z, vpa, run_labels)
+
+ PyPlot.subplot(1, n_runs, run_ind)
+ @views plot_unnormalised(f[:,:,is,iframe], this_z, this_vpa,
+ n[:,is,iframe], upar[:,is,iframe],
+ vth[:,is,iframe], ev_n, ev_u, ev_p,
+ title=run_label, plot_log=true)
+ end
  end
+ fig = PyPlot.figure(figsize=(4,4))
+ myanim = matplotlib_animation.FuncAnimation(fig, make_frame_log, frames=nframes)
  outfile = string(prefix, "_logf_unnorm_vs_vpa_z", spec_string, ".gif")
- gif(anim, outfile, fps=5)
+ myanim.save(outfile, writer=matplotlib_animation.PillowWriter(fps=30))
  end
  if pp.animate_deltaf_vs_vpa_z
  # make a gif animation of δf(vpa,z,t)
@@ -1170,15 +1224,34 @@ function plot_unnormalised(f, z_grid, vpa_grid, density, upar, vth, evolve_densi
  f_unnorm = f
  end
 
+ ## The following commented out section does not work at the moment because
+ ## Plots.heatmap() does not support 2d coordinates.
+ ## https://github.com/JuliaPlots/Plots.jl/pull/4298 would add this feature...
+ #if plot_log
+ # @. f_unnorm = log(abs(f_unnorm))
+ # cmlog(cmlin::ColorGradient) = RGB[cmlin[x] for x=LinRange(0,1,30)]
+ # cmap = cgrad(:deep, scale=:log) |> cmlog
+ #else
+ # cmap = :deep
+ #end
+
+ #p = plot(; xlabel="z", ylabel="vpa", c=cmap, kwargs...)
+
+ #heatmap(z2d, dzdt2d, f_unnorm)
+
+ # Use PyPlot directly instead. Unfortunately makes animation a pain...
  if plot_log
- @. f_unnorm = log(abs(f_unnorm))
- cmlog(cmlin::ColorGradient) = RGB[cmlin[x] for x=LinRange(0,1,30)]
- cmap = cgrad(:deep, scale=:log) |> cmlog
+ vmin = minimum(x for x in f if x > 0.0)
+ norm = PyPlot.matplotlib.colors.LogNorm(vmin=vmin, vmax=maximum(f))
  else
- cmap = :deep
+ norm = nothing
  end
+ p = PyPlot.pcolormesh(z2d, dzdt2d, f_unnorm; norm=norm, cmap="viridis_r")
+ PyPlot.xlabel("z")
+ PyPlot.ylabel("vpa")
+ PyPlot.colorbar()
 
- return heatmap(z2d, dzdt2d, f_unnorm; xlabel="z", ylabel="vpa", c=cmap, kwargs...)
+ return p
 end
 
 end