From afae800630330d4014d34a33e4e43df3cc7776f8 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 25 Jun 2022 12:58:24 +0100
Subject: [PATCH 01/17] Option to animate ppar in post processing

---
 src/input_structs.jl         |  2 ++
 src/post_processing.jl       | 10 +++++++++-
 src/post_processing_input.jl |  4 +++-
 3 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/src/input_structs.jl b/src/input_structs.jl
index f560015c46..65dbe70d2d 100644
--- a/src/input_structs.jl
+++ b/src/input_structs.jl
@@ -275,6 +275,8 @@ struct pp_input
     animate_dens_vs_z::Bool
     # if animate_upar_vs_z = true, create animation of species parallel flow vs z at different time slices
     animate_upar_vs_z::Bool
+    # if animate_ppar_vs_z = true, create animation of species parallel pressure vs z at different time slices
+    animate_ppar_vs_z::Bool
     # if animate_f_vs_z_vpa = true, create animation of f(z,vpa) at different time slices
     animate_f_vs_vpa_z::Bool
     # if animate_f_vs_z_vpa0 = true, create animation of f(z,vpa0) at different time slices
diff --git a/src/post_processing.jl b/src/post_processing.jl
index 77f02f14bc..8a2bfb3171 100644
--- a/src/post_processing.jl
+++ b/src/post_processing.jl
@@ -549,13 +549,21 @@ function plot_moments(density, delta_density, density_fldline_avg,
             gif(anim, outfile, fps=5)
         end
         if pp.animate_upar_vs_z
-            # make a gif animation of ϕ(z) at different times
+            # make a gif animation of upar(z) at different times
             anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
                 @views plot(z, parallel_flow[:,is,i], xlabel="z", ylabel="upars/vt", ylims = (upar_min,upar_max))
             end
             outfile = string(run_name, "_upar_vs_z_spec", spec_string, ".gif")
             gif(anim, outfile, fps=5)
         end
+        if pp.animate_ppar_vs_z
+            # make a gif animation of ppar(z) at different times
+            anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
+                @views plot(z, parallel_pressure[:,is,i], xlabel="z", ylabel="ppars", ylims = (ppar_min,ppar_max))
+            end
+            outfile = string(run_name, "_ppar_vs_z_spec", spec_string, ".gif")
+            gif(anim, outfile, fps=5)
+        end
     end
     println("done.")
 end
diff --git a/src/post_processing_input.jl b/src/post_processing_input.jl
index 15725651f1..e2a0eadb94 100644
--- a/src/post_processing_input.jl
+++ b/src/post_processing_input.jl
@@ -37,6 +37,8 @@ const plot_qpar_vs_z_t = false
 const animate_dens_vs_z = true
 # if animate_upar_vs_z = true, create animation of species parallel flow(z) at different time slices
 const animate_upar_vs_z = false
+# if animate_ppar_vs_z = true, create animation of species parallel pressure(z) at different time slices
+const animate_ppar_vs_z = false
 # if animate_f_vs_vpa_z = true, create animation of f(vpa,z) at different time slices
 const animate_f_vs_vpa_z = false
 # if animate_deltaf_vs_vpa_z = true, create animation of δf(vpa,z) at different time slices
@@ -69,7 +71,7 @@ const ir0 = -1
 pp = pp_input(calculate_frequencies, plot_phi0_vs_t, plot_phi_vs_z_t,
     animate_phi_vs_z, plot_dens0_vs_t, plot_upar0_vs_t, plot_ppar0_vs_t, plot_qpar0_vs_t,
     plot_dens_vs_z_t, plot_upar_vs_z_t, plot_ppar_vs_z_t, plot_qpar_vs_z_t,
-    animate_dens_vs_z, animate_upar_vs_z,
+    animate_dens_vs_z, animate_upar_vs_z, animate_ppar_vs_z,
     animate_f_vs_vpa_z, animate_f_vs_vpa_z0, animate_f_vs_z0_vpa,
     animate_deltaf_vs_vpa_z, animate_deltaf_vs_vpa_z0, animate_deltaf_vs_vpa_z0,
     nwrite_movie, itime_min, itime_max, ivpa0, iz0, ir0)

From e4f49e55c88ba15ba91f8cd91ad5bf98e5e82a32 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 5 Jul 2022 15:32:27 +0100
Subject: [PATCH 02/17] Option to animate qpar_vs_z

---
 src/input_structs.jl         | 2 ++
 src/post_processing.jl       | 8 ++++++++
 src/post_processing_input.jl | 4 +++-
 3 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/src/input_structs.jl b/src/input_structs.jl
index 65dbe70d2d..4178a8e54c 100644
--- a/src/input_structs.jl
+++ b/src/input_structs.jl
@@ -277,6 +277,8 @@ struct pp_input
     animate_upar_vs_z::Bool
     # if animate_ppar_vs_z = true, create animation of species parallel pressure vs z at different time slices
     animate_ppar_vs_z::Bool
+    # if animate_qpar_vs_z = true, create animation of species parallel heat flux vs z at different time slices
+    animate_qpar_vs_z::Bool
     # if animate_f_vs_z_vpa = true, create animation of f(z,vpa) at different time slices
     animate_f_vs_vpa_z::Bool
     # if animate_f_vs_z_vpa0 = true, create animation of f(z,vpa0) at different time slices
diff --git a/src/post_processing.jl b/src/post_processing.jl
index 8a2bfb3171..54551f5163 100644
--- a/src/post_processing.jl
+++ b/src/post_processing.jl
@@ -564,6 +564,14 @@ function plot_moments(density, delta_density, density_fldline_avg,
             outfile = string(run_name, "_ppar_vs_z_spec", spec_string, ".gif")
             gif(anim, outfile, fps=5)
         end
+        if pp.animate_qpar_vs_z
+            # make a gif animation of ppar(z) at different times
+            anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
+                @views plot(z, parallel_heat_flux[:,is,i], xlabel="z", ylabel="qpars", ylims = (qpar_min,qpar_max))
+            end
+            outfile = string(run_name, "_qpar_vs_z_spec", spec_string, ".gif")
+            gif(anim, outfile, fps=5)
+        end
     end
     println("done.")
 end
diff --git a/src/post_processing_input.jl b/src/post_processing_input.jl
index e2a0eadb94..4a63d39628 100644
--- a/src/post_processing_input.jl
+++ b/src/post_processing_input.jl
@@ -39,6 +39,8 @@ const animate_dens_vs_z = true
 const animate_upar_vs_z = false
 # if animate_ppar_vs_z = true, create animation of species parallel pressure(z) at different time slices
 const animate_ppar_vs_z = false
+# if animate_qpar_vs_z = true, create animation of species parallel heat flux(z) at different time slices
+const animate_qpar_vs_z = false
 # if animate_f_vs_vpa_z = true, create animation of f(vpa,z) at different time slices
 const animate_f_vs_vpa_z = false
 # if animate_deltaf_vs_vpa_z = true, create animation of δf(vpa,z) at different time slices
@@ -71,7 +73,7 @@ const ir0 = -1
 pp = pp_input(calculate_frequencies, plot_phi0_vs_t, plot_phi_vs_z_t,
     animate_phi_vs_z, plot_dens0_vs_t, plot_upar0_vs_t, plot_ppar0_vs_t, plot_qpar0_vs_t,
     plot_dens_vs_z_t, plot_upar_vs_z_t, plot_ppar_vs_z_t, plot_qpar_vs_z_t,
-    animate_dens_vs_z, animate_upar_vs_z, animate_ppar_vs_z,
+    animate_dens_vs_z, animate_upar_vs_z, animate_ppar_vs_z, animate_qpar_vs_z,
     animate_f_vs_vpa_z, animate_f_vs_vpa_z0, animate_f_vs_z0_vpa,
     animate_deltaf_vs_vpa_z, animate_deltaf_vs_vpa_z0, animate_deltaf_vs_vpa_z0,
     nwrite_movie, itime_min, itime_max, ivpa0, iz0, ir0)

From d4f140ea77ca8127395e1af9c7dfa1fdcc252151 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 25 Aug 2022 17:17:21 +0100
Subject: [PATCH 03/17] Options to make plots of thermal speed in post
 processing

---
 src/analysis.jl              | 20 +++++++++++++++++---
 src/input_structs.jl         |  4 ++++
 src/post_processing.jl       | 34 ++++++++++++++++++++++++++++++----
 src/post_processing_input.jl |  8 ++++++--
 4 files changed, 57 insertions(+), 9 deletions(-)

diff --git a/src/analysis.jl b/src/analysis.jl
index 71b5e5f2fa..c73b4a4298 100644
--- a/src/analysis.jl
+++ b/src/analysis.jl
@@ -29,7 +29,8 @@ end
 
 """
 """
-function analyze_moments_data(density, parallel_flow, parallel_pressure, parallel_heat_flux, ntime, n_species, nz, z_wgts, Lz)
+function analyze_moments_data(density, parallel_flow, parallel_pressure, thermal_speed,
+                              parallel_heat_flux, ntime, n_species, nz, z_wgts, Lz)
     print("Analyzing velocity moments data...")
     density_fldline_avg = allocate_float(n_species, ntime)
     for is ∈ 1:n_species
@@ -49,6 +50,12 @@ function analyze_moments_data(density, parallel_flow, parallel_pressure, paralle
             ppar_fldline_avg[is,i] = field_line_average(view(parallel_pressure,:,is,i), z_wgts, Lz)
         end
     end
+    vth_fldline_avg = allocate_float(n_species, ntime)
+    for is ∈ 1:n_species
+        for i ∈ 1:ntime
+            vth_fldline_avg[is,i] = field_line_average(view(thermal_speed,:,is,i), z_wgts, Lz)
+        end
+    end
     qpar_fldline_avg = allocate_float(n_species, ntime)
     for is ∈ 1:n_species
         for i ∈ 1:ntime
@@ -76,6 +83,13 @@ function analyze_moments_data(density, parallel_flow, parallel_pressure, paralle
             @. delta_ppar[iz,is,:] = parallel_pressure[iz,is,:] - ppar_fldline_avg[is,:]
         end
     end
+    # delta_vth = vth_s - <vth_s> is the fluctuating thermal_speed
+    delta_vth = allocate_float(nz,n_species,ntime)
+    for is ∈ 1:n_species
+        for iz ∈ 1:nz
+            @. delta_vth[iz,is,:] = thermal_speed[iz,is,:] - vth_fldline_avg[is,:]
+        end
+    end
     # delta_qpar = qpar_s - <qpar_s> is the fluctuating parallel heat flux
     delta_qpar = allocate_float(nz,n_species,ntime)
     for is ∈ 1:n_species
@@ -84,8 +98,8 @@ function analyze_moments_data(density, parallel_flow, parallel_pressure, paralle
         end
     end
     println("done.")
-    return density_fldline_avg, upar_fldline_avg, ppar_fldline_avg, qpar_fldline_avg,
-           delta_density, delta_upar, delta_ppar, delta_qpar
+    return density_fldline_avg, upar_fldline_avg, ppar_fldline_avg, vth_fldline_avg, qpar_fldline_avg,
+           delta_density, delta_upar, delta_ppar, delta_vth, delta_qpar
 end
 
 """
diff --git a/src/input_structs.jl b/src/input_structs.jl
index 4178a8e54c..5f92b7e762 100644
--- a/src/input_structs.jl
+++ b/src/input_structs.jl
@@ -261,6 +261,8 @@ struct pp_input
     plot_upar0_vs_t::Bool
     # if plot_ppar0_vs_t = true, create plots of species ppar(z0) vs time
     plot_ppar0_vs_t::Bool
+    # if plot_vth0_vs_t = true, create plots of species vth(z0) vs time
+    plot_vth0_vs_t::Bool
     # if plot_qpar0_vs_t = true, create plots of species qpar(z0) vs time
     plot_qpar0_vs_t::Bool
     # if plot_dens_vs_z_t = true, create plot of species density vs z and time
@@ -277,6 +279,8 @@ struct pp_input
     animate_upar_vs_z::Bool
     # if animate_ppar_vs_z = true, create animation of species parallel pressure vs z at different time slices
     animate_ppar_vs_z::Bool
+    # if animate_vth_vs_z = true, create animation of species thermal speed vs z at different time slices
+    animate_vth_vs_z::Bool
     # if animate_qpar_vs_z = true, create animation of species parallel heat flux vs z at different time slices
     animate_qpar_vs_z::Bool
     # if animate_f_vs_z_vpa = true, create animation of f(z,vpa) at different time slices
diff --git a/src/post_processing.jl b/src/post_processing.jl
index 54551f5163..e702810091 100644
--- a/src/post_processing.jl
+++ b/src/post_processing.jl
@@ -151,14 +151,15 @@ function plot_1D_1V_diagnostics(run_name, fid, nwrite_movie, itime_min, itime_ma
                 z, iz0, run_name, fitted_delta_phi, pp)
     # load velocity moments data
     # analyze the velocity moments data
-    density_fldline_avg, upar_fldline_avg, ppar_fldline_avg, qpar_fldline_avg,
-        delta_density, delta_upar, delta_ppar, delta_qpar =
-        analyze_moments_data(density, parallel_flow, parallel_pressure, parallel_heat_flux,
-                             ntime, n_species, nz, z_wgts, Lz)
+    density_fldline_avg, upar_fldline_avg, ppar_fldline_avg, vth_fldline_avg, qpar_fldline_avg,
+        delta_density, delta_upar, delta_ppar, delta_vth, delta_qpar =
+        analyze_moments_data(density, parallel_flow, parallel_pressure, thermal_speed,
+                             parallel_heat_flux, ntime, n_species, nz, z_wgts, Lz)
     # create the requested plots of the moments
     plot_moments(density, delta_density, density_fldline_avg,
         parallel_flow, delta_upar, upar_fldline_avg,
         parallel_pressure, delta_ppar, ppar_fldline_avg,
+        thermal_speed, delta_vth, vth_fldline_avg,
         parallel_heat_flux, delta_qpar, qpar_fldline_avg,
         pp, run_name, time, itime_min, itime_max,
         nwrite_movie, z, iz0, n_species)
@@ -433,6 +434,7 @@ end
 function plot_moments(density, delta_density, density_fldline_avg,
     parallel_flow, delta_upar, upar_fldline_avg,
     parallel_pressure, delta_ppar, ppar_fldline_avg,
+    thermal_speed, delta_vth, vth_fldline_avg,
     parallel_heat_flux, delta_qpar, qpar_fldline_avg,
     pp, run_name, time, itime_min, itime_max, nwrite_movie,
     z, iz0, n_species)
@@ -500,6 +502,22 @@ function plot_moments(density, delta_density, density_fldline_avg,
             outfile = string(run_name, "_fldline_avg_ppar_vs_t_spec", spec_string, ".pdf")
             savefig(outfile)
         end
+        vth_min = minimum(thermal_speed[:,is,:])
+        vth_max = maximum(thermal_speed[:,is,:])
+        if pp.plot_vth0_vs_t
+            # plot the time trace of n_s(z=z0)
+            @views plot(time, thermal_speed[iz0,is,:])
+            outfile = string(run_name, "_vth0_vs_t_spec", spec_string, ".pdf")
+            savefig(outfile)
+            # plot the time trace of n_s(z=z0)-density_fldline_avg
+            @views plot(time, abs.(delta_vth[iz0,is,:]), yaxis=:log)
+            outfile = string(run_name, "_delta_vth0_vs_t_spec", spec_string, ".pdf")
+            savefig(outfile)
+            # plot the time trace of vth_fldline_avg
+            @views plot(time, vth_fldline_avg[is,:], xlabel="time", ylabel="<vths/cₛ₀>", ylims=(vth_min,vth_max))
+            outfile = string(run_name, "_fldline_avg_vth_vs_t_spec", spec_string, ".pdf")
+            savefig(outfile)
+        end
         qpar_min = minimum(parallel_heat_flux[:,is,:])
         qpar_max = maximum(parallel_heat_flux[:,is,:])
         if pp.plot_qpar0_vs_t
@@ -564,6 +582,14 @@ function plot_moments(density, delta_density, density_fldline_avg,
             outfile = string(run_name, "_ppar_vs_z_spec", spec_string, ".gif")
             gif(anim, outfile, fps=5)
         end
+        if pp.animate_vth_vs_z
+            # make a gif animation of vth(z) at different times
+            anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
+                @views plot(z, thermal_speed[:,is,i], xlabel="z", ylabel="vths", ylims = (vth_min,vth_max))
+            end
+            outfile = string(run_name, "_vth_vs_z_spec", spec_string, ".gif")
+            gif(anim, outfile, fps=5)
+        end
         if pp.animate_qpar_vs_z
             # make a gif animation of ppar(z) at different times
             anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
diff --git a/src/post_processing_input.jl b/src/post_processing_input.jl
index 4a63d39628..328640f4dd 100644
--- a/src/post_processing_input.jl
+++ b/src/post_processing_input.jl
@@ -23,6 +23,8 @@ const plot_dens0_vs_t = true
 const plot_upar0_vs_t = false
 # if plot_ppar0_vs_t = true, create plots of species ppar(z0) vs time
 const plot_ppar0_vs_t = false
+# if plot_vth0_vs_t = true, create plots of species vth(z0) vs time
+const plot_vth0_vs_t = false
 # if plot_qpar0_vs_t = true, create plots of species qpar(z0) vs time
 const plot_qpar0_vs_t = false
 # if plot_dens_vs_z_t = true, create heatmap of species density vs z and time
@@ -39,6 +41,8 @@ const animate_dens_vs_z = true
 const animate_upar_vs_z = false
 # if animate_ppar_vs_z = true, create animation of species parallel pressure(z) at different time slices
 const animate_ppar_vs_z = false
+# if animate_vth_vs_z = true, create animation of species thermal_velocity(z) at different time slices
+const animate_vth_vs_z = false
 # if animate_qpar_vs_z = true, create animation of species parallel heat flux(z) at different time slices
 const animate_qpar_vs_z = false
 # if animate_f_vs_vpa_z = true, create animation of f(vpa,z) at different time slices
@@ -71,9 +75,9 @@ const iz0 = -1
 const ir0 = -1
 
 pp = pp_input(calculate_frequencies, plot_phi0_vs_t, plot_phi_vs_z_t,
-    animate_phi_vs_z, plot_dens0_vs_t, plot_upar0_vs_t, plot_ppar0_vs_t, plot_qpar0_vs_t,
+    animate_phi_vs_z, plot_dens0_vs_t, plot_upar0_vs_t, plot_ppar0_vs_t, plot_vth0_vs_t, plot_qpar0_vs_t,
     plot_dens_vs_z_t, plot_upar_vs_z_t, plot_ppar_vs_z_t, plot_qpar_vs_z_t,
-    animate_dens_vs_z, animate_upar_vs_z, animate_ppar_vs_z, animate_qpar_vs_z,
+    animate_dens_vs_z, animate_upar_vs_z, animate_ppar_vs_z, animate_vth_vs_z, animate_qpar_vs_z,
     animate_f_vs_vpa_z, animate_f_vs_vpa_z0, animate_f_vs_z0_vpa,
     animate_deltaf_vs_vpa_z, animate_deltaf_vs_vpa_z0, animate_deltaf_vs_vpa_z0,
     nwrite_movie, itime_min, itime_max, ivpa0, iz0, ir0)

From 93b27feb7976301bc233e1764130c52404966a9e Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 4 Aug 2022 12:57:35 +0100
Subject: [PATCH 04/17] Function for adding line showing physical v_parallel=0
 on z/vpa 2d plots

---
 src/post_processing.jl | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/src/post_processing.jl b/src/post_processing.jl
index e702810091..5a8ad89c9e 100644
--- a/src/post_processing.jl
+++ b/src/post_processing.jl
@@ -800,4 +800,25 @@ end
 #    println("advection_test_1d rms error: ", rmserr)
 #end
 
+"""
+Add a thin, red, dashed line showing v_parallel=(vth*w_parallel+upar)=0 to a 2d plot
+with axes (z,vpa).
+"""
+function draw_v_parallel_zero!(plt::Plots.Plot, z::AbstractVector, upar, vth,
+                               evolve_upar::Bool, evolve_ppar::Bool)
+    if evolve_ppar && evolve_upar
+        zero_value = @. -upar/vth
+    elseif evolve_upar
+        zero_value = @. -upar
+    else
+        zero_value = zeros(size(upar))
+    end
+    plot!(plt, z, zero_value, color=:red, linestyle=:dash, linewidth=1,
+          xlims=xlims(plt), ylims=ylims(plt), label="")
+end
+function draw_v_parallel_zero!(z::AbstractVector, upar, vth, evolve_upar::Bool,
+                               evolve_ppar::Bool)
+    draw_v_parallel_zero!(Plots.CURRENT_PLOT, z, upar, vth, evolve_upar, evolve_ppar)
+end
+
 end

From af3567453db71f22ed7e0fb314268b9dfdfae5ee Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 4 Aug 2022 12:58:13 +0100
Subject: [PATCH 05/17] Optionally make some plots at each output step to
 monitor progress

Can be useful to monitor an ongoing simulation. This cannot be done by
reading the output file, because opening the NetCDF file from two
programs at once may corrupt it and crash either/both programs. Instead,
if `runtime_plots = true` is passed in the input, a .png file with a
bunch of plots of moments and distribution functions is produced at each
output step. The plots are created in serial, and may take some time
(~10s?) for high-resolution simulations.
---
 debug_test/sound_wave_tests.jl |  3 +-
 src/input_structs.jl           |  1 +
 src/moment_kinetics.jl         |  9 +++--
 src/moment_kinetics_input.jl   |  4 ++-
 src/time_advance.jl            | 61 ++++++++++++++++++++++++++++++++++
 5 files changed, 71 insertions(+), 7 deletions(-)

diff --git a/debug_test/sound_wave_tests.jl b/debug_test/sound_wave_tests.jl
index d614f7a12f..d430033c06 100644
--- a/debug_test/sound_wave_tests.jl
+++ b/debug_test/sound_wave_tests.jl
@@ -90,7 +90,8 @@ test_input_chebyshev_split_2_moments =
 test_input_chebyshev_split_3_moments =
     merge(test_input_chebyshev_split_2_moments,
           Dict("run_name" => "chebyshev_pseudospectral_split_3_moments",
-               "evolve_moments_parallel_pressure" => true))
+               "evolve_moments_parallel_pressure" => true,
+               "runtime_plots" => true))
 
 
 """
diff --git a/src/input_structs.jl b/src/input_structs.jl
index 5f92b7e762..8c9bb91ba1 100644
--- a/src/input_structs.jl
+++ b/src/input_structs.jl
@@ -34,6 +34,7 @@ struct time_input
     use_semi_lagrange::Bool
     n_rk_stages::mk_int
     split_operators::Bool
+    runtime_plots::Bool
 end
 
 """
diff --git a/src/moment_kinetics.jl b/src/moment_kinetics.jl
index 1310b837ff..c9b5a977ae 100644
--- a/src/moment_kinetics.jl
+++ b/src/moment_kinetics.jl
@@ -43,15 +43,14 @@ include("energy_equation.jl")
 include("force_balance.jl")
 include("source_terms.jl")
 include("numerical_dissipation.jl")
-include("time_advance.jl")
-
-include("moment_kinetics_input.jl")
-include("scan_input.jl")
-
 include("analysis.jl")
 include("load_data.jl")
 include("post_processing_input.jl")
 include("post_processing.jl")
+include("time_advance.jl")
+
+include("moment_kinetics_input.jl")
+include("scan_input.jl")
 
 using TimerOutputs
 using TOML
diff --git a/src/moment_kinetics_input.jl b/src/moment_kinetics_input.jl
index 6c992fd088..cc8f9752ca 100644
--- a/src/moment_kinetics_input.jl
+++ b/src/moment_kinetics_input.jl
@@ -135,6 +135,7 @@ function mk_input(scan_input=Dict())
     # Heun's method, SSP RK3 and 4-stage SSP RK3)
     n_rk_stages = get(scan_input, "n_rk_stages", 4)
     split_operators = get(scan_input, "split_operators", false)
+    runtime_plots = get(scan_input, "runtime_plots", false)
 
     # overwrite some default parameters related to the r grid
     # ngrid is number of grid points per element
@@ -184,7 +185,8 @@ function mk_input(scan_input=Dict())
     ########## end user inputs. do not modify following code! ###############
     #########################################################################
 
-    t = time_input(nstep, dt, nwrite, use_semi_lagrange, n_rk_stages, split_operators)
+    t = time_input(nstep, dt, nwrite, use_semi_lagrange, n_rk_stages, split_operators,
+                   runtime_plots)
     # replace mutable structures with immutable ones to optimize performance
     # and avoid possible misunderstandings
     z_advection_immutable = advection_input(z.advection.option, z.advection.constant_speed,
diff --git a/src/time_advance.jl b/src/time_advance.jl
index 8af4495bc5..b847af1217 100644
--- a/src/time_advance.jl
+++ b/src/time_advance.jl
@@ -38,6 +38,8 @@ using ..semi_lagrange: setup_semi_lagrange
 @debug_detect_redundant_block_synchronize using ..communication: debug_detect_redundant_is_active
 
 using Dates
+using Plots
+using ..post_processing: draw_v_parallel_zero!
 
 """
 """
@@ -385,6 +387,65 @@ function time_advance!(pdf, scratch, t, t_input, vpa, z, r, vpa_spectral, z_spec
             write_data_to_ascii(pdf.norm, moments, fields, vpa, z, r, t, composition.n_species, io)
             # write initial data to binary file (netcdf)
             write_data_to_binary(pdf.norm, moments, fields, t, composition.n_species, cdf, iwrite)
+            # Hack to save *.pdf of current pdf
+            if t_input.runtime_plots
+                @serial_region begin
+                    #pyplot()
+                    cmlog(cmlin::ColorGradient) = RGB[cmlin[x] for x=LinRange(0,1,30)]
+                    logdeep = cgrad(:deep, scale=:log) |> cmlog
+                    f_plots = [
+                        heatmap(z.grid, vpa.grid, pdf.norm[:,:,1,is],
+                                xlim=(z.grid[1] - z.L / 100.0, z.grid[end] + z.L / 100.0),
+                                ylim=(vpa.grid[1] - vpa.L / 100.0, vpa.grid[end] + vpa.L / 100.0),
+                                xlabel="z", ylabel="vpa", c=:deep, colorbar=false)
+                        for is ∈ 1:composition.n_species]
+                    for (is, p) in enumerate(f_plots)
+                        @views draw_v_parallel_zero!(p, z.grid, moments.upar[:,1,is],
+                                                     moments.vth[:,1,is],
+                                                     moments.evolve_upar,
+                                                     moments.evolve_ppar)
+                    end
+                    logf_plots = [
+                        heatmap(z.grid, vpa.grid, log.(abs.(pdf.norm[:,:,1,is])),
+                                xlim=(z.grid[1] - z.L / 100.0, z.grid[end] + z.L / 100.0),
+                                ylim=(vpa.grid[1] - vpa.L / 100.0, vpa.grid[end] + vpa.L / 100.0),
+                                xlabel="z", ylabel="vpa", fillcolor=logdeep, colorbar=false)
+                        for is ∈ 1:composition.n_species]
+                    for (is, p) in enumerate(logf_plots)
+                        @views draw_v_parallel_zero!(p, z.grid, moments.upar[:,1,is],
+                                                     moments.vth[:,1,is],
+                                                     moments.evolve_upar,
+                                                     moments.evolve_ppar)
+                    end
+                    f0_plots = [
+                        plot(vpa.grid, pdf.norm[:,1,1,is], xlabel="vpa", ylabel="f0", legend=false)
+                        for is ∈ 1:composition.n_species]
+                    fL_plots = [
+                        plot(vpa.grid, pdf.norm[:,end,1,is], xlabel="vpa", ylabel="fL", legend=false)
+                        for is ∈ 1:composition.n_species]
+                    density_plots = [
+                        plot(z.grid, moments.dens[:,1,is], xlabel="z", ylabel="density", legend=false)
+                        for is ∈ 1:composition.n_species]
+                    upar_plots = [
+                        plot(z.grid, moments.upar[:,1,is], xlabel="z", ylabel="upar", legend=false)
+                        for is ∈ 1:composition.n_species]
+                    ppar_plots = [
+                        plot(z.grid, moments.ppar[:,1,is], xlabel="z", ylabel="ppar", legend=false)
+                        for is ∈ 1:composition.n_species]
+                    vth_plots = [
+                        plot(z.grid, moments.vth[:,1,is], xlabel="z", ylabel="vth", legend=false)
+                        for is ∈ 1:composition.n_species]
+                    qpar_plots = [
+                        plot(z.grid, moments.qpar[:,1,is], xlabel="z", ylabel="qpar", legend=false)
+                        for is ∈ 1:composition.n_species]
+                    # Put all plots into subplots of a single figure
+                    plot(f_plots..., logf_plots..., f0_plots..., fL_plots...,
+                         density_plots..., upar_plots..., ppar_plots..., vth_plots...,
+                         qpar_plots...,
+                         layout=(9,composition.n_species), size=(800,3600), plot_title="$t")
+                    savefig("latest_plots.png")
+                end
+            end
             iwrite += 1
 
             # Restore to same region type as in rk_update!() so that execution after a

From 3ab49761d5a0ca7ed27a14886db80b761f98957f Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 5 Aug 2022 14:53:49 +0100
Subject: [PATCH 06/17] Save and load evolve_density and evolve_upar

These can be useful for post-processing.
---
 src/file_io.jl         | 27 +++++++++++++++++++++------
 src/load_data.jl       | 21 ++++++++++++++++++++-
 src/moment_kinetics.jl |  1 +
 src/post_processing.jl |  3 ++-
 4 files changed, 44 insertions(+), 8 deletions(-)

diff --git a/src/file_io.jl b/src/file_io.jl
index 1c78cf7fe7..90ed8154cc 100644
--- a/src/file_io.jl
+++ b/src/file_io.jl
@@ -56,7 +56,7 @@ end
 open the necessary output files
 """
 function setup_file_io(output_dir, run_name, vpa, z, r, composition,
-                       collisions, evolve_ppar)
+                       collisions, evolve_density, evolve_upar, evolve_ppar)
     begin_serial_region()
     @serial_region begin
         # Only read/write from first process in each 'block'
@@ -69,7 +69,7 @@ function setup_file_io(output_dir, run_name, vpa, z, r, composition,
         mom_io = open_output_file(out_prefix, "moments_vs_t")
         fields_io = open_output_file(out_prefix, "fields_vs_t")
         cdf = setup_netcdf_io(out_prefix, r, z, vpa, composition, collisions,
-                              evolve_ppar)
+                              evolve_density, evolve_upar, evolve_ppar)
         #return ios(ff_io, mom_io, fields_io), cdf
         return ios(mom_io, fields_io), cdf
     end
@@ -110,11 +110,11 @@ function define_dimensions!(fid, nvpa, nz, nr, n_species, n_ion_species=nothing,
 end
 
 """
-    define_static_variables!(vpa,vperp,z,r,composition,collisions,evolve_ppar)
+    define_static_variables!(vpa,vperp,z,r,composition,collisions,evolve_density,evolve_upar,evolve_ppar)
 
 Define static (i.e. time-independent) variables for an output file.
 """
-function define_static_variables!(fid,vpa,z,r,composition,collisions,evolve_ppar)
+function define_static_variables!(fid,vpa,z,r,composition,collisions,evolve_density,evolve_upar,evolve_ppar)
     # create and write the "r" variable to file
     varname = "r"
     attributes = Dict("description" => "radial coordinate")
@@ -168,6 +168,20 @@ function define_static_variables!(fid,vpa,z,r,composition,collisions,evolve_ppar
     vartype = mk_float
     var = defVar(fid, varname, vartype, dims, attrib=attributes)
     var[:] = collisions.charge_exchange
+    # create and write the "evolve_density" variable to file
+    varname = "evolve_density"
+    attributes = Dict("description" => "flag indicating if the density is separately evolved")
+    vartype = mk_int
+    dims = ("n_species",)
+    var = defVar(fid, varname, vartype, dims, attrib=attributes)
+    var[:] = evolve_density
+    # create and write the "evolve_upar" variable to file
+    varname = "evolve_upar"
+    attributes = Dict("description" => "flag indicating if the parallel flow is separately evolved")
+    vartype = mk_int
+    dims = ("n_species",)
+    var = defVar(fid, varname, vartype, dims, attrib=attributes)
+    var[:] = evolve_upar
     # create and write the "evolve_ppar" variable to file
     varname = "evolve_ppar"
     attributes = Dict("description" => "flag indicating if the parallel pressure is separately evolved")
@@ -240,7 +254,8 @@ end
 """
 setup file i/o for netcdf
 """
-function setup_netcdf_io(prefix, r, z, vpa, composition, collisions, evolve_ppar)
+function setup_netcdf_io(prefix, r, z, vpa, composition, collisions, evolve_density,
+                         evolve_upar, evolve_ppar)
     # the netcdf file will be given by output_dir/run_name with .cdf appended
     filename = string(prefix,".cdf")
     # if a netcdf file with the requested name already exists, remove it
@@ -253,7 +268,7 @@ function setup_netcdf_io(prefix, r, z, vpa, composition, collisions, evolve_ppar
     define_dimensions!(fid, vpa.n, z.n, r.n, composition.n_species,
                        composition.n_ion_species, composition.n_neutral_species)
     ### create and write static variables to file ###
-    define_static_variables!(fid,vpa,z,r,composition,collisions,evolve_ppar)
+    define_static_variables!(fid,vpa,z,r,composition,collisions,evolve_density,evolve_upar,evolve_ppar)
     ### create variables for time-dependent quantities and store them ###
     ### in a struct for later access ###
     cdf_time, cdf_f, cdf_phi, cdf_density, cdf_upar, cdf_ppar, cdf_qpar, cdf_vth =
diff --git a/src/load_data.jl b/src/load_data.jl
index 8587e15af9..f70a100e9b 100644
--- a/src/load_data.jl
+++ b/src/load_data.jl
@@ -112,6 +112,24 @@ function load_moments_data(fid)
     thermal_speed = cdfvar.var[:,:,:,:]
     # define the number of species
     n_species = size(cdfvar,3)
+    # define a handle for the flag indicating if the density should be separately advanced
+    cdfvar = fid["evolve_density"]
+    # load the parallel pressure evolution flag
+    evolve_density_int = cdfvar.var[:]
+    if evolve_density_int[1] == 1
+        evolve_density = true
+    else
+        evolve_density = false
+    end
+    # define a handle for the flag indicating if the parallel pressure should be separately advanced
+    cdfvar = fid["evolve_upar"]
+    # load the parallel pressure evolution flag
+    evolve_upar_int = cdfvar.var[:]
+    if evolve_upar_int[1] == 1
+        evolve_upar = true
+    else
+        evolve_upar = false
+    end
     # define a handle for the flag indicating if the parallel pressure should be separately advanced
     cdfvar = fid["evolve_ppar"]
     # load the parallel pressure evolution flag
@@ -122,7 +140,8 @@ function load_moments_data(fid)
         evolve_ppar = false
     end
     println("done.")
-    return density, parallel_flow, parallel_pressure, parallel_heat_flux, thermal_speed, n_species, evolve_ppar
+    return density, parallel_flow, parallel_pressure, parallel_heat_flux, thermal_speed,
+           n_species, evolve_density, evolve_upar, evolve_ppar
 end
 
 """
diff --git a/src/moment_kinetics.jl b/src/moment_kinetics.jl
index c9b5a977ae..5fc8969800 100644
--- a/src/moment_kinetics.jl
+++ b/src/moment_kinetics.jl
@@ -272,6 +272,7 @@ function setup_moment_kinetics(input_dict::Dict; backup_filename=nothing,
 
     # setup i/o
     io, cdf = setup_file_io(output_dir, run_name, vpa, z, r, composition, collisions,
+                            moments.evolve_density, moments.evolve_upar,
                             moments.evolve_ppar)
     # write initial data to ascii files
     write_data_to_ascii(pdf.norm, moments, fields, vpa, z, r, code_time, composition.n_species, io)
diff --git a/src/post_processing.jl b/src/post_processing.jl
index 5a8ad89c9e..456c808523 100644
--- a/src/post_processing.jl
+++ b/src/post_processing.jl
@@ -69,7 +69,8 @@ function analyze_and_plot_data(path)
     phi = load_fields_data(fid)
     # load full (z,r,species,t) velocity moments data
     density, parallel_flow, parallel_pressure, parallel_heat_flux,
-        thermal_speed, n_species, evolve_ppar = load_moments_data(fid)
+        thermal_speed, n_species, evolve_density, evolve_upar, evolve_ppar =
+        load_moments_data(fid)
     # load full (vpa,z,r,species,t) particle distribution function (pdf) data
     ff = load_pdf_data(fid)
 

From 83dd53361d789cf93daaf78835b5dbecebfd63f7 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 5 Aug 2022 15:00:35 +0100
Subject: [PATCH 07/17] Option to plot unnormalised f against z and
 unnormalised v_parallel

Also requires updated version of Plots.jl (to be added later if/when
https://github.com/JuliaPlots/Plots.jl/pull/4298 is merged).
---
 src/input_structs.jl         |  5 ++-
 src/post_processing.jl       | 73 ++++++++++++++++++++++++++++++++++--
 src/post_processing_input.jl |  9 +++--
 3 files changed, 80 insertions(+), 7 deletions(-)

diff --git a/src/input_structs.jl b/src/input_structs.jl
index 8c9bb91ba1..9602536a1f 100644
--- a/src/input_structs.jl
+++ b/src/input_structs.jl
@@ -284,8 +284,11 @@ struct pp_input
     animate_vth_vs_z::Bool
     # if animate_qpar_vs_z = true, create animation of species parallel heat flux vs z at different time slices
     animate_qpar_vs_z::Bool
-    # if animate_f_vs_z_vpa = true, create animation of f(z,vpa) at different time slices
+    # if animate_f_vs_vpa_z = true, create animation of f(z,vpa) at different time slices
     animate_f_vs_vpa_z::Bool
+    # if animate_f_unnormalized = true, create animation of
+    # f_unnorm(v_parallel_unnorm,z) at different time slices
+    animate_f_unnormalized::Bool
     # if animate_f_vs_z_vpa0 = true, create animation of f(z,vpa0) at different time slices
     animate_f_vs_vpa0_z::Bool
     # if animate_f_vs_z0_vpa = true, create animation of f(z0,vpa) at different time slices
diff --git a/src/post_processing.jl b/src/post_processing.jl
index 456c808523..e56cd712fa 100644
--- a/src/post_processing.jl
+++ b/src/post_processing.jl
@@ -18,6 +18,7 @@ using ..quadrature: composite_simpson_weights
 using ..array_allocation: allocate_float
 using ..file_io: open_output_file
 using ..type_definitions: mk_float, mk_int
+using ..initial_conditions: vpagrid_to_dzdt
 using ..load_data: open_netcdf_file
 using ..load_data: load_coordinate_data, load_fields_data, load_moments_data, load_pdf_data
 using ..analysis: analyze_fields_data, analyze_moments_data, analyze_pdf_data
@@ -83,7 +84,7 @@ function analyze_and_plot_data(path)
         parallel_heat_flux[:,ir0,:,:],
         thermal_speed[:,ir0,:,:],
         ff[:,:,ir0,:,:],
-        n_species, evolve_ppar, nvpa, vpa, vpa_wgts,
+        n_species, evolve_density, evolve_upar, evolve_ppar, nvpa, vpa, vpa_wgts,
         nz, z, z_wgts, Lz, ntime, time)
 
     close(fid)
@@ -138,8 +139,13 @@ end
 """
 function plot_1D_1V_diagnostics(run_name, fid, nwrite_movie, itime_min, itime_max, ivpa0, iz0, ir0, r,
     phi, density, parallel_flow, parallel_pressure, parallel_heat_flux,
-    thermal_speed, ff, n_species, evolve_ppar, nvpa, vpa, vpa_wgts,
-    nz, z, z_wgts, Lz, ntime, time)
+    thermal_speed, ff, n_species, evolve_density, evolve_upar, evolve_ppar, nvpa, vpa,
+    vpa_wgts, nz, z, z_wgts, Lz, ntime, time)
+
+    # plot_unnormalised() requires PyPlot, so ensure it is used for all plots for
+    # consistency
+    pyplot()
+
     # analyze the fields data
     phi_fldline_avg, delta_phi = analyze_fields_data(phi, ntime, nz, z_wgts, Lz)
     # use a fit to calculate and write to file the damping rate and growth rate of the
@@ -223,6 +229,16 @@ function plot_1D_1V_diagnostics(run_name, fid, nwrite_movie, itime_min, itime_ma
             outfile = string(run_name, "_f_vs_z_vpa_final", spec_string, ".pdf")
             savefig(outfile)
         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
+                @views plot_unnormalised(ff[:,:,is,i], z, vpa, density[:,is,i],
+                                         parallel_flow[:,is,i], thermal_speed[:,is,i],
+                                         evolve_density, evolve_upar, evolve_ppar)
+            end
+            outfile = string(run_name, "_f_unnorm_vs_vpa_z", spec_string, ".gif")
+            gif(anim, outfile, fps=5)
+        end
         if pp.animate_deltaf_vs_vpa_z
             # make a gif animation of δf(vpa,z,t)
             anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
@@ -822,4 +838,55 @@ function draw_v_parallel_zero!(z::AbstractVector, upar, vth, evolve_upar::Bool,
     draw_v_parallel_zero!(Plots.CURRENT_PLOT, z, upar, vth, evolve_upar, evolve_ppar)
 end
 
+"""
+Plot the unnormalised distribution function against unnormalised ('lab space')
+coordinates.
+
+Inputs should depend only on z and vpa.
+
+Note this function requires using the PyPlot backend to support 2d coordinates being
+passed to `heatmap()`.
+"""
+function plot_unnormalised(f, z_grid, vpa_grid, density, upar, vth, evolve_density,
+                           evolve_upar, evolve_ppar; plot_log=false)
+
+    if backend_name() != :pyplot
+        error("PyPlot backend is required for plot_unnormalised(). Call pyplot() "
+              * "first.")
+    end
+
+    nvpa, nz = size(f)
+    z2d = zeros(nvpa, nz)
+    dzdt2d = zeros(nvpa, nz)
+    for iz ∈ 1:nz
+        @views z2d[:,iz] .= z_grid[iz]
+        @views dzdt2d[:,iz] .= vpagrid_to_dzdt(vpa_grid, vth[iz], upar[iz], evolve_ppar,
+                                             evolve_upar)
+    end
+
+    if evolve_ppar
+        f_unnorm = similar(f)
+        for iz ∈ 1:nz, ivpa ∈ 1:nvpa
+            f_unnorm[ivpa,iz] = @. f[ivpa,iz] * density[iz] / vth[iz]
+        end
+    elseif evolve_density
+        f_unnorm = similar(f)
+        for iz ∈ 1:nz, ivpa ∈ 1:nvpa
+            f_unnorm[ivpa,iz] = @. f[ivpa,iz] * density[iz]
+        end
+    else
+        f_unnorm = f
+    end
+
+    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
+
+    return heatmap(z2d, dzdt2d, f_unnorm, xlabel="z", ylabel="vpa", c=cmap)
+end
+
 end
diff --git a/src/post_processing_input.jl b/src/post_processing_input.jl
index 328640f4dd..ac22ff260c 100644
--- a/src/post_processing_input.jl
+++ b/src/post_processing_input.jl
@@ -47,6 +47,9 @@ const animate_vth_vs_z = false
 const animate_qpar_vs_z = false
 # if animate_f_vs_vpa_z = true, create animation of f(vpa,z) at different time slices
 const animate_f_vs_vpa_z = false
+# if animate_f_unnormalized = true, create animation of f_unnorm(v_parallel_unnorm,z) at
+# different time slices
+const animate_f_unnormalized = false
 # if animate_deltaf_vs_vpa_z = true, create animation of δf(vpa,z) at different time slices
 const animate_deltaf_vs_vpa_z = false
 # if animate_f_vs_vpa_z0 = true, create animation of f(vpa0,z) at different time slices
@@ -78,8 +81,8 @@ pp = pp_input(calculate_frequencies, plot_phi0_vs_t, plot_phi_vs_z_t,
     animate_phi_vs_z, plot_dens0_vs_t, plot_upar0_vs_t, plot_ppar0_vs_t, plot_vth0_vs_t, plot_qpar0_vs_t,
     plot_dens_vs_z_t, plot_upar_vs_z_t, plot_ppar_vs_z_t, plot_qpar_vs_z_t,
     animate_dens_vs_z, animate_upar_vs_z, animate_ppar_vs_z, animate_vth_vs_z, animate_qpar_vs_z,
-    animate_f_vs_vpa_z, animate_f_vs_vpa_z0, animate_f_vs_z0_vpa,
-    animate_deltaf_vs_vpa_z, animate_deltaf_vs_vpa_z0, animate_deltaf_vs_vpa_z0,
-    nwrite_movie, itime_min, itime_max, ivpa0, iz0, ir0)
+    animate_f_vs_vpa_z, animate_f_unnormalized, animate_f_vs_vpa_z0,
+    animate_f_vs_z0_vpa, animate_deltaf_vs_vpa_z, animate_deltaf_vs_vpa_z0,
+    animate_deltaf_vs_vpa_z0, nwrite_movie, itime_min, itime_max, ivpa0, iz0, ir0)
 
 end

From f9ac1ab8435c94a76ce9c44563e3d1c5e5ef3aa1 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 31 Aug 2022 22:18:11 +0100
Subject: [PATCH 08/17] Update post_processing.jl so multiple sims can be
 plotted together

---
 src/post_processing.jl | 626 ++++++++++++++++++++++++++++++-----------
 1 file changed, 459 insertions(+), 167 deletions(-)

diff --git a/src/post_processing.jl b/src/post_processing.jl
index e56cd712fa..58b24ac330 100644
--- a/src/post_processing.jl
+++ b/src/post_processing.jl
@@ -24,6 +24,8 @@ using ..load_data: load_coordinate_data, load_fields_data, load_moments_data, lo
 using ..analysis: analyze_fields_data, analyze_moments_data, analyze_pdf_data
 using ..velocity_moments: integrate_over_vspace
 
+const default_compare_prefix = "comparison_plots/compare"
+
 """
 Calculate a moving average
 
@@ -50,44 +52,97 @@ function moving_average(v::AbstractVector, n::mk_int)
 end
 
 """
+    get_tuple_of_return_values(func, arg_tuples...)
+
+Suppose `func(args...)` returns a tuple of return values, then
+`get_tuple_of_return_values(func, arg_tuples...)` returns a tuple (with an entry for each
+return value of `func`) of tuples (one for each argument in each of `arg_tuples...`) of
+return values.
+"""
+function get_tuple_of_return_values(func, arg_tuples...)
+
+    if isempty(arg_tuples)
+        return ()
+    end
+    n_args_tuple = Tuple(length(a) for a ∈ arg_tuples)
+    if !all(n==n_args_tuple[1] for n ∈ n_args_tuple)
+        error("All argument tuples passed to `get_tuple_of_return_values()` must have "
+              * "the same length")
+    end
+    n_args = n_args_tuple[1]
+
+    collected_args = Tuple(Tuple(a[i] for a ∈ arg_tuples) for i ∈ 1:n_args)
+
+    wrong_way_tuple = Tuple(func(args...) for args ∈ collected_args)
+
+    n_return_values = length(wrong_way_tuple[1])
+
+    return Tuple(Tuple(wrong_way_tuple[i][j] for i ∈ 1:n_args)
+                 for j ∈ 1:n_return_values)
+end
+
+"""
+    analyze_and_plot_data(path...)
+
+Make some plots for the simulation at `path`. If more than one argument is passed to
+`path`, plot all the simulations together.
+
+If a single value is passed for `path` the plots/movies are created in the same
+directory as the run, and given names based on the name of the run. If multiple values
+are passed, the plots/movies are given names beginning with `compare_` and are created
+in the current directory.
 """
-function analyze_and_plot_data(path)
+function analyze_and_plot_data(path...)
     # Create run_name from the path to the run directory
-    path = realpath(path)
-    if isdir(path)
-        run_name = joinpath(path, basename(path))
-    else
-        run_name = splitext(path)[1]
+    run_names = Vector{String}(undef,0)
+    for p ∈ path
+        p = realpath(p)
+        if isdir(p)
+            push!(run_names, joinpath(p, basename(p)))
+        else
+            push!(run_names, splitext(p)[1])
+        end
     end
+
     # open the netcdf file and give it the handle 'fid'
-    fid = open_netcdf_file(run_name)
+    nc_files = Tuple(open_netcdf_file(x) for x ∈ run_names)
+
+    # Truncate to just the file names to make better titles
+    run_labels = Tuple(basename(x) for x ∈ run_names)
+
     # load space-time coordinate data
-    nvpa, vpa, vpa_wgts, nz, z, z_wgts, Lz,
-     nr, r, r_wgts, Lr, ntime, time = load_coordinate_data(fid)
+    nvpa, vpa, vpa_wgts, nz, z, z_wgts, Lz, nr, r, r_wgts, Lr, ntime, time =
+        get_tuple_of_return_values(load_coordinate_data, nc_files)
+
     # initialise the post-processing input options
-    nwrite_movie, itime_min, itime_max, ivpa0, iz0, ir0 = init_postprocessing_options(pp, nvpa, nz, nr, ntime)
+    nwrite_movie, itime_min, itime_max, ivpa0, iz0, ir0 =
+        init_postprocessing_options(pp, minimum(nvpa), minimum(nz), minimum(nr),
+                                    minimum(ntime))
     # load full (z,r,t) fields data
-    phi = load_fields_data(fid)
+    phi = Tuple(load_fields_data(f)[:,ir0,:] for f ∈ nc_files)
+
     # load full (z,r,species,t) velocity moments data
     density, parallel_flow, parallel_pressure, parallel_heat_flux,
         thermal_speed, n_species, evolve_density, evolve_upar, evolve_ppar =
-        load_moments_data(fid)
+        get_tuple_of_return_values(load_moments_data, nc_files)
+    density = Tuple(n[:,ir0,:,:] for n ∈ density)
+    parallel_flow = Tuple(upar[:,ir0,:,:] for upar ∈ parallel_flow)
+    parallel_pressure = Tuple(ppar[:,ir0,:,:] for ppar ∈ parallel_pressure)
+    parallel_heat_flux = Tuple(qpar[:,ir0,:,:] for qpar ∈ parallel_heat_flux)
+    thermal_speed = Tuple(vth[:,ir0,:,:] for vth ∈ thermal_speed)
+
     # load full (vpa,z,r,species,t) particle distribution function (pdf) data
-    ff = load_pdf_data(fid)
+    ff = Tuple(load_pdf_data(f)[:,:,ir0,:,:] for f ∈ nc_files)
 
     #evaluate 1D-1V diagnostics at fixed ir0
-    plot_1D_1V_diagnostics(run_name, fid, nwrite_movie, itime_min, itime_max, ivpa0, iz0, ir0, r,
-        phi[:,ir0,:],
-        density[:,ir0,:,:],
-        parallel_flow[:,ir0,:,:],
-        parallel_pressure[:,ir0,:,:],
-        parallel_heat_flux[:,ir0,:,:],
-        thermal_speed[:,ir0,:,:],
-        ff[:,:,ir0,:,:],
-        n_species, evolve_density, evolve_upar, evolve_ppar, nvpa, vpa, vpa_wgts,
-        nz, z, z_wgts, Lz, ntime, time)
-
-    close(fid)
+    plot_1D_1V_diagnostics(run_names, run_labels, nc_files, nwrite_movie, itime_min,
+        itime_max, ivpa0, iz0, ir0, r, phi, density, parallel_flow, parallel_pressure,
+        parallel_heat_flux, thermal_speed, ff, n_species, evolve_density, evolve_upar,
+        evolve_ppar, nvpa, vpa, vpa_wgts, nz, z, z_wgts, Lz, ntime, time)
+
+    for f ∈ nc_files
+        close(f)
+    end
 
 end
 
@@ -137,74 +192,102 @@ end
 
 """
 """
-function plot_1D_1V_diagnostics(run_name, fid, nwrite_movie, itime_min, itime_max, ivpa0, iz0, ir0, r,
-    phi, density, parallel_flow, parallel_pressure, parallel_heat_flux,
-    thermal_speed, ff, n_species, evolve_density, evolve_upar, evolve_ppar, nvpa, vpa,
-    vpa_wgts, nz, z, z_wgts, Lz, ntime, time)
+function plot_1D_1V_diagnostics(run_names, run_labels, nc_files, nwrite_movie,
+        itime_min, itime_max, ivpa0, iz0, ir0, r, phi, density, parallel_flow,
+        parallel_pressure, parallel_heat_flux, thermal_speed, ff, n_species,
+        evolve_density, evolve_upar, evolve_ppar, nvpa, vpa, vpa_wgts, nz, z, z_wgts,
+        Lz, ntime, time)
+
+    n_runs = length(run_names)
 
     # plot_unnormalised() requires PyPlot, so ensure it is used for all plots for
     # consistency
     pyplot()
 
     # analyze the fields data
-    phi_fldline_avg, delta_phi = analyze_fields_data(phi, ntime, nz, z_wgts, Lz)
+    phi_fldline_avg, delta_phi = get_tuple_of_return_values(analyze_fields_data, phi,
+                                                            ntime, nz, z_wgts, Lz)
+
+    function as_tuple(x)
+        return Tuple(x for _ ∈ 1:n_runs)
+    end
+
     # use a fit to calculate and write to file the damping rate and growth rate of the
     # perturbed electrostatic potential
     frequency, growth_rate, shifted_time, fitted_delta_phi =
-        calculate_and_write_frequencies(fid, run_name, ntime, time, z, itime_min,
-                                        itime_max, iz0, delta_phi, pp)
+        get_tuple_of_return_values(calculate_and_write_frequencies, nc_files, run_names,
+            ntime, time, z, as_tuple(itime_min), as_tuple(itime_max), as_tuple(iz0),
+            delta_phi, as_tuple(pp))
     # create the requested plots of the fields
     plot_fields(phi, delta_phi, time, itime_min, itime_max, nwrite_movie,
-                z, iz0, run_name, fitted_delta_phi, pp)
+                z, iz0, run_names, run_labels, fitted_delta_phi, pp)
     # load velocity moments data
     # analyze the velocity moments data
     density_fldline_avg, upar_fldline_avg, ppar_fldline_avg, vth_fldline_avg, qpar_fldline_avg,
         delta_density, delta_upar, delta_ppar, delta_vth, delta_qpar =
-        analyze_moments_data(density, parallel_flow, parallel_pressure, thermal_speed,
-                             parallel_heat_flux, ntime, n_species, nz, z_wgts, Lz)
+        get_tuple_of_return_values(analyze_moments_data, density, parallel_flow,
+            parallel_pressure, thermal_speed, parallel_heat_flux, ntime, n_species, nz,
+            z_wgts, Lz)
     # create the requested plots of the moments
     plot_moments(density, delta_density, density_fldline_avg,
         parallel_flow, delta_upar, upar_fldline_avg,
         parallel_pressure, delta_ppar, ppar_fldline_avg,
         thermal_speed, delta_vth, vth_fldline_avg,
         parallel_heat_flux, delta_qpar, qpar_fldline_avg,
-        pp, run_name, time, itime_min, itime_max,
+        pp, run_names, run_labels, time, itime_min, itime_max,
         nwrite_movie, z, iz0, n_species)
     # load particle distribution function (pdf) data
     # analyze the pdf data
     f_fldline_avg, delta_f, dens_moment, upar_moment, ppar_moment =
-        analyze_pdf_data(ff, vpa, nvpa, nz, n_species, ntime, vpa_wgts, z_wgts,
-                         Lz, thermal_speed, evolve_ppar)
+        get_tuple_of_return_values(analyze_pdf_data, ff, vpa, nvpa, nz, n_species,
+            ntime, vpa_wgts, z_wgts, Lz, thermal_speed, evolve_ppar)
 
     println("Plotting distribution function data...")
     cmlog(cmlin::ColorGradient) = RGB[cmlin[x] for x=LinRange(0,1,30)]
     logdeep = cgrad(:deep, scale=:log) |> cmlog
-    for is ∈ 1:n_species
-        if n_species > 1
+    n_species_max = maximum(n_species)
+    if n_runs == 1
+        prefix = run_names[1]
+        legend = false
+    else
+        prefix = default_compare_prefix
+        legend = true
+    end
+    for is ∈ 1:n_species_max
+        if n_species_max > 1
             spec_string = string("_spec", string(is))
         else
             spec_string = ""
         end
         # plot difference between evolved density and ∫dvpa f; only possibly different if density removed from
         # normalised distribution function at run-time
-        @views plot(time, density[iz0,is,:] .- dens_moment[iz0,is,:])
-        outfile = string(run_name, "_intf0_vs_t", spec_string, ".pdf")
+        plot(legend=legend)
+        for (t, n, n_int, run_label) ∈ zip(time, density, dens_moment, run_labels)
+            @views plot!(t, n[iz0,is,:] .- n_int[iz0,is,:], label=run_label)
+        end
+        outfile = string(prefix, "_intf0_vs_t", spec_string, ".pdf")
         savefig(outfile)
         # if evolve_upar = true, plot ∫dwpa wpa * f, which should equal zero
         # otherwise, this plots ∫dvpa vpa * f, which is dens*upar
-        intwf0_max = maximum(abs.(upar_moment[iz0,is,:]))
-        if intwf0_max < 1.0e-15
-            @views plot(time, upar_moment[iz0,is,:], ylims = (-1.0e-15, 1.0e-15))
-        else
-            @views plot(time, upar_moment[iz0,is,:])
+        plot(legend=legend)
+        for (t, upar_int, run_label) ∈ zip(time, upar_moment, run_labels)
+            intwf0_max = maximum(abs.(upar_int[iz0,is,:]))
+            if intwf0_max < 1.0e-15
+                @views plot!(t, upar_int[iz0,is,:], ylims = (-1.0e-15, 1.0e-15), label=run_label)
+            else
+                @views plot!(t, upar_int[iz0,is,:], label=run_label)
+            end
         end
-        outfile = string(run_name, "_intwf0_vs_t", spec_string, ".pdf")
+        outfile = string(prefix, "_intwf0_vs_t", spec_string, ".pdf")
         savefig(outfile)
         # plot difference between evolved parallel pressure and ∫dvpa vpa^2 f;
         # only possibly different if density and thermal speed removed from
         # normalised distribution function at run-time
-        @views plot(time, parallel_pressure[iz0,is,:] .- ppar_moment[iz0,is,:])
-        outfile = string(run_name, "_intw2f0_vs_t", spec_string, ".pdf")
+        plot(legend=legend)
+        for (t, ppar, ppar_int, run_label) ∈ zip(time, parallel_pressure, ppar_moment, run_labels)
+            @views plot(t, ppar[iz0,is,:] .- ppar_int[iz0,is,:], label=run_label)
+        end
+        outfile = string(prefix, "_intw2f0_vs_t", spec_string, ".pdf")
         savefig(outfile)
         #fmin = minimum(ff[:,:,is,:])
         #fmax = maximum(ff[:,:,is,:])
@@ -212,64 +295,128 @@ function plot_1D_1V_diagnostics(run_name, fid, nwrite_movie, itime_min, itime_ma
             # make a gif animation of ln f(vpa,z,t)
             anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
                 #heatmap(z, vpa, log.(abs.(ff[:,:,i])), xlabel="z", ylabel="vpa", clims = (fmin,fmax), c = :deep)
-                @views heatmap(z, vpa, log.(abs.(ff[:,:,is,i])), xlabel="z", ylabel="vpa", fillcolor = logdeep)
+                subplots = (@views heatmap(z, vpa, log.(abs.(f[:,:,is,i])), xlabel="z",
+                                           ylabel="vpa", fillcolor = logdeep, title=run_label)
+                            for (f, this_z, this_vpa, run_label) ∈ zip(ff, z, vpa, run_labels))
+                plot(subplots..., layout=(1,n_runs), size=(400*n_runs, 400))
             end
-            outfile = string(run_name, "_logf_vs_vpa_z", spec_string, ".gif")
+            outfile = string(prefix, "_logf_vs_vpa_z", spec_string, ".gif")
             gif(anim, outfile, fps=5)
             # make a gif animation of f(vpa,z,t)
             anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
                 #heatmap(z, vpa, log.(abs.(ff[:,:,i])), xlabel="z", ylabel="vpa", clims = (fmin,fmax), c = :deep)
-                @views heatmap(z, vpa, ff[:,:,is,i], xlabel="z", ylabel="vpa", c = :deep, interpolation = :cubic)
+                subplots = (@views heatmap(this_z, this_vpa, f[:,:,is,i], xlabel="z",
+                                           ylabel="vpa", c = :deep, interpolation =
+                                           :cubic, title=run_label)
+                            for (f, this_z, this_vpa, run_label) ∈ zip(ff, z, vpa, run_labels))
+                plot(subplots..., layout=(1,n_runs), size=(400*n_runs, 400))
             end
-            outfile = string(run_name, "_f_vs_vpa_z", spec_string, ".gif")
+            outfile = string(prefix, "_f_vs_vpa_z", spec_string, ".gif")
             gif(anim, outfile, fps=5)
             # make pdf of f(vpa,z,t_final) for each species
             str = string("spec ", string(is), " pdf")
-            @views heatmap(z, vpa, ff[:,:,is,end], xlabel="vpa", ylabel="z", c = :deep, interpolation = :cubic, title=str)
-            outfile = string(run_name, "_f_vs_z_vpa_final", spec_string, ".pdf")
+            subplots = (@views heatmap(this_z, this_vpa, f[:,:,is,end], xlabel="vpa", ylabel="z",
+                                       c = :deep, interpolation = :cubic,
+                                       title=string(run_label, str))
+                        for (f, this_z, this_vpa, run_label) ∈ zip(ff, z, vpa, run_labels))
+            plot(subplots..., layout=(1,n_runs), size=(400*n_runs, 400))
+            outfile = string(prefix, "_f_vs_z_vpa_final", spec_string, ".pdf")
             savefig(outfile)
+
+            anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
+                plot(legend=legend)
+                for (f, this_vpa, run_label) ∈ zip(ff, vpa, run_labels)
+                    @views plot!(this_vpa, f[:,1,is,i], xlabel="vpa", ylabel="f(z=0)",
+                                 label=run_label)
+                end
+            end
+            outfile = string(prefix, "_f0_vs_vpa", spec_string, ".gif")
+            gif(anim, outfile, fps=5)
+
+            anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
+                plot(legend=legend)
+                for (f, this_vpa, run_label) ∈ zip(ff, vpa, run_labels)
+                    @views plot!(this_vpa, f[:,end,is,i], xlabel="vpa", ylabel="f(z=L)",
+                                 label=run_label)
+                end
+            end
+            outfile = string(prefix, "_fL_vs_vpa", spec_string, ".gif")
+            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
-                @views plot_unnormalised(ff[:,:,is,i], z, vpa, density[:,is,i],
-                                         parallel_flow[:,is,i], thermal_speed[:,is,i],
-                                         evolve_density, evolve_upar, evolve_ppar)
+                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(run_name, "_f_unnorm_vs_vpa_z", spec_string, ".gif")
+            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)
         end
         if pp.animate_deltaf_vs_vpa_z
             # make a gif animation of δf(vpa,z,t)
             anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
-                @views heatmap(z, vpa, delta_f[:,:,is,i], xlabel="z", ylabel="vpa", c = :deep, interpolation = :cubic)
+                subplots = (@views heatmap(this_z, this_vpa, delta_f[:,:,is,i],
+                                       xlabel="z", ylabel="vpa", c = :deep,
+                                       interpolation = :cubic, title=run_label)
+                            for (df, this_z, this_vpa, run_label) ∈
+                                zip(delta_f, z, vpa, run_labels))
+                plot(subplots..., layout=(1,n_runs), size=(400*n_runs, 400))
             end
-            outfile = string(run_name, "_deltaf_vs_vpa_z", spec_string, ".gif")
+            outfile = string(prefix, "_deltaf_vs_vpa_z", spec_string, ".gif")
             gif(anim, outfile, fps=5)
         end
         if pp.animate_f_vs_vpa_z0
-            fmin = minimum(ff[ivpa0,:,is,:])
-            fmax = maximum(ff[ivpa0,:,is,:])
+            fmin = minimum(minimum(f[ivpa0,:,is,:]) for f ∈ ff)
+            fmax = maximum(maximum(f[ivpa0,:,is,:]) for f ∈ ff)
             # make a gif animation of f(vpa0,z,t)
             anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
-                @views plot(z, ff[ivpa0,:,is,i], ylims = (fmin,fmax))
+                plot(legend=legend)
+                for (f, this_z, run_label) ∈ zip(ff, z, run_labels)
+                    @views plot!(this_z, f[ivpa0,:,is,i], ylims = (fmin,fmax),
+                                 label=run_label)
+                end
             end
-            outfile = string(run_name, "_f_vs_z", spec_string, ".gif")
+            outfile = string(prefix, "_f_vs_z", spec_string, ".gif")
             gif(anim, outfile, fps=5)
         end
         if pp.animate_deltaf_vs_vpa_z0
-            fmin = minimum(delta_f[ivpa0,:,is,:])
-            fmax = maximum(delta_f[ivpa0,:,is,:])
+            fmin = minimum(minimum(df[ivpa0,:,is,:]) for df ∈ delta_f)
+            fmax = maximum(maximum(df[ivpa0,:,is,:]) for df ∈ delta_f)
             # make a gif animation of f(vpa0,z,t)
             anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
-                @views plot(z, delta_f[ivpa0,:,is,i], ylims = (fmin,fmax))
+                plot(legend=legend)
+                for (df, this_z, run_label) ∈ zip(delta_f, z, run_labels)
+                    @views plot!(this_z, df[ivpa0,:,is,i], ylims = (fmin,fmax),
+                                 label=run_label)
+                end
             end
-            outfile = string(run_name, "_deltaf_vs_z", spec_string, ".gif")
+            outfile = string(prefix, "_deltaf_vs_z", spec_string, ".gif")
             gif(anim, outfile, fps=5)
         end
         if pp.animate_f_vs_vpa_z0
-            fmin = minimum(ff[:,iz0,is,:])
-            fmax = maximum(ff[:,iz0,is,:])
+            fmin = minimum(minimum(f[:,iz0,is,:]) for f ∈ ff)
+            fmax = maximum(maximum(f[:,iz0,is,:]) for f ∈ ff)
 
             # if is == 1
             #     tmp = copy(ff)
@@ -280,7 +427,7 @@ function plot_1D_1V_diagnostics(run_name, fid, nwrite_movie, itime_min, itime_ma
             #     end
             #     plot(time, bohm_integral, xlabel="time", label="Bohm integral")
             #     plot!(time, density[1,1,:], label="nᵢ(zmin)")
-            #     outfile = string(run_name, "_Bohm_criterion.pdf")
+            #     outfile = string(prefix, "_Bohm_criterion.pdf")
             #     savefig(outfile)
             #     println()
             #     if bohm_integral[end] <= density[1,1,end]
@@ -296,19 +443,27 @@ function plot_1D_1V_diagnostics(run_name, fid, nwrite_movie, itime_min, itime_ma
             # make a gif animation of f(vpa,z0,t)
             anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
                 #@views plot(vpa, ff[iz0,:,is,i], ylims = (fmin,fmax))
-                @views plot(vpa, ff[:,iz0,is,i])
+                plot(legend=legend)
+                for (f, this_vpa, run_label) ∈ zip(ff, vpa, run_labels)
+                    @views plot!(this_vpa, f[:,iz0,is,i], label=run_label)
+                end
             end
-            outfile = string(run_name, "_f_vs_vpa", spec_string, ".gif")
+            outfile = string(prefix, "_f_vs_vpa", spec_string, ".gif")
             gif(anim, outfile, fps=5)
         end
         if pp.animate_deltaf_vs_vpa_z0
-            fmin = minimum(delta_f[:,iz0,is,:])
-            fmax = maximum(delta_f[:,iz0,is,:])
+            fmin = minimum(minimum(df[:,iz0,is,:]) for df ∈ delta_f)
+            fmax = maximum(maximum(df[:,iz0,is,:]) for df ∈ delta_f)
             # make a gif animation of f(vpa,z0,t)
             anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
-                @views plot(vpa, delta_f[:,iz0,is,i], ylims = (fmin,fmax))
+                plot(legend=legend)
+                for (df, this_vpa, fn, fx, run_label) ∈
+                        zip(delta_f, vpa, fmin, fmax, run_labels)
+                    @views plot!(this_vpa, delta_f[:,iz0,is,i], ylims = (fn,fx),
+                                label=run_label)
+                end
             end
-            outfile = string(run_name, "_deltaf_vs_vpa", spec_string, ".gif")
+            outfile = string(prefix, "_deltaf_vs_vpa", spec_string, ".gif")
             gif(anim, outfile, fps=5)
         end
     end
@@ -400,47 +555,76 @@ end
 """
 """
 function plot_fields(phi, delta_phi, time, itime_min, itime_max, nwrite_movie,
-    z, iz0, run_name, fitted_delta_phi, pp)
+    z, iz0, run_names, run_labels, fitted_delta_phi, pp)
 
     println("Plotting fields data...")
-    phimin = minimum(phi)
-    phimax = maximum(phi)
+
+    n_runs = length(run_names)
+    if n_runs == 1
+        prefix = run_names[1]
+        legend = false
+    else
+        prefix = default_compare_prefix
+        legend = true
+    end
+
+    phimin = minimum(minimum(p) for p ∈ phi)
+    phimax = maximum(maximum(p) for p ∈ phi)
+
     if pp.plot_phi0_vs_t
         # plot the time trace of phi(z=z0)
         #plot(time, log.(phi[i,:]), yscale = :log10)
-        @views plot(time, phi[iz0,:])
-        outfile = string(run_name, "_phi0_vs_t.pdf")
+        plot(legend=legend)
+        for (t, p, run_label) ∈ zip(time, phi, run_labels)
+            @views plot!(t, p[iz0,:], label=run_label)
+        end
+        outfile = string(prefix, "_phi0_vs_t.pdf")
         savefig(outfile)
-        # plot the time trace of phi(z=z0)-phi_fldline_avg
-        @views plot(time, abs.(delta_phi[iz0,:]), xlabel="t*Lz/vti", ylabel="δϕ", yaxis=:log)
-        if pp.calculate_frequencies
-            plot!(time, abs.(fitted_delta_phi))
+        plot(legend=legend)
+        for (t, dp, fit, run_label) ∈ zip(time, delta_phi, fitted_delta_phi, run_labels)
+            # plot the time trace of phi(z=z0)-phi_fldline_avg
+            @views plot!(t, abs.(dp[iz0,:]), xlabel="t*Lz/vti", ylabel="δϕ", yaxis=:log,
+                         label="$run_label δϕ")
+            if pp.calculate_frequencies
+                plot!(t, abs.(fit), linestyle=:dash, label="$run_label fit")
+            end
         end
-        outfile = string(run_name, "_delta_phi0_vs_t.pdf")
+        outfile = string(prefix, "_delta_phi0_vs_t.pdf")
         savefig(outfile)
     end
     if pp.plot_phi_vs_z_t
         # make a heatmap plot of ϕ(z,t)
-        heatmap(time, z, phi, xlabel="time", ylabel="z", title="ϕ", c = :deep)
-        outfile = string(run_name, "_phi_vs_z_t.pdf")
+        subplots = (heatmap(t, this_z, p, xlabel="time", ylabel="z", title=run_label, c =
+                            :deep)
+                    for (t, this_z, p, run_label) ∈ zip(time, z, phi, run_labels))
+        plot(subplots..., layout=(1,n_runs), size=(400*n_runs, 400))
+        outfile = string(prefix, "_phi_vs_z_t.pdf")
         savefig(outfile)
     end
     if pp.animate_phi_vs_z
         # make a gif animation of ϕ(z) at different times
         anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
-            @views plot(z, phi[:,i], xlabel="z", ylabel="ϕ", ylims = (phimin,phimax))
+            plot(legend=legend)
+            for (t, this_z, p, run_label) ∈ zip(time, z, phi, run_labels)
+                @views plot!(this_z, p[:,i], xlabel="z", ylabel="ϕ",
+                             ylims=(phimin, phimax), label=run_label)
+            end
         end
-        outfile = string(run_name, "_phi_vs_z.gif")
+        outfile = string(prefix, "_phi_vs_z.gif")
         gif(anim, outfile, fps=5)
     end
     # nz = length(z)
     # izmid = cld(nz,2)
     # plot(z[izmid:end], phi[izmid:end,end] .- phi[izmid,end], xlabel="z/Lz - 1/2", ylabel="eϕ/Te", label = "data", linewidth=2)
     # plot!(exp.(-(phi[cld(nz,2),end] .- phi[izmid:end,end])) .* erfi.(sqrt.(abs.(phi[cld(nz,2),end] .- phi[izmid:end,end])))/sqrt(pi)/0.688, phi[izmid:end,end] .- phi[izmid,end], label = "analytical", linewidth=2)
-    # outfile = string(run_name, "_harrison_comparison.pdf")
+    # outfile = string(prefix, "_harrison_comparison.pdf")
     # savefig(outfile)
-    plot(z, phi[:,end], xlabel="z/Lz", ylabel="eϕ/Te", label="", linewidth=2)
-    outfile = string(run_name, "_phi_final.pdf")
+    plot(legend=legend)
+    for (t, this_z, p, run_label) ∈ zip(time, z, phi, run_labels)
+        plot!(this_z, p[:,end], xlabel="z/Lz", ylabel="eϕ/Te", label=run_label,
+             linewidth=2)
+    end
+    outfile = string(prefix, "_phi_final.pdf")
     savefig(outfile)
 
     println("done.")
@@ -453,166 +637,274 @@ function plot_moments(density, delta_density, density_fldline_avg,
     parallel_pressure, delta_ppar, ppar_fldline_avg,
     thermal_speed, delta_vth, vth_fldline_avg,
     parallel_heat_flux, delta_qpar, qpar_fldline_avg,
-    pp, run_name, time, itime_min, itime_max, nwrite_movie,
+    pp, run_names, run_labels, time, itime_min, itime_max, nwrite_movie,
     z, iz0, n_species)
+
     println("Plotting velocity moments data...")
+
+    n_runs = length(run_names)
+    if n_runs == 1
+        prefix = run_names[1]
+        legend = false
+    else
+        prefix = default_compare_prefix
+        legend = true
+    end
+
     # plot the species-summed, field-line averaged vs time
-    denstot = copy(density_fldline_avg)
-    denstot .= sum(density_fldline_avg,dims=1)
-    @. denstot /= denstot[1,1]
-    denstot_min = minimum(denstot[1,:]) - 0.1
-    denstot_max = maximum(denstot[1,:]) + 0.1
-    @views plot(time, denstot[1,:], ylims=(denstot_min,denstot_max), xlabel="time", ylabel="∑ⱼn̅ⱼ(t)/∑ⱼn̅ⱼ(0)", label="", linewidth=2)
-    outfile = string(run_name, "_denstot_vs_t.pdf")
+    denstot = Tuple(sum(n_fldline_avg,dims=1)[1,:]
+                    for n_fldline_avg ∈ density_fldline_avg)
+    for d in denstot
+        d ./= d[1]
+    end
+    denstot_min = minimum(minimum(dtot) for dtot in denstot) - 0.1
+    denstot_max = maximum(maximum(dtot) for dtot in denstot) + 0.1
+    plot(legend=legend)
+    for (t, dtot, run_label) ∈ zip(time, denstot, run_labels)
+        @views plot!(t, dtot[1,:], ylims=(denstot_min,denstot_max), xlabel="time",
+                     ylabel="∑ⱼn̅ⱼ(t)/∑ⱼn̅ⱼ(0)", linewidth=2, label=run_label)
+    end
+    outfile = string(prefix, "_denstot_vs_t.pdf")
     savefig(outfile)
-    for is ∈ 1:n_species
+    for is ∈ 1:maximum(n_species)
         spec_string = string(is)
-        dens_min = minimum(density[:,is,:])
-        dens_max = maximum(density[:,is,:])
+        dens_min = minimum(minimum(n[:,is,:]) for n ∈ density)
+        dens_max = maximum(maximum(n[:,is,:]) for n ∈ density)
         if pp.plot_dens0_vs_t
             # plot the time trace of n_s(z=z0)
-            @views plot(time, density[iz0,is,:])
-            outfile = string(run_name, "_dens0_vs_t_spec", spec_string, ".pdf")
+            plot(legend=legend)
+            for (t, n, run_label) ∈ zip(time, density, run_labels)
+                @views plot!(t, n[iz0,is,:], label=run_label)
+            end
+            outfile = string(prefix, "_dens0_vs_t_spec", spec_string, ".pdf")
             savefig(outfile)
             # plot the time trace of n_s(z=z0)-density_fldline_avg
-            @views plot(time, abs.(delta_density[iz0,is,:]), yaxis=:log)
-            outfile = string(run_name, "_delta_dens0_vs_t_spec", spec_string, ".pdf")
+            plot(legend=legend)
+            for (t, dn, run_label) ∈ zip(time, delta_density, run_labels)
+                @views plot!(t, abs.(dn[iz0,is,:]), yaxis=:log, label=run_label)
+            end
+            outfile = string(prefix, "_delta_dens0_vs_t_spec", spec_string, ".pdf")
             savefig(outfile)
             # plot the time trace of density_fldline_avg
-            @views plot(time, density_fldline_avg[is,:], xlabel="time", ylabel="<ns/Nₑ>", ylims=(dens_min,dens_max))
-            outfile = string(run_name, "_fldline_avg_dens_vs_t_spec", spec_string, ".pdf")
+            plot(legend=legend)
+            for (t, n_avg, run_label) ∈ zip(time, density_fldline_avg, run_labels)
+                @views plot!(t, n_avg[is,:], xlabel="time", ylabel="<ns/Nₑ>",
+                             ylims=(dens_min,dens_max), label=run_label)
+            end
+            outfile = string(prefix, "_fldline_avg_dens_vs_t_spec", spec_string, ".pdf")
             savefig(outfile)
             # plot the deviation from conservation of density_fldline_avg
-            @views plot(time, density_fldline_avg[is,:] .- density_fldline_avg[is,1], xlabel="time", ylabel="<(ns-ns(0))/Nₑ>")
-            outfile = string(run_name, "_conservation_dens_spec", spec_string, ".pdf")
+            plot(legend=legend)
+            for (t, n_avg, run_label) ∈ zip(time, density_fldline_avg, run_labels)
+                @views plot!(t, n_avg[is,:] .- n_avg[is,1], xlabel="time",
+                             ylabel="<(ns-ns(0))/Nₑ>", label=run_label)
+            end
+            outfile = string(prefix, "_conservation_dens_spec", spec_string, ".pdf")
             savefig(outfile)
         end
-        upar_min = minimum(parallel_flow[:,is,:])
-        upar_max = maximum(parallel_flow[:,is,:])
+        upar_min = minimum(minimum(upar[:,is,:]) for upar ∈ parallel_flow)
+        upar_max = maximum(maximum(upar[:,is,:]) for upar ∈ parallel_flow)
         if pp.plot_upar0_vs_t
             # plot the time trace of n_s(z=z0)
-            @views plot(time, parallel_flow[iz0,is,:])
-            outfile = string(run_name, "_upar0_vs_t_spec", spec_string, ".pdf")
+            plot(legend=legend)
+            for (t, upar, run_label) ∈ zip(time, parallel_flow, run_labels)
+                @views plot!(t, upar[iz0,is,:], label=run_label)
+            end
+            outfile = string(prefix, "_upar0_vs_t_spec", spec_string, ".pdf")
             savefig(outfile)
             # plot the time trace of n_s(z=z0)-density_fldline_avg
-            @views plot(time, abs.(delta_upar[iz0,is,:]), yaxis=:log)
-            outfile = string(run_name, "_delta_upar0_vs_t_spec", spec_string, ".pdf")
+            plot(legend=legend)
+            for (t, dupar, run_label) ∈ zip(time, delta_upar, run_labels)
+                @views plot!(t, abs.(du[iz0,is,:]), yaxis=:log, label=run_label)
+            end
+            outfile = string(prefix, "_delta_upar0_vs_t_spec", spec_string, ".pdf")
             savefig(outfile)
             # plot the time trace of ppar_fldline_avg
-            @views plot(time, upar_fldline_avg[is,:], xlabel="time", ylabel="<upars/sqrt(2Te/ms)>", ylims=(upar_min,upar_max))
-            outfile = string(run_name, "_fldline_avg_upar_vs_t_spec", spec_string, ".pdf")
+            plot(legend=legend)
+            for (t, upar_avg, run_label) ∈ zip(time, upar_fldline_avg, run_labels)
+                @views plot!(t, upar_avg[is,:], xlabel="time",
+                             ylabel="<upars/sqrt(2Te/ms)>", ylims=(upar_min,upar_max),
+                             label=run_label)
+            end
+            outfile = string(prefix, "_fldline_avg_upar_vs_t_spec", spec_string, ".pdf")
             savefig(outfile)
         end
-        ppar_min = minimum(parallel_pressure[:,is,:])
-        ppar_max = maximum(parallel_pressure[:,is,:])
+        ppar_min = minimum(minimum(ppar[:,is,:]) for ppar ∈ parallel_pressure)
+        ppar_max = maximum(maximum(ppar[:,is,:]) for ppar ∈ parallel_pressure)
         if pp.plot_ppar0_vs_t
             # plot the time trace of n_s(z=z0)
-            @views plot(time, parallel_pressure[iz0,is,:])
-            outfile = string(run_name, "_ppar0_vs_t_spec", spec_string, ".pdf")
+            plot(legend=legend)
+            for (t, ppar, run_label) ∈ zip(time, parallel_pressure, run_labels)
+                @views plot!(t, ppar[iz0,is,:], label=run_label)
+            end
+            outfile = string(prefix, "_ppar0_vs_t_spec", spec_string, ".pdf")
             savefig(outfile)
             # plot the time trace of n_s(z=z0)-density_fldline_avg
-            @views plot(time, abs.(delta_ppar[iz0,is,:]), yaxis=:log)
-            outfile = string(run_name, "_delta_ppar0_vs_t_spec", spec_string, ".pdf")
+            plot(legend=legend)
+            for (t, dppar, run_label) ∈ zip(time, delta_ppar, run_labels)
+                @views plot!(t, abs.(dppar[iz0,is,:]), yaxis=:log, label=run_label)
+            end
+            outfile = string(prefix, "_delta_ppar0_vs_t_spec", spec_string, ".pdf")
             savefig(outfile)
             # plot the time trace of ppar_fldline_avg
-            @views plot(time, ppar_fldline_avg[is,:], xlabel="time", ylabel="<ppars/NₑTₑ>", ylims=(ppar_min,ppar_max))
-            outfile = string(run_name, "_fldline_avg_ppar_vs_t_spec", spec_string, ".pdf")
+            plot(legend=legend)
+            for (t, ppar_avg, run_label) ∈ zip(time, ppar_fldline_avg, run_labels)
+                @views plot!(t, ppar_avg[is,:], xlabel="time", ylabel="<ppars/NₑTₑ>",
+                             ylims=(ppar_min,ppar_max), label=run_label)
+            end
+            outfile = string(prefix, "_fldline_avg_ppar_vs_t_spec", spec_string, ".pdf")
             savefig(outfile)
         end
-        vth_min = minimum(thermal_speed[:,is,:])
-        vth_max = maximum(thermal_speed[:,is,:])
+        vth_min = minimum(minimum(vth[:,is,:]) for vth ∈ thermal_speed)
+        vth_max = maximum(maximum(vth[:,is,:]) for vth ∈ thermal_speed)
         if pp.plot_vth0_vs_t
             # plot the time trace of n_s(z=z0)
-            @views plot(time, thermal_speed[iz0,is,:])
-            outfile = string(run_name, "_vth0_vs_t_spec", spec_string, ".pdf")
+            plot(legend=legend)
+            for (t, vth, run_label) ∈ zip(time, thermal_speed, run_labels)
+                @views plot!(t, vth[iz0,is,:], label=run_label)
+            end
+            outfile = string(prefix, "_vth0_vs_t_spec", spec_string, ".pdf")
             savefig(outfile)
             # plot the time trace of n_s(z=z0)-density_fldline_avg
-            @views plot(time, abs.(delta_vth[iz0,is,:]), yaxis=:log)
-            outfile = string(run_name, "_delta_vth0_vs_t_spec", spec_string, ".pdf")
+            plot(legend=legend)
+            for (t, dvth, run_label) ∈ zip(time, delta_vth, run_labels)
+                @views plot!(t, abs.(dvth[iz0,is,:]), yaxis=:log, label=run_label)
+            end
+            outfile = string(prefix, "_delta_vth0_vs_t_spec", spec_string, ".pdf")
             savefig(outfile)
             # plot the time trace of vth_fldline_avg
-            @views plot(time, vth_fldline_avg[is,:], xlabel="time", ylabel="<vths/cₛ₀>", ylims=(vth_min,vth_max))
-            outfile = string(run_name, "_fldline_avg_vth_vs_t_spec", spec_string, ".pdf")
+            plot(legend=legend)
+            for (t, vth_avg, run_label) ∈ zip(time, vth_fldline_avg, run_labels)
+                @views plot!(t, vth_avg[is,:], xlabel="time", ylabel="<vths/cₛ₀>",
+                             ylims=(vth_min,vth_max), label=run_label)
+            end
+            outfile = string(prefix, "_fldline_avg_vth_vs_t_spec", spec_string, ".pdf")
             savefig(outfile)
         end
-        qpar_min = minimum(parallel_heat_flux[:,is,:])
-        qpar_max = maximum(parallel_heat_flux[:,is,:])
+        qpar_min = minimum(minimum(qpar[:,is,:]) for qpar ∈ parallel_heat_flux)
+        qpar_max = maximum(maximum(qpar[:,is,:]) for qpar ∈ parallel_heat_flux)
         if pp.plot_qpar0_vs_t
             # plot the time trace of n_s(z=z0)
-            @views plot(time, parallel_heat_flux[iz0,is,:])
-            outfile = string(run_name, "_qpar0_vs_t_spec", spec_string, ".pdf")
+            plot(legend=legend)
+            for (t, qpar, run_label) ∈ zip(time, parallel_heat_flux, run_labels)
+                @views plot!(t, qpar[iz0,is,:], label=run_label)
+            end
+            outfile = string(prefix, "_qpar0_vs_t_spec", spec_string, ".pdf")
             savefig(outfile)
             # plot the time trace of n_s(z=z0)-density_fldline_avg
-            @views plot(time, abs.(delta_qpar[iz0,is,:]), yaxis=:log)
-            outfile = string(run_name, "_delta_qpar0_vs_t_spec", spec_string, ".pdf")
+            plot(legend=legend)
+            for (t, dqpar, run_label) ∈ zip(time, delta_qpar, run_labels)
+                @views plot!(t, abs.(dqpar[iz0,is,:]), yaxis=:log, label=run_label)
+            end
+            outfile = string(prefix, "_delta_qpar0_vs_t_spec", spec_string, ".pdf")
             savefig(outfile)
             # plot the time trace of ppar_fldline_avg
-            @views plot(time, qpar_fldline_avg[is,:], xlabel="time", ylabel="<qpars/NₑTₑvth>", ylims=(qpar_min,qpar_max))
-            outfile = string(run_name, "_fldline_avg_qpar_vs_t_spec", spec_string, ".pdf")
+            plot(legend=legend)
+            for (t, qpar_avg, run_label) ∈ zip(time, qpar_fldline_avg, run_labels)
+                @views plot!(t, qpar_avg[is,:], xlabel="time", ylabel="<qpars/NₑTₑvth>",
+                             ylims=(qpar_min,qpar_max), label=run_label)
+            end
+            outfile = string(prefix, "_fldline_avg_qpar_vs_t_spec", spec_string, ".pdf")
             savefig(outfile)
         end
         if pp.plot_dens_vs_z_t
             # make a heatmap plot of n_s(z,t)
-            heatmap(time, z, density[:,is,:], xlabel="time", ylabel="z", title="ns/Nₑ", c = :deep)
-            outfile = string(run_name, "_dens_vs_z_t_spec", spec_string, ".pdf")
+            subplots = (heatmap(t, this_z, n[:,is,:], xlabel="time", ylabel="z",
+                                title=run_label, c = :deep)
+                        for (t, this_z, n, run_label) ∈ zip(time, z, density, run_labels))
+            plot(subplots..., layout=(1,n_runs), size=(400*n_runs, 400))
+            outfile = string(prefix, "_dens_vs_z_t_spec", spec_string, ".pdf")
             savefig(outfile)
         end
         if pp.plot_upar_vs_z_t
             # make a heatmap plot of upar_s(z,t)
-            heatmap(time, z, parallel_flow[:,is,:], xlabel="time", ylabel="z", title="upars/vt", c = :deep)
-            outfile = string(run_name, "_upar_vs_z_t_spec", spec_string, ".pdf")
+            subplots = (heatmap(t, this_z, upar[:,is,:], xlabel="time", ylabel="z",
+                                title=run_label, c = :deep)
+                        for (t, this_z, upar, run_label) ∈ zip(time, z, parallel_flow,
+                                                              run_labels))
+            plot(subplots..., layout=(1,n_runs), size=(400*n_runs, 400))
+            outfile = string(prefix, "_upar_vs_z_t_spec", spec_string, ".pdf")
             savefig(outfile)
         end
         if pp.plot_ppar_vs_z_t
             # make a heatmap plot of upar_s(z,t)
-            heatmap(time, z, parallel_pressure[:,is,:], xlabel="time", ylabel="z", title="ppars/NₑTₑ", c = :deep)
-            outfile = string(run_name, "_ppar_vs_z_t_spec", spec_string, ".pdf")
+            subplots = (heatmap(t, this_z, ppar[:,is,:], xlabel="time", ylabel="z",
+                                title=run_label, c = :deep)
+                        for (t, this_z, ppar, run_label) ∈
+                            zip(time, z, parallel_pressure, run_labels))
+            plot(subplots..., layout=(1,n_runs), size=(400*n_runs, 400))
+            outfile = string(prefix, "_ppar_vs_z_t_spec", spec_string, ".pdf")
             savefig(outfile)
         end
         if pp.plot_qpar_vs_z_t
             # make a heatmap plot of upar_s(z,t)
-            heatmap(time, z, parallel_heat_flux[:,is,:], xlabel="time", ylabel="z", title="qpars/NₑTₑvt", c = :deep)
-            outfile = string(run_name, "_qpar_vs_z_t_spec", spec_string, ".pdf")
+            subplots = (heatmap(t, this_z, qpar[:,is,:], xlabel="time", ylabel="z",
+                                title=run_label, c = :deep)
+                        for (t, this_z, qpar, run_label) ∈
+                            zip(time, z, parallel_heat_flux, run_labels))
+            plot(subplots..., layout=(1,n_runs), size=(400*n_runs, 400))
+            outfile = string(prefix, "_qpar_vs_z_t_spec", spec_string, ".pdf")
             savefig(outfile)
         end
         if pp.animate_dens_vs_z
             # make a gif animation of ϕ(z) at different times
             anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
-                @views plot(z, density[:,is,i], xlabel="z", ylabel="nᵢ/Nₑ", ylims = (dens_min,dens_max))
+                plot(legend=legend)
+                for (t, this_z, n, run_label) ∈ zip(time, z, density, run_labels)
+                    @views plot!(this_z, n[:,is,i], xlabel="z", ylabel="nᵢ/Nₑ",
+                                 ylims=(dens_min, dens_max), label=run_label)
+                end
             end
-            outfile = string(run_name, "_dens_vs_z_spec", spec_string, ".gif")
+            outfile = string(prefix, "_dens_vs_z_spec", spec_string, ".gif")
             gif(anim, outfile, fps=5)
         end
         if pp.animate_upar_vs_z
             # make a gif animation of upar(z) at different times
             anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
-                @views plot(z, parallel_flow[:,is,i], xlabel="z", ylabel="upars/vt", ylims = (upar_min,upar_max))
+                plot(legend=legend)
+                for (t, this_z, upar, run_label) ∈ zip(time, z, parallel_flow, run_labels)
+                    @views plot!(this_z, upar[:,is,i], xlabel="z", ylabel="upars/vt",
+                                 ylims=(upar_min, upar_max), label=run_label)
+                end
             end
-            outfile = string(run_name, "_upar_vs_z_spec", spec_string, ".gif")
+            outfile = string(prefix, "_upar_vs_z_spec", spec_string, ".gif")
             gif(anim, outfile, fps=5)
         end
         if pp.animate_ppar_vs_z
             # make a gif animation of ppar(z) at different times
             anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
-                @views plot(z, parallel_pressure[:,is,i], xlabel="z", ylabel="ppars", ylims = (ppar_min,ppar_max))
+                plot(legend=legend)
+                for (t, this_z, ppar, run_label) ∈ zip(time, z, parallel_pressure, run_labels)
+                    @views plot!(this_z, ppar[:,is,i], xlabel="z", ylabel="ppars",
+                                 ylims=(ppar_min, ppar_max), label=run_label)
+                end
             end
-            outfile = string(run_name, "_ppar_vs_z_spec", spec_string, ".gif")
+            outfile = string(prefix, "_ppar_vs_z_spec", spec_string, ".gif")
             gif(anim, outfile, fps=5)
         end
         if pp.animate_vth_vs_z
             # make a gif animation of vth(z) at different times
             anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
-                @views plot(z, thermal_speed[:,is,i], xlabel="z", ylabel="vths", ylims = (vth_min,vth_max))
+                plot(legend=legend)
+                for (t, this_z, vth, run_label) ∈ zip(time, z, thermal_speed, run_labels)
+                    @views plot!(this_z, vth[:,is,i], xlabel="z", ylabel="vths",
+                                 ylims=(vth_min, vth_max), label=run_label)
+                end
             end
-            outfile = string(run_name, "_vth_vs_z_spec", spec_string, ".gif")
+            outfile = string(prefix, "_vth_vs_z_spec", spec_string, ".gif")
             gif(anim, outfile, fps=5)
         end
         if pp.animate_qpar_vs_z
             # make a gif animation of ppar(z) at different times
             anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
-                @views plot(z, parallel_heat_flux[:,is,i], xlabel="z", ylabel="qpars", ylims = (qpar_min,qpar_max))
+                plot(legend=legend)
+                for (t, this_z, qpar, run_label) ∈ zip(time, z, parallel_heat_flux,
+                                                      run_labels)
+                    @views plot!(this_z, qpar[:,is,i], xlabel="z", ylabel="qpars",
+                                 ylims=(qpar_min, qpar_max), label=run_label)
+                end
             end
-            outfile = string(run_name, "_qpar_vs_z_spec", spec_string, ".gif")
+            outfile = string(prefix, "_qpar_vs_z_spec", spec_string, ".gif")
             gif(anim, outfile, fps=5)
         end
     end
@@ -848,7 +1140,7 @@ Note this function requires using the PyPlot backend to support 2d coordinates b
 passed to `heatmap()`.
 """
 function plot_unnormalised(f, z_grid, vpa_grid, density, upar, vth, evolve_density,
-                           evolve_upar, evolve_ppar; plot_log=false)
+                           evolve_upar, evolve_ppar; plot_log=false, kwargs...)
 
     if backend_name() != :pyplot
         error("PyPlot backend is required for plot_unnormalised(). Call pyplot() "
@@ -886,7 +1178,7 @@ function plot_unnormalised(f, z_grid, vpa_grid, density, upar, vth, evolve_densi
         cmap = :deep
     end
 
-    return heatmap(z2d, dzdt2d, f_unnorm, xlabel="z", ylabel="vpa", c=cmap)
+    return heatmap(z2d, dzdt2d, f_unnorm; xlabel="z", ylabel="vpa", c=cmap, kwargs...)
 end
 
 end

From 5d884dc8aeae12dba45cacd947b35ae0ab8984ca Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 1 Sep 2022 11:45:18 +0100
Subject: [PATCH 09/17] Fix export in post_processing.jl

`export analyze_and_plot` -> `export analyze_and_plot_data`
---
 src/post_processing.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/post_processing.jl b/src/post_processing.jl
index 58b24ac330..7b52754bb9 100644
--- a/src/post_processing.jl
+++ b/src/post_processing.jl
@@ -2,7 +2,7 @@
 """
 module post_processing
 
-export analyze_and_plot
+export analyze_and_plot_data
 
 # packages
 using Plots

From 842fe2586202f541ad0e538795aaf7a828d3aa69 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 13 Sep 2022 21:45:56 +0100
Subject: [PATCH 10/17] 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
https://github.com/JuliaPlots/Plots.jl/pull/4298
---
 Manifest.toml          |  50 +++++++++++-----
 Project.toml           |   2 +
 src/post_processing.jl | 131 ++++++++++++++++++++++++++++++++---------
 3 files changed, 138 insertions(+), 45 deletions(-)

diff --git a/Manifest.toml b/Manifest.toml
index 8eb638d6cd..928c719d48 100644
--- a/Manifest.toml
+++ b/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"]
diff --git a/Project.toml b/Project.toml
index 4d9925b301..ee03684af5 100644
--- a/Project.toml
+++ b/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"
diff --git a/src/post_processing.jl b/src/post_processing.jl
index 7b52754bb9..5e77638009 100644
--- a/src/post_processing.jl
+++ b/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

From c8b1a3372a6855d780499fd8e4925c82df9a50ec Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 13 Sep 2022 22:35:51 +0100
Subject: [PATCH 11/17] pip-install matplotlib in CI jobs

Needed when using PyCall to import matplotlib.
---
 .github/workflows/debug_checks.yml  | 1 +
 .github/workflows/documentation.yml | 4 +++-
 .github/workflows/examples.yml      | 1 +
 .github/workflows/longtest.yml      | 1 +
 .github/workflows/parallel_test.yml | 1 +
 .github/workflows/test.yml          | 2 ++
 6 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/debug_checks.yml b/.github/workflows/debug_checks.yml
index d1822c9399..dd7d3d82d1 100644
--- a/.github/workflows/debug_checks.yml
+++ b/.github/workflows/debug_checks.yml
@@ -24,6 +24,7 @@ jobs:
       - uses: julia-actions/julia-buildpkg@v1
       - name: Test examples
         run: |
+          pip install --user matplotlib
           julia --project -e 'ENV["JULIA_MPI_BINARY"]="system"; using Pkg; Pkg.build("MPI"; verbose=true)'
           # Need to use openmpi so that the following arguments work:
           # * `--mca rmaps_base_oversubscribe 1` allows oversubscription (more processes
diff --git a/.github/workflows/documentation.yml b/.github/workflows/documentation.yml
index 41c9be3dfe..5ae248c5ec 100644
--- a/.github/workflows/documentation.yml
+++ b/.github/workflows/documentation.yml
@@ -16,7 +16,9 @@ jobs:
         with:
           version: '1'
       - name: Install dependencies
-        run: julia --project=docs/ -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()'
+        run: |
+          pip install --user matplotlib
+          julia --project=docs/ -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()'
       - name: Build and deploy
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} # Authenticate with GitHub Actions token
diff --git a/.github/workflows/examples.yml b/.github/workflows/examples.yml
index 3c5567a7d7..53466dd79b 100644
--- a/.github/workflows/examples.yml
+++ b/.github/workflows/examples.yml
@@ -21,4 +21,5 @@ jobs:
       - uses: julia-actions/julia-buildpkg@v1
       - name: Test examples
         run: |
+          pip install --user matplotlib
           julia --project -O3 -e 'using moment_kinetics; for (root, dirs, files) in walkdir("examples") for file in files if endswith(file, ".toml") println(joinpath(root, file)); run_moment_kinetics(joinpath(root, file)) end end end'
diff --git a/.github/workflows/longtest.yml b/.github/workflows/longtest.yml
index e98f39b5df..0f7ca4000c 100644
--- a/.github/workflows/longtest.yml
+++ b/.github/workflows/longtest.yml
@@ -28,5 +28,6 @@ jobs:
         # in order to pass customised arguments to `Pkg.test()`
         #
       - run: |
+          pip install --user matplotlib
           julia --check-bounds=yes --color=yes --depwarn=yes --inline=yes --project=@. -e 'import Pkg; Pkg.test(; coverage = :true, test_args=["--long"])'
         shell: bash
diff --git a/.github/workflows/parallel_test.yml b/.github/workflows/parallel_test.yml
index 6066609b78..c02aa710c7 100644
--- a/.github/workflows/parallel_test.yml
+++ b/.github/workflows/parallel_test.yml
@@ -23,6 +23,7 @@ jobs:
           arch: x64
       - uses: julia-actions/julia-buildpkg@v1
       - run: |
+          pip install --user matplotlib
           julia --project -e 'ENV["JULIA_MPI_BINARY"]="system"; using Pkg; Pkg.build("MPI"; verbose=true)'
           julia -O3 --check-bounds=no precompile.jl --debug 1
           # Need to use openmpi so that the following arguments work:
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 7b92703061..781ade317a 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -19,4 +19,6 @@ jobs:
           version: '1.7'
           arch: x64
       - uses: julia-actions/julia-buildpkg@v1
+      - run: |
+          pip install --user matplotlib
       - uses: julia-actions/julia-runtest@v1

From b5bec8b7e0049ac2b1cd512b6dfc907d2d4e6220 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 13 Sep 2022 22:35:51 +0100
Subject: [PATCH 12/17] Use system Python for PyCall on all platforms in CI
 jobs

Using the system Python is the default on Linux, but not macOS (or
Windows). This commit explicitly sets the ENV["PYTHON"] variable so that
all platforms use the system Python. This was needed to get PyCall to
work correctly on macOS - previously Julia's internal Conda-installed
Python had linking errors when using matplotlib.
---
 .github/workflows/debug_checks.yml  | 6 +++++-
 .github/workflows/documentation.yml | 2 +-
 .github/workflows/examples.yml      | 6 +++++-
 .github/workflows/longtest.yml      | 6 +++++-
 .github/workflows/parallel_test.yml | 6 +++++-
 .github/workflows/test.yml          | 6 +++++-
 6 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/.github/workflows/debug_checks.yml b/.github/workflows/debug_checks.yml
index dd7d3d82d1..a1f1b8f633 100644
--- a/.github/workflows/debug_checks.yml
+++ b/.github/workflows/debug_checks.yml
@@ -17,14 +17,18 @@ jobs:
       - uses: mpi4py/setup-mpi@v1
         with:
           mpi: 'openmpi'
+      - uses: actions/setup-python@v2
       - uses: julia-actions/setup-julia@v1
         with:
           version: '1.7'
           arch: x64
       - uses: julia-actions/julia-buildpkg@v1
+        env:
+          # Use the system Python for PyCall - avoids library linking error on macOS
+          PYTHON: "${{ env.pythonLocation }}/bin/python"
       - name: Test examples
         run: |
-          pip install --user matplotlib
+          pip3 install --user matplotlib
           julia --project -e 'ENV["JULIA_MPI_BINARY"]="system"; using Pkg; Pkg.build("MPI"; verbose=true)'
           # Need to use openmpi so that the following arguments work:
           # * `--mca rmaps_base_oversubscribe 1` allows oversubscription (more processes
diff --git a/.github/workflows/documentation.yml b/.github/workflows/documentation.yml
index 5ae248c5ec..70822766da 100644
--- a/.github/workflows/documentation.yml
+++ b/.github/workflows/documentation.yml
@@ -17,7 +17,7 @@ jobs:
           version: '1'
       - name: Install dependencies
         run: |
-          pip install --user matplotlib
+          pip3 install --user matplotlib
           julia --project=docs/ -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()'
       - name: Build and deploy
         env:
diff --git a/.github/workflows/examples.yml b/.github/workflows/examples.yml
index 53466dd79b..5cb77ac875 100644
--- a/.github/workflows/examples.yml
+++ b/.github/workflows/examples.yml
@@ -14,12 +14,16 @@ jobs:
 
     steps:
       - uses: actions/checkout@v2
+      - uses: actions/setup-python@v2
       - uses: julia-actions/setup-julia@v1
         with:
           version: '1.7'
           arch: x64
       - uses: julia-actions/julia-buildpkg@v1
+        env:
+          # Use the system Python for PyCall - avoids library linking error on macOS
+          PYTHON: "${{ env.pythonLocation }}/bin/python"
       - name: Test examples
         run: |
-          pip install --user matplotlib
+          pip3 install --user matplotlib
           julia --project -O3 -e 'using moment_kinetics; for (root, dirs, files) in walkdir("examples") for file in files if endswith(file, ".toml") println(joinpath(root, file)); run_moment_kinetics(joinpath(root, file)) end end end'
diff --git a/.github/workflows/longtest.yml b/.github/workflows/longtest.yml
index 0f7ca4000c..a06d9cd89c 100644
--- a/.github/workflows/longtest.yml
+++ b/.github/workflows/longtest.yml
@@ -18,16 +18,20 @@ jobs:
 
     steps:
       - uses: actions/checkout@v2
+      - uses: actions/setup-python@v2
       - uses: julia-actions/setup-julia@v1
         with:
           version: '1.7'
           arch: x64
       - uses: julia-actions/julia-buildpkg@v1
+        env:
+          # Use the system Python for PyCall - avoids library linking error on macOS
+          PYTHON: "${{ env.pythonLocation }}/bin/python"
         # The following is copied and simplified from
         # https://github.com/julia-actions/julia-runtest/blob/master/action.yml
         # in order to pass customised arguments to `Pkg.test()`
         #
       - run: |
-          pip install --user matplotlib
+          pip3 install --user matplotlib
           julia --check-bounds=yes --color=yes --depwarn=yes --inline=yes --project=@. -e 'import Pkg; Pkg.test(; coverage = :true, test_args=["--long"])'
         shell: bash
diff --git a/.github/workflows/parallel_test.yml b/.github/workflows/parallel_test.yml
index c02aa710c7..d9ecd80ca0 100644
--- a/.github/workflows/parallel_test.yml
+++ b/.github/workflows/parallel_test.yml
@@ -17,13 +17,17 @@ jobs:
       - uses: mpi4py/setup-mpi@v1
         with:
           mpi: 'openmpi'
+      - uses: actions/setup-python@v2
       - uses: julia-actions/setup-julia@v1
         with:
           version: '1.7'
           arch: x64
       - uses: julia-actions/julia-buildpkg@v1
+        env:
+          # Use the system Python for PyCall - avoids library linking error on macOS
+          PYTHON: "${{ env.pythonLocation }}/bin/python"
       - run: |
-          pip install --user matplotlib
+          pip3 install --user matplotlib
           julia --project -e 'ENV["JULIA_MPI_BINARY"]="system"; using Pkg; Pkg.build("MPI"; verbose=true)'
           julia -O3 --check-bounds=no precompile.jl --debug 1
           # Need to use openmpi so that the following arguments work:
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 781ade317a..87f88fe842 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -14,11 +14,15 @@ jobs:
 
     steps:
       - uses: actions/checkout@v2
+      - uses: actions/setup-python@v2
       - uses: julia-actions/setup-julia@v1
         with:
           version: '1.7'
           arch: x64
       - uses: julia-actions/julia-buildpkg@v1
+        env:
+          # Use the system Python for PyCall - avoids library linking error on macOS
+          PYTHON: "${{ env.pythonLocation }}/bin/python"
       - run: |
-          pip install --user matplotlib
+          pip3 install --user matplotlib
       - uses: julia-actions/julia-runtest@v1

From e5ae0789f14a94f10950f0e2199452a02cba4d9a Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 14 Sep 2022 15:28:20 +0100
Subject: [PATCH 13/17] Fix figure sizes

When creating figures for multiple runs, we need to set the figure size
explicitly for 2d plots (it should get wider when there are more runs to
avoid squashing the figures). This commit changes the settings so that
the size of the plot for each run matches the Plots.jl default size for
a single plot.
---
 src/post_processing.jl | 26 +++++++++++++-------------
 1 file changed, 13 insertions(+), 13 deletions(-)

diff --git a/src/post_processing.jl b/src/post_processing.jl
index 5e77638009..e949d2dceb 100644
--- a/src/post_processing.jl
+++ b/src/post_processing.jl
@@ -302,7 +302,7 @@ function plot_1D_1V_diagnostics(run_names, run_labels, nc_files, nwrite_movie,
                 subplots = (@views heatmap(z, vpa, log.(abs.(f[:,:,is,i])), xlabel="z",
                                            ylabel="vpa", fillcolor = logdeep, title=run_label)
                             for (f, this_z, this_vpa, run_label) ∈ zip(ff, z, vpa, run_labels))
-                plot(subplots..., layout=(1,n_runs), size=(400*n_runs, 400))
+                plot(subplots..., layout=(1,n_runs), size=(600*n_runs, 400))
             end
             outfile = string(prefix, "_logf_vs_vpa_z", spec_string, ".gif")
             gif(anim, outfile, fps=5)
@@ -313,7 +313,7 @@ function plot_1D_1V_diagnostics(run_names, run_labels, nc_files, nwrite_movie,
                                            ylabel="vpa", c = :deep, interpolation =
                                            :cubic, title=run_label)
                             for (f, this_z, this_vpa, run_label) ∈ zip(ff, z, vpa, run_labels))
-                plot(subplots..., layout=(1,n_runs), size=(400*n_runs, 400))
+                plot(subplots..., layout=(1,n_runs), size=(600*n_runs, 400))
             end
             outfile = string(prefix, "_f_vs_vpa_z", spec_string, ".gif")
             gif(anim, outfile, fps=5)
@@ -323,7 +323,7 @@ function plot_1D_1V_diagnostics(run_names, run_labels, nc_files, nwrite_movie,
                                        c = :deep, interpolation = :cubic,
                                        title=string(run_label, str))
                         for (f, this_z, this_vpa, run_label) ∈ zip(ff, z, vpa, run_labels))
-            plot(subplots..., layout=(1,n_runs), size=(400*n_runs, 400))
+            plot(subplots..., layout=(1,n_runs), size=(600*n_runs, 400))
             outfile = string(prefix, "_f_vs_z_vpa_final", spec_string, ".pdf")
             savefig(outfile)
 
@@ -360,7 +360,7 @@ function plot_1D_1V_diagnostics(run_names, run_labels, nc_files, nwrite_movie,
             #                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))
+            #    plot(subplots..., layout=(1,n_runs), size=(600*n_runs, 400))
             #end
             #outfile = string(prefix, "_f_unnorm_vs_vpa_z", spec_string, ".gif")
             #gif(anim, outfile, fps=5)
@@ -374,7 +374,7 @@ function plot_1D_1V_diagnostics(run_names, run_labels, nc_files, nwrite_movie,
             #                    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))
+            #    plot(subplots..., layout=(1,n_runs), size=(600*n_runs, 400))
             #end
             #outfile = string(prefix, "_logf_unnorm_vs_vpa_z", spec_string, ".gif")
             #gif(anim, outfile, fps=5)
@@ -400,7 +400,7 @@ function plot_1D_1V_diagnostics(run_names, run_labels, nc_files, nwrite_movie,
                                              title=run_label)
                 end
             end
-            fig = PyPlot.figure(1, figsize=(4,4))
+            fig = PyPlot.figure(1, figsize=(6,4))
             fig.clf()
             myanim = matplotlib_animation.FuncAnimation(fig, make_frame, frames=nframes)
             outfile = string(prefix, "_f_unnorm_vs_vpa_z", spec_string, ".gif")
@@ -422,7 +422,7 @@ function plot_1D_1V_diagnostics(run_names, run_labels, nc_files, nwrite_movie,
                                              title=run_label, plot_log=true)
                 end
             end
-            fig = PyPlot.figure(figsize=(4,4))
+            fig = PyPlot.figure(figsize=(6,4))
             myanim = matplotlib_animation.FuncAnimation(fig, make_frame_log, frames=nframes)
             outfile = string(prefix, "_logf_unnorm_vs_vpa_z", spec_string, ".gif")
             myanim.save(outfile, writer=matplotlib_animation.PillowWriter(fps=30))
@@ -435,7 +435,7 @@ function plot_1D_1V_diagnostics(run_names, run_labels, nc_files, nwrite_movie,
                                        interpolation = :cubic, title=run_label)
                             for (df, this_z, this_vpa, run_label) ∈
                                 zip(delta_f, z, vpa, run_labels))
-                plot(subplots..., layout=(1,n_runs), size=(400*n_runs, 400))
+                plot(subplots..., layout=(1,n_runs), size=(600*n_runs, 400))
             end
             outfile = string(prefix, "_deltaf_vs_vpa_z", spec_string, ".gif")
             gif(anim, outfile, fps=5)
@@ -651,7 +651,7 @@ function plot_fields(phi, delta_phi, time, itime_min, itime_max, nwrite_movie,
         subplots = (heatmap(t, this_z, p, xlabel="time", ylabel="z", title=run_label, c =
                             :deep)
                     for (t, this_z, p, run_label) ∈ zip(time, z, phi, run_labels))
-        plot(subplots..., layout=(1,n_runs), size=(400*n_runs, 400))
+        plot(subplots..., layout=(1,n_runs), size=(600*n_runs, 400))
         outfile = string(prefix, "_phi_vs_z_t.pdf")
         savefig(outfile)
     end
@@ -866,7 +866,7 @@ function plot_moments(density, delta_density, density_fldline_avg,
             subplots = (heatmap(t, this_z, n[:,is,:], xlabel="time", ylabel="z",
                                 title=run_label, c = :deep)
                         for (t, this_z, n, run_label) ∈ zip(time, z, density, run_labels))
-            plot(subplots..., layout=(1,n_runs), size=(400*n_runs, 400))
+            plot(subplots..., layout=(1,n_runs), size=(600*n_runs, 400))
             outfile = string(prefix, "_dens_vs_z_t_spec", spec_string, ".pdf")
             savefig(outfile)
         end
@@ -876,7 +876,7 @@ function plot_moments(density, delta_density, density_fldline_avg,
                                 title=run_label, c = :deep)
                         for (t, this_z, upar, run_label) ∈ zip(time, z, parallel_flow,
                                                               run_labels))
-            plot(subplots..., layout=(1,n_runs), size=(400*n_runs, 400))
+            plot(subplots..., layout=(1,n_runs), size=(600*n_runs, 400))
             outfile = string(prefix, "_upar_vs_z_t_spec", spec_string, ".pdf")
             savefig(outfile)
         end
@@ -886,7 +886,7 @@ function plot_moments(density, delta_density, density_fldline_avg,
                                 title=run_label, c = :deep)
                         for (t, this_z, ppar, run_label) ∈
                             zip(time, z, parallel_pressure, run_labels))
-            plot(subplots..., layout=(1,n_runs), size=(400*n_runs, 400))
+            plot(subplots..., layout=(1,n_runs), size=(600*n_runs, 400))
             outfile = string(prefix, "_ppar_vs_z_t_spec", spec_string, ".pdf")
             savefig(outfile)
         end
@@ -896,7 +896,7 @@ function plot_moments(density, delta_density, density_fldline_avg,
                                 title=run_label, c = :deep)
                         for (t, this_z, qpar, run_label) ∈
                             zip(time, z, parallel_heat_flux, run_labels))
-            plot(subplots..., layout=(1,n_runs), size=(400*n_runs, 400))
+            plot(subplots..., layout=(1,n_runs), size=(600*n_runs, 400))
             outfile = string(prefix, "_qpar_vs_z_t_spec", spec_string, ".pdf")
             savefig(outfile)
         end

From 1ed8c4748344fac1d85934ec2219c50b7946d492 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 16 Sep 2022 15:18:30 +0100
Subject: [PATCH 14/17] Fix coordinate arguments for multiple-run plotting in
 animate_f_vs_vpa_z

---
 src/post_processing.jl | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/post_processing.jl b/src/post_processing.jl
index e949d2dceb..f6b51a19ac 100644
--- a/src/post_processing.jl
+++ b/src/post_processing.jl
@@ -299,8 +299,9 @@ function plot_1D_1V_diagnostics(run_names, run_labels, nc_files, nwrite_movie,
             # make a gif animation of ln f(vpa,z,t)
             anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
                 #heatmap(z, vpa, log.(abs.(ff[:,:,i])), xlabel="z", ylabel="vpa", clims = (fmin,fmax), c = :deep)
-                subplots = (@views heatmap(z, vpa, log.(abs.(f[:,:,is,i])), xlabel="z",
-                                           ylabel="vpa", fillcolor = logdeep, title=run_label)
+                subplots = (@views heatmap(this_z, this_vpa, log.(abs.(f[:,:,is,i])),
+                                           xlabel="z", ylabel="vpa", fillcolor =
+                                           logdeep, title=run_label)
                             for (f, this_z, this_vpa, run_label) ∈ zip(ff, z, vpa, run_labels))
                 plot(subplots..., layout=(1,n_runs), size=(600*n_runs, 400))
             end

From 3862674f205f165b8e2ac24fafad75e29ffac119 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 19 Sep 2022 11:00:18 +0100
Subject: [PATCH 15/17] Clear PyPlot figure after making f_unnorm plots

After switching between PyPlot interface and Plots interface, the Figure
apparently does not get cleaned up without an explicit call to
Plots.closeall().
---
 src/post_processing.jl | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/post_processing.jl b/src/post_processing.jl
index f6b51a19ac..a4788e7ac1 100644
--- a/src/post_processing.jl
+++ b/src/post_processing.jl
@@ -402,10 +402,10 @@ function plot_1D_1V_diagnostics(run_names, run_labels, nc_files, nwrite_movie,
                 end
             end
             fig = PyPlot.figure(1, figsize=(6,4))
-            fig.clf()
             myanim = matplotlib_animation.FuncAnimation(fig, make_frame, frames=nframes)
             outfile = string(prefix, "_f_unnorm_vs_vpa_z", spec_string, ".gif")
             myanim.save(outfile, writer=matplotlib_animation.PillowWriter(fps=30))
+            PyPlot.clf()
 
             function make_frame_log(i)
                 PyPlot.clf()
@@ -427,6 +427,10 @@ function plot_1D_1V_diagnostics(run_names, run_labels, nc_files, nwrite_movie,
             myanim = matplotlib_animation.FuncAnimation(fig, make_frame_log, frames=nframes)
             outfile = string(prefix, "_logf_unnorm_vs_vpa_z", spec_string, ".gif")
             myanim.save(outfile, writer=matplotlib_animation.PillowWriter(fps=30))
+
+            # Ensure PyPlot figure is cleared
+            closeall()
+
         end
         if pp.animate_deltaf_vs_vpa_z
             # make a gif animation of δf(vpa,z,t)

From cc46f599af16f732f6f225cd33f680a300d1ec79 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 19 Sep 2022 11:01:47 +0100
Subject: [PATCH 16/17] Make plots of unnormalized f at the wall boundaries

Includes some refactoring of functions used to get/plot unnormalized f
so that they can be reused in more ways.
---
 src/post_processing.jl | 120 +++++++++++++++++++++++++++++------------
 1 file changed, 85 insertions(+), 35 deletions(-)

diff --git a/src/post_processing.jl b/src/post_processing.jl
index a4788e7ac1..cfab5a1ebe 100644
--- a/src/post_processing.jl
+++ b/src/post_processing.jl
@@ -395,10 +395,11 @@ function plot_1D_1V_diagnostics(run_names, run_labels, nc_files, nwrite_movie,
                                       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)
+                    @views f_unnorm, z2d, dzdt2d = get_unnormalised_f_coords_2d(
+                        f[:,:,is,iframe], this_z, this_vpa, n[:,is,iframe],
+                        upar[:,is,iframe], vth[:,is,iframe], ev_n, ev_u, ev_p)
+                    plot_unnormalised_f2d(f_unnorm, z2d, dzdt2d; title=run_label,
+                                          plot_log=false)
                 end
             end
             fig = PyPlot.figure(1, figsize=(6,4))
@@ -417,10 +418,11 @@ function plot_1D_1V_diagnostics(run_names, run_labels, nc_files, nwrite_movie,
                                       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)
+                    @views f_unnorm, z2d, dzdt2d = get_unnormalised_f_coords_2d(
+                        f[:,:,is,iframe], this_z, this_vpa, n[:,is,iframe],
+                        upar[:,is,iframe], vth[:,is,iframe], ev_n, ev_u, ev_p)
+                    plot_unnormalised_f2d(f_unnorm, z2d, dzdt2d; title=run_label,
+                                          plot_log=true)
                 end
             end
             fig = PyPlot.figure(figsize=(6,4))
@@ -431,6 +433,35 @@ function plot_1D_1V_diagnostics(run_names, run_labels, nc_files, nwrite_movie,
             # Ensure PyPlot figure is cleared
             closeall()
 
+            anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
+                plot(legend=legend)
+                for (f, n, upar, vth, ev_n, ev_u, ev_p, this_vpa, run_label) ∈
+                    zip(ff, density, parallel_flow, thermal_speed, evolve_density,
+                        evolve_upar, evolve_ppar, vpa, run_labels)
+                    @views f_unnorm, dzdt = get_unnormalised_f_dzdt_1d(
+                        f[:,1,is,i], this_vpa, n[1,is,i], upar[1,is,i], vth[1,is,i],
+                        ev_n, ev_u, ev_p)
+                    @views plot!(dzdt, f_unnorm, xlabel="vpa", ylabel="f_unnorm(z=0)",
+                                 label=run_label)
+                end
+            end
+            outfile = string(prefix, "_f0_unnorm_vs_vpa", spec_string, ".gif")
+            gif(anim, outfile, fps=5)
+
+            anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
+                plot(legend=legend)
+                for (f, n, upar, vth, ev_n, ev_u, ev_p, this_vpa, run_label) ∈
+                    zip(ff, density, parallel_flow, thermal_speed, evolve_density,
+                        evolve_upar, evolve_ppar, vpa, run_labels)
+                    @views f_unnorm, dzdt = get_unnormalised_f_dzdt_1d(
+                        f[:,end,is,i], this_vpa, n[end,is,i], upar[end,is,i],
+                        vth[end,is,i], ev_n, ev_u, ev_p)
+                    @views plot!(dzdt, f_unnorm, xlabel="vpa", ylabel="f_unnorm(z=L)",
+                                 label=run_label)
+                end
+            end
+            outfile = string(prefix, "_fL_unnorm_vs_vpa", spec_string, ".gif")
+            gif(anim, outfile, fps=5)
         end
         if pp.animate_deltaf_vs_vpa_z
             # make a gif animation of δf(vpa,z,t)
@@ -1190,43 +1221,62 @@ function draw_v_parallel_zero!(z::AbstractVector, upar, vth, evolve_upar::Bool,
 end
 
 """
-Plot the unnormalised distribution function against unnormalised ('lab space')
-coordinates.
-
-Inputs should depend only on z and vpa.
+Get the unnormalised distribution function and unnormalised ('lab space') dzdt
+coordinate at a point in space.
 
-Note this function requires using the PyPlot backend to support 2d coordinates being
-passed to `heatmap()`.
+Inputs should depend only on vpa.
 """
-function plot_unnormalised(f, z_grid, vpa_grid, density, upar, vth, evolve_density,
-                           evolve_upar, evolve_ppar; plot_log=false, kwargs...)
+function get_unnormalised_f_dzdt_1d(f, vpa_grid, density, upar, vth, evolve_density,
+                                    evolve_upar, evolve_ppar)
 
-    if backend_name() != :pyplot
-        error("PyPlot backend is required for plot_unnormalised(). Call pyplot() "
-              * "first.")
+    dzdt = vpagrid_to_dzdt(vpa_grid, vth, upar, evolve_ppar, evolve_upar)
+
+    if evolve_ppar
+        f_unnorm = @. f * density / vth
+    elseif evolve_density
+        f_unnorm = @. f * density
+    else
+        f_unnorm = copy(f)
     end
 
+    return f_unnorm, dzdt
+end
+
+"""
+Get the unnormalised distribution function and unnormalised ('lab space') coordinates.
+
+Inputs should depend only on z and vpa.
+"""
+function get_unnormalised_f_coords_2d(f, z_grid, vpa_grid, density, upar, vth,
+                                      evolve_density, evolve_upar, evolve_ppar)
+
     nvpa, nz = size(f)
     z2d = zeros(nvpa, nz)
     dzdt2d = zeros(nvpa, nz)
+    f_unnorm = similar(f)
     for iz ∈ 1:nz
         @views z2d[:,iz] .= z_grid[iz]
-        @views dzdt2d[:,iz] .= vpagrid_to_dzdt(vpa_grid, vth[iz], upar[iz], evolve_ppar,
-                                             evolve_upar)
+        f_unnorm[:,iz], dzdt2d[:,iz] =
+            get_unnormalised_f_dzdt_1d(f[:,iz], vpa_grid, density[iz], upar[iz],
+                                       vth[iz], evolve_density, evolve_upar,
+                                       evolve_ppar)
     end
 
-    if evolve_ppar
-        f_unnorm = similar(f)
-        for iz ∈ 1:nz, ivpa ∈ 1:nvpa
-            f_unnorm[ivpa,iz] = @. f[ivpa,iz] * density[iz] / vth[iz]
-        end
-    elseif evolve_density
-        f_unnorm = similar(f)
-        for iz ∈ 1:nz, ivpa ∈ 1:nvpa
-            f_unnorm[ivpa,iz] = @. f[ivpa,iz] * density[iz]
-        end
-    else
-        f_unnorm = f
+    return f_unnorm, z2d, dzdt2d
+end
+
+"""
+Make a 2d plot of an unnormalised f on unnormalised coordinates, as returned from
+get_unnormalised_f_coords()
+
+Note this function requires using the PyPlot backend to support 2d coordinates being
+passed to `heatmap()`.
+"""
+function plot_unnormalised_f2d(f_unnorm, z2d, dzdt2d; plot_log=false, kwargs...)
+
+    if backend_name() != :pyplot
+        error("PyPlot backend is required for plot_unnormalised(). Call pyplot() "
+              * "first.")
     end
 
     ## The following commented out section does not work at the moment because
@@ -1246,8 +1296,8 @@ function plot_unnormalised(f, z_grid, vpa_grid, density, upar, vth, evolve_densi
 
     # Use PyPlot directly instead. Unfortunately makes animation a pain...
     if plot_log
-        vmin = minimum(x for x in f if x > 0.0)
-        norm = PyPlot.matplotlib.colors.LogNorm(vmin=vmin, vmax=maximum(f))
+        vmin = minimum(x for x in f_unnorm if x > 0.0)
+        norm = PyPlot.matplotlib.colors.LogNorm(vmin=vmin, vmax=maximum(f_unnorm))
     else
         norm = nothing
     end

From 52c30df4fc6ad483b5a3a4067f53ee3f4f9c242e Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 20 Sep 2022 09:23:00 +0100
Subject: [PATCH 17/17] Fix width of multi-run unnormalized pdf plots

---
 src/post_processing.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/post_processing.jl b/src/post_processing.jl
index cfab5a1ebe..cf77c04299 100644
--- a/src/post_processing.jl
+++ b/src/post_processing.jl
@@ -402,7 +402,7 @@ function plot_1D_1V_diagnostics(run_names, run_labels, nc_files, nwrite_movie,
                                           plot_log=false)
                 end
             end
-            fig = PyPlot.figure(1, figsize=(6,4))
+            fig = PyPlot.figure(1, figsize=(6*n_runs,4))
             myanim = matplotlib_animation.FuncAnimation(fig, make_frame, frames=nframes)
             outfile = string(prefix, "_f_unnorm_vs_vpa_z", spec_string, ".gif")
             myanim.save(outfile, writer=matplotlib_animation.PillowWriter(fps=30))
@@ -425,7 +425,7 @@ function plot_1D_1V_diagnostics(run_names, run_labels, nc_files, nwrite_movie,
                                           plot_log=true)
                 end
             end
-            fig = PyPlot.figure(figsize=(6,4))
+            fig = PyPlot.figure(figsize=(6*n_runs,4))
             myanim = matplotlib_animation.FuncAnimation(fig, make_frame_log, frames=nframes)
             outfile = string(prefix, "_logf_unnorm_vs_vpa_z", spec_string, ".gif")
             myanim.save(outfile, writer=matplotlib_animation.PillowWriter(fps=30))